out_state

agent.py

@@ -18,9 +18,12 @@ from tools import (
 )
 
 # ─────────────────────────── Configuration ───────────────────────────────
-
+SYSTEM_PROMPT = """
+You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+"""
 MAX_TOOL_CALLS = 5
 
+
 # ─────────────────────────── Helper utilities ────────────────────────────
 
 # ─────────────────────────── Agent state ⬇ ───────────────────────────────
@@ -33,7 +36,7 @@ MAX_TOOL_CALLS = 5
 # ─────────────────────────── Graph wiring ───────────────────────────────
 
 def build_graph():
-    graph = StateGraph(AgentState)
+    
 
     llm = ChatOpenAI(model_name="gpt-4.1-mini", temperature=0.3)
 
@@ -45,6 +48,6 @@ def build_graph():
         analyze_code_tool,
     ]
     # llm = llm.bind_tools(llm_tools)
-    agent = create_react_agent(llm, llm_tools)
+    agent = create_react_agent(model=llm, tools=llm_tools, prompt=SYSTEM_PROMPT)
 
     return agent 

app.py

@@ -33,12 +33,9 @@ class BasicAgent:
 
         # IMPORTANT: invoke() returns a **new** state instance (or an AddableValuesDict),
         # not the object we pass in.  Use the returned value to fetch final_answer.
-        out_state = self.graph.invoke(init_state)
-
-        if isinstance(out_state, dict):        # AddableValuesDict behaves like a dict
-            return out_state.get("final_answer", "No answer.")
-        else:                                  # If future versions return the dataclass
-            return getattr(out_state, "final_answer", "No answer.")
+        out_state = self.graph.invoke({"messages": [HumanMessage(content=question)]})
+        print("out_state: ", out_state)
+        return out_state
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):

state.py

@@ -1,24 +1,16 @@
-from dataclasses import dataclass, field
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, TypedDict, Annotated
 import json
-from dataclasses import dataclass, field, asdict
 from langchain.schema import SystemMessage, HumanMessage, AIMessage, BaseMessage
+from langgraph.graph.message import add_messages
 
 
-@dataclass
-class AgentState:
+class AgentState(TypedDict):
     """Single source‑of‑truth context for one user query run."""
-
     user_question: str
-    task_id: Optional[str] = None
-    messages: List[BaseMessage] = field(default_factory=list)
-
-    next_action: Optional[str] = None      # wiki | ocr | audio | final
-    query: Optional[str] = None            # wiki search term
-    snippet: Optional[str] = None          # code snippet
-    tool_calls: int = 0
-
-    final_answer: Optional[str] = None
-
-    def add(self, *msgs: BaseMessage):
-        self.messages.extend(msgs)
+    task_id: Optional[str] 
+    messages: Annotated[List[BaseMessage], add_messages]
+    next_action: Optional[str]      # wiki | ocr | audio | final
+    query: Optional[str]            # wiki search term
+    snippet: Optional[str]          # code snippet
+    tool_calls: int
+    final_answer: Optional[str]