Update app.py

app.py

@@ -14,8 +14,31 @@ from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 # Create a ToolNode that knows about your web_search function
+import json
 
-# (Keep Constants as is)
+def parse_tool_json(text: str) -> dict | None:
+    """
+    Given a string like '{"tool":"web_search","query":"..."}'
+    or '"{\"tool\":\"web_search\",\"query\":\"...\"}"', return
+    the parsed dict. Otherwise, return None.
+    """
+    t = text.strip()
+    # If it’s wrapped in single or double quotes, remove them:
+    if (t.startswith('"') and t.endswith('"')) or (t.startswith("'") and t.endswith("'")):
+        t = t[1:-1]
+    try:
+        obj = json.loads(t)
+        if isinstance(obj, dict) and "tool" in obj:
+            return obj
+    except Exception:
+        return None
+    return None
+
+
+# (Keep Constan
+# 
+# 
+# ts as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
@@ -39,19 +62,20 @@ def respond_to_input(user_input: str) -> str:
     """
     # (A) First message: describe your tools
     system_msg = SystemMessage(
-        content=(
-            "You are an assistant with access to the following tools:\n"
-            "  1) web_search(query: str) → Returns the top search results for the query.\n"
-            "  2) parse_excel(path: str, sheet_name: str) → Reads an Excel file and returns its contents.\n"
-            "  3) ocr_image(path: str) → Runs OCR on an image and returns any detected text.\n\n"
-            "When you need up‐to‐date info, respond exactly with JSON:\n"
-            '  { "tool": "web_search", "query": "<search terms>" }\n'
-            "If you need to read an Excel file, respond:\n"
-            '  { "tool": "parse_excel", "path": "<file.xlsx>", "sheet_name": "<SheetName>" }\n'
-            "If you need OCR, respond:\n"
-            '  { "tool": "ocr_image", "path": "<image.png>" }\n'
-            "Otherwise, reply only with your final answer as plain text."
-        )
+    content=(
+        "You are an assistant with access to exactly these tools:\n"
+        "  1) web_search(query:str)\n"
+        "  2) parse_excel(path:str,sheet_name:str)\n"
+        "  3) ocr_image(path:str)\n\n"
+        "⚠️ **IMPORTANT** ⚠️: If (and only if) you need to call one of these tools, "
+        "output exactly one JSON object and nothing else. For example:\n"
+        "```json\n"
+        '{"tool":"web_search","query":"Mercedes Sosa albums 2000-2009"}\n'
+        "```\n"
+        "That JSON must start at the very first character of your response and end at the very last character—"
+        "no quotes, no code fences, no extra explanation. \n\n"
+        "If you do NOT need any tool, simply reply with your final answer as plain text (no JSON)."
+    )
     )
 
     # (B) initial state has only the system prompt
@@ -61,12 +85,33 @@ def respond_to_input(user_input: str) -> str:
     final_state = compiled_graph.invoke(initial_state, user_input)
 
     # (D) Pull out the last AIMessage from final_state["messages"]:
-    assistant_messages = [
-        msg.content
-        for msg in final_state["messages"]
-        if isinstance(msg, AIMessage)
-    ]
-    return assistant_messages[-1] if assistant_messages else ""
+    last_msg = None
+    for msg in final_state["messages"][::-1]:
+        if isinstance(msg, AIMessage):
+            last_msg = msg.content
+            break
+
+    # Try to parse it as a tool‐call dict:
+    tool_dict = parse_tool_json(last_msg or "")
+    if tool_dict is not None:
+        # run the tool and then feed it back into the LLM:
+        result = tool_node.run(tool_dict)
+        # Append the tool’s result as an AIMessage, then invoke again:
+        new_state = {
+            "messages": [
+                *final_state["messages"],
+                AIMessage(content=result)
+            ]
+        }
+        second_pass = compiled_graph.invoke(new_state, "")
+        # Get that second pass’s last AIMessage:
+        for msg in second_pass["messages"][::-1]:
+            if isinstance(msg, AIMessage):
+                return msg.content
+        return ""
+    else:
+        # No valid tool JSON → treat last_msg as final answer
+        return last_msg or ""
 
 
 class BasicAgent: