Update agent.py

agent.py

@@ -2,11 +2,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, create_react_agent
+from langchain.agents import create_tool_calling_agent, ConversationalAgent, AgentExecutor, initialize_agent, create_react_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 from langchain_core.tools import tool
 from langchain_core.messages import HumanMessage, SystemMessage
+from langchain.memory import ConversationBufferMemory
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 from langgraph.prebuilt import ToolNode
 from langgraph.prebuilt import tools_condition
@@ -237,19 +238,7 @@ def divide(a: float, b: float) -> float:
 # - subtract: Subtract A by B with passing A as the first argument
 # - divide: Divide A by B with passing A as the first argument
 
-# You have access to the following tools:
-# - serper_websearch: web search the content of the query by passing the query as input with Serper Search Engine
-# - duckduck_websearch: web search the content of the query by passing the query as input with DuckDuckGo Search Engine
-# - visit_webpage: visit the given webpage url by passing the url as input
-# - wiki_search: wiki search the content of the query by passing the query as input if the question asks for wiki search it
-# - text_splitter: split text into chunks
-# - youtube_transcript: fetch the transcript of the Youtube video by passing the video url as input if the question asks for watching a Youtube video
-# - read_file: read the content of the attached file by passing the TASK-ID as input
-# - excel_read: read the content of the attached excel file by passing the TASK-ID as input
-# - csv_read: read the content of the attached csv file by passing the TASK-ID as input
-# - mp3_listen: listen to the content of the attached mp3 file by passing the TASK-ID as input
-# - image_caption: understand the visual content of the attached image by passing the TASK-ID as input
-# - run_python: run the python code
+
 
 # ("human", f"Question: {question}\nReport to validate: {final_answer}")
 class BasicAgent:
@@ -257,21 +246,10 @@ class BasicAgent:
         self.model = ChatGoogleGenerativeAI(
             model="gemini-2.0-flash-lite",
             temperature=0,
-            max_tokens=128,
-            timeout=None,
-            max_retries=2,
+            max_tokens=1024,
+            candidate_count=1,
             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: 
@@ -282,48 +260,44 @@ class BasicAgent:
                 If you are asked for a comma separated list, apply the above rules depending of whether the element to put in the list is a number or a string.
 
                 You have access to the following tools:
-                {tools}
-                Here are the tools you can use: {tool_names}
+                - serper_websearch: web search the content of the query by passing the query as input with Serper Search Engine
+                - duckduck_websearch: web search the content of the query by passing the query as input with DuckDuckGo Search Engine
+                - visit_webpage: visit the given webpage url by passing the url as input
+                - wiki_search: wiki search the content of the query by passing the query as input if the question asks for wiki search it
+                - text_splitter: split text into chunks
+                - youtube_transcript: fetch the transcript of the Youtube video by passing the video url as input if the question asks for watching a Youtube video
+                - read_file: read the content of the attached file by passing the TASK-ID as input
+                - excel_read: read the content of the attached excel file by passing the TASK-ID as input
+                - csv_read: read the content of the attached csv file by passing the TASK-ID as input
+                - mp3_listen: listen to the content of the attached mp3 file by passing the TASK-ID as input
+                - image_caption: understand the visual content of the attached image by passing the TASK-ID as input
+                - run_python: run the python code
+                
                 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]
                 Note: python_tool is called when the question mentions the term "Python" or any math calculation.
