Update agent.py

agent.py

@@ -169,6 +169,7 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration
 
 from tavily import TavilyClient
 
+# GPT_4o_API_KEY  = os.environ["DeepSeek_API_KEY"]
 GPT_4o_API_KEY  = os.environ["GPT_4o_API_KEY"]
 TAVILY_API_KEY  = os.environ["TAVILY_API_KEY"]
 Pixtral12B_API  = os.environ["Pixtral12B_API"]
@@ -183,24 +184,17 @@ model = ChatOpenAI(openai_api_key=GPT_4o_API_KEY,
 tavily = TavilyClient(api_key=TAVILY_API_KEY)
 
 
-
 # === 1. AGENT STATE ===
-
-class GaiaAgentState(TypedDict):
-    task_id: int
-    question: str
-    files: List[str]
-    level: int
-    plan: List[Dict[str, Any]]
-    tool_results: List[Any]
-    final_answer: str
-    
-# def input_parser(state: GaiaAgentState) -> dict:
-#     return {
-#       "question": state["question"],
-#       "files":    state.get("files", []),
-#       "level":    state.get("level", 1)
-#     }
+@dataclass
+class GaiaAgentState:
+    raw_payload: Dict[str, Any]
+    question: str = ""
+    files:    List[str] = field(default_factory=list)
+    level:    int = 1
+    # downstream nodes can write into these:
+    plan:         List[Dict[str,Any]] = field(default_factory=list)
+    tool_results: List[Any]         = field(default_factory=list)
+    final_answer: str = ""
 
 # === 2. TOOL STUBS ===
 
@@ -259,11 +253,9 @@ def encode_image_to_base64(image_path):
 def vision_caption_via_api(image_path: str, max_length: Optional[int] = 64) -> str:
     """
     Generate a natural‑language caption for the image at `image_path`.
-
     Args:
       image_path: Local filesystem path (or URL fetched locally) to an image file.
       max_length: Maximum number of tokens in the generated caption (default: 64).
-
     Returns:
       A concise, human‑readable caption describing the image content.
     """
@@ -298,7 +290,7 @@ def vision_caption_via_api(image_path: str, max_length: Optional[int] = 64) -> s
         "messages": messages
     }
 
-    response = requests.post(url, headers=headers, json=payload)
+    response = requests.post(API_URL, headers=headers, json=payload)
     response.raise_for_status()
 
     # The HF endpoint returns JSON like {"generated_text": "...caption..."}
