fix AgentState

agent.py

@@ -1,8 +1,33 @@
+import os
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.checkpoint.memory import MemorySaver
+from langchain.globals import set_debug
+from langchain.globals import set_verbose
+from langgraph.prebuilt import create_react_agent
+from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import AgentState
 
-class KateState(AgentState):
-    task_id: str
-    question: str
-    file_name: str
+from smolagents import DuckDuckGoSearchTool
+from smolagents import PythonInterpreterTool
+from tools import analyze_audio
+from tools import analyze_excel
+from tools import analyze_image
+from tools import analyze_video
+from tools import download_file_for_task
+from tools import read_file_contents
+from tools import search_arxiv
+from tools import search_tavily
+from tools import search_wikipedia
+from tools import SmolagentToolWrapper
+from tools import tavily_extract_tool
+from utils import get_llm
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
+if not GOOGLE_API_KEY:
+    raise ValueError("GOOGLE_API_KEY environment variable is not set.")
+
+AGENT_MODEL_NAME = os.getenv("AGENT_MODEL_NAME", "gemini-2.0-flash")
 
 KATE_AGENT_PROMPT = """
 You are Kate's Advanced AI Assistant designed to solve complex tasks efficiently.
@@ -25,19 +50,79 @@ QUESTION: Capital of France? FINAL ANSWER: Paris
 QUESTION: 2+2? FINAL ANSWER: 4
 """
 
+class KateState(AgentState):
+    task_id: str
+    question: str
+    file_name: str
+
+
 class KateMultiModalAgent:
     def __init__(self, model_name: str | None = None):
-        self.model_name = model_name or AGENT_MODEL_NAME
-        self.llm = self._get_llm()
-        self.tools = self._get_tools()
-        self.agent = self._create_agent()
-        
-    def _create_agent(self):
-        return create_react_agent(
-            self.llm,
-            tools=self.tools,
+        if model_name is None:
+            model_name = AGENT_MODEL_NAME
+        llm = self._get_llm(model_name)
+        tools = self._get_tools()
+        self.agent = create_react_agent(
+            llm,
+            tools=tools,
             state_schema=KateState,
             state_modifier=KATE_AGENT_PROMPT,
-            checkpointer=MemorySaver(),
-            max_iterations=10
+            checkpointer=MemorySaver()
+        )
+
+    def _get_llm(self, model_name: str):
+        return get_llm(
+            llm_provider_api_key=GOOGLE_API_KEY,
+            model_name=model_name,
+        )
+
+    def _get_tools(self):
+        tools = [
+            SmolagentToolWrapper(DuckDuckGoSearchTool()),
+            SmolagentToolWrapper(PythonInterpreterTool()),
+            download_file_for_task,
+            read_file_contents,
+            analyze_audio,
+            analyze_image,
+            analyze_excel,
+            analyze_video,
+            search_arxiv,
+            search_tavily,
+            search_wikipedia,
+            tavily_extract_tool,
+        ]
+        return ToolNode(tools)
+
+    async def __call__(
+        self, task_id: str, question: str, file_name: str | None = None
+    ) -> str:
+        config = RunnableConfig(
+            recursion_limit=64,
+            configurable={"thread_id": task_id}
         )
+
+        if not file_name:
+            file_name = "None - no file present"
+
+        message = HumanMessage(
+            content=[
+                {
+                    "type": "text",
+                    "text": f"Task Id: {task_id}, Question: {question}, Filename: {file_name}. If a filename is present (and is not 'None'), download the file for the task that's referenced in the question. If there isn't a filename present, please use tools where applicable."
+                }
+            ]
+        )
+
+        answer = await self.agent.ainvoke(
+            {
+                "messages": [message],
+                "question": question,
+                "task_id": task_id,
+                "file_name": file_name
+            }, config)
+
+        final_answer = answer['messages'][-1].content
+        if "FINAL ANSWER: " in final_answer:
+            return final_answer.split("FINAL ANSWER: ", 1)[1].strip()
+        else:
+            return final_answer

app.py

@@ -5,6 +5,8 @@ import requests
 import pandas as pd
 from agent import KateMultiModalAgent
 
+agent = KateMultiModalAgent()
+
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 AGENT_NAME = "Kate's Advanced Agent"