Update agent.py

agent.py

@@ -1,18 +1,11 @@
 import os
 from typing import TypedDict, List, Dict, Any, Optional
-from langgraph.prebuilt import create_react_agent
-from langgraph.graph import StateGraph, START, END
 from langchain.agents import create_tool_calling_agent, AgentExecutor
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.tools import tool
 from langchain_core.messages import HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
 
-class AgentState(TypedDict):
-    messages: List
-    current_question: str
-    final_answer: str
-
 # 1. Web Browsing
 from langchain_community.tools import DuckDuckGoSearchRun
 from langchain_community.document_loaders import ImageCaptionLoader
@@ -78,7 +71,7 @@ def image_caption(dir: str) -> str:
 # ("human", f"Question: {question}\nReport to validate: {final_answer}")
 class BasicAgent:
     def __init__(self):
-        model = ChatGoogleGenerativeAI(
+        self.model = ChatGoogleGenerativeAI(
             model="gemini-2.0-flash",
             temperature=0,
             max_tokens=1024,
@@ -87,15 +80,6 @@ class BasicAgent:
             google_api_key="AIzaSyAxVUPaGJIgdxB46ZR0RWPKSjB9a63Z80o",
             # other params...
         )
-        # model = ChatAnthropic(
-        #     model="claude-3-5-sonnet-20240620",
-        #     temperature=0,
-        #     max_tokens=20000,
-        #     timeout=None,
-        #     max_retries=2,
-        #     api_key=os.getenv("ANTHROPIC_API_KEY"),
-        #     # other params...
-        # )
         # System Prompt for few shot prompting
         self.sys_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].
@@ -112,17 +96,12 @@ class BasicAgent:
             ("system", self.sys_prompt),
             ("human", "Question: {input}")
         ])
-        self.agent = create_tool_calling_agent(model, tools, prompt)
+        self.agent = create_tool_calling_agent(self.model, tools, prompt)
         self.agent_exe = AgentExecutor(agent=self.agent, tools=tools, verbose=True)
         print("BasicAgent initialized.")
     
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        prompt_msg = [
-            ("system", self.sys_prompt),
-            ("human", f"Question: {question}")
-        ]
-        # response = self.model.invoke(prompt_msg)
         response = self.agent_exe.invoke({"input": question})
         fixed_answer = response['message'][-1].content
         # fixed_answer = "This is a default answer."