@@ -360,11 +352,9 @@ TOOLS: Dict[str, Any] = {
 def planner_node(state: GaiaAgentState) -> dict:
   # Build metadata for each file
   files_metadata = []
-  print(33333333333333, state)
-  print(55555555555555,state['files'])
-  for f_path in state['files']:
+  # print(55555555555555,state.files)
+  for f_path in state.files:
   # for f_path in state.files:
-      print(5555555555555555555555555555555555555555555555555555555)
       mime, _ = mimetypes.guess_type(f_path)
       kind = "image" if mime and mime.startswith("image/") else "text"
       files_metadata.append({
@@ -373,9 +363,9 @@ def planner_node(state: GaiaAgentState) -> dict:
       })
 
   human_content = json.dumps({
-      "question": state['question'],
+      "question": state.question,
       "files":    files_metadata,
-      "level":    state['level']
+      "level":    state.level
   })
 
   messages = [
@@ -389,25 +379,23 @@ def planner_node(state: GaiaAgentState) -> dict:
 
 def tool_executor(state: GaiaAgentState) -> dict:
   tool_results = []
-  for step in state['plan']:
+  for step in state.plan:
     tool_obj = TOOLS[step["tool"]]
     out = tool_obj.run(tool_input=step["args"])
     tool_results.append(out)
-  return {'tool_results': out}
+  return {'tool_results': tool_results}
 
 
 # ===== 2. The node implementation =====
 def synthesizer(state: GaiaAgentState) -> dict:
 
     # Format tool outputs as bullet points
-    snippet_text = "\n".join(f"- {s}" for s in state["tool_results"])
+    snippet_text = "\n".join(f"- {s}" for s in state.tool_results)
 
     human_payload = f"""
-      Question: {state['question']}
-
+      Question: {state.question}
       Tool outputs:
       {snippet_text}
-
       Using the above, answer the question.
       """
     messages = [
@@ -432,26 +420,51 @@ def synthesizer(state: GaiaAgentState) -> dict:
     return {"final_answer": final}
 
 
+# builder = StateGraph(dict)
+from dataclasses import dataclass, field
+
+
+@dataclass
+class GaiaAgentState:
+    raw_payload: Dict[str, Any]
+    # these will get populated by the input_parser node:
+    question: str = ""
+    files:    List[str] = field(default_factory=list)
+    level:    int = 1
+    # downstream nodes can write into these:
+    plan:         List[Dict[str,Any]] = field(default_factory=list)
+    tool_results: List[Any]         = field(default_factory=list)
+    final_answer: str = ""
+
+
+    
+
+# 2) Input‑parser node: unpacks raw_payload→working state
+def input_parser(state: GaiaAgentState) -> dict:
+    return {
+        "question": state.raw_payload["question"],
+        "files":    state.raw_payload.get("files", []),
+        "level":    state.raw_payload.get("level", 1),
+    }
+
+
+
 
-# === 4. BUILD & COMPILE GRAPH ===
-# agent.py
 class EsiAgent:
     def __init__(self):
         # 1. Build & compile your graph
         builder = StateGraph(GaiaAgentState)
 
-        # builder.add_node("parse",       input_parser)
+        builder.add_node("parse",       input_parser)
         builder.add_node("planner",       planner_node)
         builder.add_node("execute",    tool_executor)
         builder.add_node("synthesize", synthesizer)
 
-        # builder.add_edge(START,   "parse")
-        # builder.add_edge("parse",   "planner")
-        builder.add_edge(START,   "planner")
+        builder.add_edge(START,   "parse")
+        builder.add_edge("parse",   "planner")
         builder.add_edge("planner",    "execute")
         builder.add_edge("execute", "synthesize")
         builder.add_edge("synthesize", END)
-        
         self.graph = builder.compile()
 
 
@@ -462,75 +475,18 @@ class EsiAgent:
         files:    optional list of file paths/URLs (default empty list)
         level:    GAIA difficulty (1,2,3)
         """
-        # 1) Build the raw_payload
-        # print(11111111111, question)
+        
         payload = {
-            "task_id": None,
-            "question": question,
-            "files":    files or [],
-            "level":    level
+          "task_id": None,
+          "question": question,
+          "files":    [],
+          "level":    1
         }
-        # initial = {
-        #       "question": question,
-        #       "files":    [],       # or your file list
-        #       "level":    1,
-        #       "task_id": 1,
-        #       "plan":     [],       # planner_node will overwrite this 
-        #       "tool_results": []    # executor will overwrite this
-        #       "final_answer": ""
-        #     }
-
-        # print(2222222222, payload)
-
-        # 2) Wrap in your custom state class
         state = GaiaAgentState(raw_payload=payload)
-        # state = GaiaAgentState(dict)
-        
-        # print(6666666666666, state)
-        # 3) Invoke the graph end-to-end
         final_state = self.graph.invoke(state)
 
         # 4) Pull out the final answer
         return final_state.final_answer
-
-
-        # final_state = graph.invoke({
-        #     "task_id": None,
-        #     "question": question,
-        #     "files": [],
-        #     "level": 1,
-        #     "plan": [],
-        #     "tool_results": [],
-        #     # plus any other initial keys your nodes expect
-        # })
-
-        
-    # def run(self, question: str) -> str:
-    #     """
-    #     Runs the agent end‑to‑end on a single question string and returns
-    #     just the answer text (stripped of any prefix).
-    #     """
-    #     # Wrap in the same state shape your graph expects
-    #     payload = {
-    #         "task_id": None,
-    #         "question": question,
-    #         "files": [],      # no attachments
-    #         "level": 1        # default to Level 1
-    #     }
-    #     # 2) Wrap in your state object
-    #     initial_state = GaiaAgentState(raw_payload=payload)
-        
-    #     # initial_state = {"messages": [HumanMessage(content=question)]}
-    #     # Invoke the graph
-    #     final_state = self.graph.invoke(initial_state)
-    #     # Extract the last message
-    #     msgs = final_state["messages"]
-    #     last = msgs[-1].content
-    #     # If you prefix with "FINAL ANSWER: " remove that
-    #     if last.startswith("FINAL ANSWER:"):
-    #         return last[len("FINAL ANSWER:"):].strip()
-    #     return last.strip()
-
 # Optional: allow running agent.py directly for quick tests
 if __name__ == "__main__":
     agent = EsiAgent()