Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import os
-
+import json
 os.environ["OPENAI_API_KEY"] = os.getenv('api_key')
 
 import math
@@ -101,7 +101,7 @@ builder = create_agent(llm, tool_registry)
 agent = builder.compile(store=store)
 
 def pvsnp(problem):
-    output = []
+    '''output = []
     for step in agent.stream(
     {"messages": "Use tools to answer:"+problem},
     stream_mode="updates",
@@ -110,14 +110,37 @@ def pvsnp(problem):
            for message in update.get("messages", []):
               message.pretty_print()
               output.append(message.pretty_print())
-    print (output)
-    return output
+    print (output)'''
+    output = agent.invoke({"messages": "Use tools to answer: "+problem})
+
+# Assuming `output` is a JSON-formatted string
+    data = json.loads(output)
+
+# Extract the answer from the final AIMessage
+    messages = data.get("messages", [])
+    answer = None
+    for msg in reversed(messages):
+        if msg.get("content"):
+            answer = msg["content"]
+            break
+
+# Extract tool names used in tool_calls
+    tools_used = []
+    for msg in messages:
+        if "tool_calls" in msg:
+            for tool_call in msg["tool_calls"]:
+                tools_used.append(tool_call["name"])
+
+    print("Answer:", answer)
+    print("Tools used:", tools_used)
+    response = "Answer: "+answer+'\n'+"Tools used: "+tools_used
+    return response
                                       
 
 iface = gr.Interface(
     fn=pvsnp, 
     inputs=gr.Textbox(label="What problem would you like to classify as P or NP?"), 
-    outputs=gr.Markdown(label="Agent's response"),  # Output as HTML
+    outputs=gr.Textbox(label="Agent's response"),  # Output as HTML
     title="PolyMath",
     description="PolyMath is an advanced AI agent that guides users through the intricate maze of computational complexity. This agent scrutinizes problem descriptions with sophisticated LLM prompts and symbolic reasoning. It classifies problems into categories such as P, NP, NP-complete, NP-hard, or beyond (e.g., PSPACE, EXPTIME), while providing clear, concise explanations of its reasoning. As part of AI Quotient’s Millennium Math Challenge, it is the first step towards solving the P vs NP problem.",
     theme = gr.themes.Ocean(),