Update agent.py

agent.py

@@ -1,11 +1,12 @@
 import os
 from dotenv import load_dotenv
 from typing import TypedDict, List, Dict, Any, Optional
+from langgraph.graph import StateGraph, START, END, MessagesState
 from langchain.agents import create_tool_calling_agent, AgentExecutor, initialize_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 from langchain_core.tools import tool
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.prompts import ChatPromptTemplate
 
 # 1. Web Browsing
@@ -231,24 +232,24 @@ def divide(a: float, b: float) -> float:
 # ("human", f"Question: {question}\nReport to validate: {final_answer}")
 class BasicAgent:
     def __init__(self):
-        # self.model = ChatGoogleGenerativeAI(
-        #     model="gemini-2.0-flash-lite",
-        #     temperature=0,
-        #     max_tokens=128,
-        #     timeout=None,
-        #     max_retries=2,
-        #     google_api_key=os.getenv("GEMINI_API_KEY"),
-        #     # other params...
-        # )
-        self.model = ChatGroq(
-            model="qwen-qwq-32b",
+        self.model = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash-lite",
             temperature=0,
             max_tokens=128,
             timeout=None,
             max_retries=2,
-            groq_api_key=os.getenv("GROQ_API_KEY")
+            google_api_key=os.getenv("GEMINI_API_KEY"),
             # other params...
         )
+        # self.model = ChatGroq(
+        #     model="qwen-qwq-32b",
+        #     temperature=0,
+        #     max_tokens=128,
+        #     timeout=None,
+        #     max_retries=2,
+        #     groq_api_key=os.getenv("GROQ_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: 
@@ -275,6 +276,9 @@ class BasicAgent:
                 If Task ID is included in the question, remember to call the relevant read tools [ie. read_file, excel_read, csv_read, mp3_listen, image_caption]
         """
         self.tools = [duckduck_websearch, serper_websearch, visit_webpage, wiki_search, text_splitter, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, python_tool]
+        self.binded_model = self.model.bind_tools(self.tools)
+        self.sys_msg = SystemMessage(content=self.sys_prompt)
+        
         self.prompt = ChatPromptTemplate.from_messages([
             ("system", self.sys_prompt),
             ("human", "{input}")
@@ -287,6 +291,7 @@ class BasicAgent:
             system_prompt=self.prompt,
             handle_parsing_errors="Check your output and make sure it conforms, use the Action/Action Input syntax"
         )
+        self.app = self.__graph_compile__()
         print("BasicAgent initialized.")
     
     def __call__(self, task: dict) -> str:
@@ -296,10 +301,33 @@ class BasicAgent:
         # fixed_answer = response['message'][-1].content
         
         if file_name == "" or file_name is None:
-            fixed_answer = self.agent.run(question)
+            question = question
         else:
-            fixed_answer = self.agent.run(f'{question} with TASK-ID: {task_id}')
+            question = f"{question} with TASK-ID: {task_id}"
+            # fixed_answer = self.agent.run(f'{question} with TASK-ID: {task_id}')
         # fixed_answer = "This is a default answer."
+        # fixed_answer = self.agent.run(question)
+        human_message = [HumanMessage(content=question)]
+        messages = self.app.invoke({"message": human_message})
+        fixed_answer = messages['messages'][-1].content
         print(f"Agent returning fixed answer: {fixed_answer}")
         time.sleep(60)
-        return fixed_answer
+        return fixed_answer
+
+    def __graph_compile__(self):
+        builder = StateGraph(MessagesState)
+        builder.add_node("assistant", self.assistant)
+        builder.add_node("tools", ToolNode(self.tools))
+        builder.add_edge(START, "assistant")
+        builder.add_conditional_edges(
+            "assistant",
+            # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+            # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
+            tools_condition,
+        )
+        builder.add_edge("tools", "assistant")
+        return builder.compile() # Compile graph
+
+    def assistant(self, state: MessagesState):
+        """Assistant Node"""
+        return {"message": [self.binded_model.invoke([self.sys_msg] + state["message"])]}