-
-                Follow this format in your response:
-                THOUGHT: [Describe your reasoning here]
-                ACTION: [Specify the action/tool to use and any relevant input]
-                OBSERVATIOn: [Result of the action/tool, provided by the system]
-                FINAL ANSWER: [Provide your final response to the user]
-
-                User Input: {input}
-                {agent_scratchpad}
         """
         self.tools = [duckduck_websearch, serper_websearch, visit_webpage, wiki_search, text_splitter, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, run_python]
-        # self.model_with_tools = self.model.bind_tools(self.tools)
-        # self.sys_msg = SystemMessage(content=self.sys_prompt)
-        
+        # Setup memory
+        self.memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+        self.agent = self.__setup__agent__()
         # self.prompt = ChatPromptTemplate.from_messages([
         #     ("system", self.sys_prompt),
         #     ("human", "{input}")
         # ])
-        self.prompt = PromptTemplate(
-            input_variables=["input", "tools", "tool_names", "agent_scratchpad"],
-            template=self.sys_prompt
-        )
+
         # self.agent = initialize_agent(
         #     tools=self.tools,
         #     llm=self.model,
         #     agent="zero-shot-react-description",  # ReAct agent type
         #     verbose=True,
         #     system_prompt=self.prompt,
-        #     handle_parsing_errors="Check your output and make sure it conforms, use the Action/Action Input syntax"
+        #     handle_parsing_errors=True,
+        #     max_iterations=30
+        #     # "Check your output and make sure it conforms, use the Action/Action Input syntax"
         # )
-        self.agent = create_react_agent(
-            llm=self.model,
-            tools=self.tools,
-            prompt=self.prompt
-        )
-        self.agent_exe = AgentExecutor(agent=self.agent, tools=self.tools, verbose=True,
-                                      handle_parsing_errors="Check your output and make sure it conforms, use the Action/Action Input syntax")
-        # self.graph = self.__graph_compile__()
         print("BasicAgent initialized.")
     
     def __call__(self, task: dict) -> str:
@@ -335,33 +309,71 @@ class BasicAgent:
         else:
             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)
-        fixed_answer = self.agent_exe.invoke({"input": question})
-        # human_message = [HumanMessage(content=question)]
-        # messages = self.graph.invoke({"messages": human_message})
-        # fixed_answer = messages['messages'][-1].content
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        time.sleep(60)
+        fixed_answer = "This is a default answer."
+        
+
+        max_retries = 3
+        base_sleep = 1
+        for attempt in range(max_retries):
+            try:
+                fixed_answer = self.agent.run(question)
+                print(f"Agent returning fixed answer: {fixed_answer}")
+                time.sleep(60)
+                return fixed_answer
+            except Exception as e:
+                sleep_time = base_sleep * (attempt + 1)
+                if attempt < max_retries - 1:
+                    print(f"Attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
+                    time.sleep(sleep_time)
+                    continue
+                return f"Error processing query after {max_retries} attempts: {str(e)}"
         return fixed_answer
 
-    def __graph_compile__(self):
-        def assistant(state: MessagesState):
-            """Assistant Node"""
-            return {"message": [self.model_with_tools.invoke(state["messages"])]}
-        
-        builder = StateGraph(MessagesState)
-        builder.add_node("assistant", 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")
-        # Compile graph
-        return builder.compile()
+    def __setup__agent__(self) -> AgentExecutor:
+        PREFIX = """
+                You are a general AI assistant that can use various tools to answer question. 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 separared 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 (eg. 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 put in the list is a number or a string.
 
-    
+                NOTE:
+                - 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]
+                - python_tool is called when the question mentions the term "Python" or any math calculation.
+        """
+        FORMAT_INSTRUCTION = """
+                To use a tool, use the following format:
+                Thought: Do I need to use a tool? Yes
+                Action: the action to take, should be one of [{tool_names}]
+                Action Input: the input to the action
+                Observation: the result of the action
+                When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
+                Thought: Do I need to use a tool? No
+                Final Answer: [your response here]
+                Begin! Remember to ALWAYS include 'Thought:', 'Action:', 'Action Input:', and 'Final Answer:' in your responses.
+        """
+        SUFFIX = """
+                Previous conversation history:
+                {chat_history}
+                New question: {input}
+                {agent_scratchpad}
+        """
+        agent = ConversationalAgent.from_llm_and_tools(
+            llm=self.model,
+            tools=self.tools,
+            prefix=PREFIX,
+            format_instructions=FORMAT_INSTRUCTIONS,
+            suffix=SUFFIX,
+            input_variables=["input", "chat_history", "agent_scratchpad", "tool_names"],
+            handle_parsing_errors=True
+        )
+        return AgentExecutor.from_agent_and_tools(
+            agent=agent,
+            tools=self.tools,
+            memory=self.memory,
+            max_iterations=5,
+            verbose=True,
+            handle_parsing_errors=True,
+            return_only_outputs=True  # This ensures we only get the final output
+        )