refactor tool validations

agent.py

@@ -217,8 +217,39 @@ def build_agent_graph(provider: str = "groq"):
         else:
             return {"messages": [system_prompt] + state["messages"]}
 
-    # ToolNode wrapper for actual tool use
-    tool_node = ToolNode(tools)
+    # Wrap tools with validation
+    def wrap_tool_with_validation(tool):
+        original_func = tool.__call__
+
+        def validated_call(*args, **kwargs):
+            response = original_func(*args, **kwargs)
+
+            try:
+                if not isinstance(response, dict):
+                    raise ValueError(
+                        f"Tool response must be a dict, got {type(response)}"
+                    )
+
+                # Check for common response keys
+                for key in ["web_results", "wiki_results", "transcript_results"]:
+                    if key in response:
+                        if not isinstance(response[key], str):
+                            raise ValueError(
+                                f"Tool response[{key}] must be string, got {type(response[key])}"
+                            )
+                        if not response[key].strip():
+                            raise ValueError(f"Tool response[{key}] is empty")
+
+                return response
+            except Exception as e:
+                return {"error": f"Tool response validation failed: {str(e)}"}
+
+        tool.__call__ = validated_call
+        return tool
+
+    # Apply validation wrapper to each tool
+    validated_tools = [wrap_tool_with_validation(tool) for tool in tools]
+    tool_node = ToolNode(validated_tools)
 
     # Define error handling node
     def error_handler_node(state: MessagesState) -> dict:

app.py

@@ -1,12 +1,13 @@
 import inspect
 import os
+from typing import Dict, List, TypedDict, cast
 
 import gradio as gr
 import pandas as pd
 import requests
-from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage
+from langgraph.graph import MessagesState
 
-# from langgraph.graph import MessagesState
 from agent import build_agent_graph
 
 # (Keep Constants as is)
@@ -23,12 +24,56 @@ class BasicAgent:
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # Wrap the question from HumanMessage from langchain_core
-        msgs = [HumanMessage(content=question)]
-        # input_state: MessagesState = {"messages": msgs}
-        result = self.graph.invoke({"messages": msgs})
-        answer = result["messages"][-1].content
-        return answer[14:]  # skip "FINAL ANSWER: "
+        try:
+            # Create properly typed messages
+            system_msg = SystemMessage(
+                content="""You are a helpful AI assistant. Format your final answer as:
+FINAL ANSWER: [your answer here]"""
+            )
+            human_msg = HumanMessage(content=question)
+            msgs: List[AnyMessage] = [system_msg, human_msg]
+
+            # Create and cast the state
+            input_state = cast(MessagesState, {"messages": msgs})
+
+            # Invoke the graph
+            result = self.graph.invoke(input_state)
+
+            # Validate response
+            if not isinstance(result, dict) or "messages" not in result:
+                raise ValueError("Invalid response structure from agent graph")
+
+            if not result["messages"]:
+                raise ValueError("Empty message list in response")
+
+            # Get the last message content
+            last_msg = result["messages"][-1]
+            if not isinstance(last_msg, (AIMessage, HumanMessage, SystemMessage)):
+                raise ValueError(f"Invalid message type: {type(last_msg)}")
+
+            answer = last_msg.content
+            if not isinstance(answer, str):
+                raise ValueError(f"Invalid answer type: {type(answer)}")
+
+            # Ensure proper formatting
+            if not answer.strip():
+                return "Error: Empty response from agent"
+
+            if "FINAL ANSWER:" not in answer:
+                # If no prefix, return as is
+                return answer
+
+            # Extract the actual answer after "FINAL ANSWER:"
+            final_answer = answer.split("FINAL ANSWER:", 1)[1].strip()
+            if not final_answer:
+                return "Error: Empty answer after FINAL ANSWER prefix"
+
+            return final_answer
+
+        except Exception as e:
+            error_msg = f"Error in agent call: {str(e)}"
+            print(error_msg)
+            return error_msg
 
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
@@ -62,6 +107,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
     # 2. Fetch Questions
     print(f"Fetching questions from: {questions_url}")
+    response = None
     try:
         response = requests.get(questions_url, timeout=15)
         response.raise_for_status()
@@ -70,13 +116,14 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             print("Fetched questions list is empty.")
             return "Fetched questions list is empty or invalid format.", None
         print(f"Fetched {len(questions_data)} questions.")
-    except requests.exceptions.RequestException as e:
-        print(f"Error fetching questions: {e}")
-        return f"Error fetching questions: {e}", None
     except requests.exceptions.JSONDecodeError as e:
+        error_text = response.text[:500] if response else "No response text available"
         print(f"Error decoding JSON response from questions endpoint: {e}")
-        print(f"Response text: {response.text[:500]}")
+        print(f"Response text: {error_text}")
         return f"Error decoding server response for questions: {e}", None
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching questions: {e}")
+        return f"Error fetching questions: {e}", None
     except Exception as e:
         print(f"An unexpected error occurred fetching questions: {e}")
         return f"An unexpected error occurred fetching questions: {e}", None
@@ -93,6 +140,17 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             continue
         try:
             submitted_answer = agent(question_text)
+            # Remove "FINAL ANSWER: " prefix if present
+            if isinstance(submitted_answer, str) and submitted_answer.startswith(
+                "FINAL ANSWER: "
+            ):
+                submitted_answer = submitted_answer[14:]
+
+            # Handle empty or invalid answers
+            if not submitted_answer:
+                print(f"Warning: Empty answer for task {task_id}")
+                submitted_answer = "Error: Agent produced empty response"
+
             answers_payload.append(
                 {"task_id": task_id, "submitted_answer": submitted_answer}
             )
@@ -103,13 +161,19 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                     "Submitted Answer": submitted_answer,
                 }
             )
+            print(f"Successfully processed task {task_id}")
         except Exception as e:
-            print(f"Error running agent on task {task_id}: {e}")
+            error_msg = f"Error running agent on task {task_id}: {str(e)}"
+            print(error_msg)
+            error_answer = f"Error: {str(e)}"
+            answers_payload.append(
+                {"task_id": task_id, "submitted_answer": error_answer}
+            )
             results_log.append(
                 {
                     "Task ID": task_id,
                     "Question": question_text,
-                    "Submitted Answer": f"AGENT ERROR: {e}",
+                    "Submitted Answer": error_answer,
                 }
             )