Update agent.py

agent.py

@@ -1,7 +1,7 @@
 import os
 from typing import TypedDict, List, Dict, Any, Optional
 from langgraph.graph import StateGraph, START, END
-from langchain_anthropic import ChatAnthropic
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.tools import tool
 from langchain_core.messages import HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
@@ -79,20 +79,27 @@ def image_caption(dir: str) -> str:
 # 2. Coding
 # 3. Multi-Modality
 
-
-
 # ("human", f"Question: {question}\nReport to validate: {final_answer}")
 class BasicAgent:
     def __init__(self):
-        model = ChatAnthropic(
-            model="claude-3-5-sonnet-20240620",
+        model = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
             temperature=0,
-            max_tokens=20000,
+            max_tokens=None,
             timeout=None,
             max_retries=2,
-            api_key=os.getenv("ANTHROPIC_API_KEY"),
+            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].