Update agent.py

agent.py

@@ -12,6 +12,7 @@ from typing import TypedDict,Annotated
 
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
+    input_file: Optional[str]  # Contains file path (PDF/PNG)
 
 
 class BasicAgent():
@@ -41,26 +42,41 @@ class BasicAgent():
         # Compile the graph
         self.agent = builder.compile()
 
+    
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        agent_prompt = f"""        
-        When providing the final answer, be as concise as possible, provide only the FINAL ANSWER. For example:
-        QUESTION: What was the actual enrollment count of the clinical trial on H. pylori in acne vulgaris patients from Jan-May 2018 as listed on the NIH website?
-        FINAL ANSWER: 90
-
-        Now:
-        QUESTION: {question}
-        FINAL ANSWER: """
-        
-        messages=[HumanMessage(content=agent_prompt)]
+        messages=[HumanMessage(content=question)]
         response = self.agent.invoke({"messages":messages})
         answer = response['messages'][-1].content
         print(f"Agent returning answer: {answer}")
         return answer
 
+    
     def assistant(self, state: AgentState):
+        image = state["input_file"]
+        textual_description_of_tool="""
+        extract_text(img_path: str) -> str:
+            Extract text from an image file using a multimodal model.
+        
+            Args:
+                img_path: A local image file path (strings).
+        
+            Returns:
+                A single string containing the concatenated text extracted from each image.
+        search_tool(query: str) -> str:
+            Search the web using the DuckDuckGoSearchRun to perform a search query and return a summarized textual result.
+
+            Args:
+                query: A string representing the search query.
+        
+            Returns:
+                A single string containing the search result or summary.
+        """        
+        sys_msg = SystemMessage(content=f"You are a helpful assistant. You have access to the following tools:\n{textual_description_of_tool} \n You may have access to some optional images. Currently the loaded image is: {image}. \n When providing the final answer, be as concise as possible, provide only the FINAL ANSWER. \n For example: QUESTION: What was the actual enrollment count of the clinical trial on H. pylori in acne vulgaris patients from Jan-May 2018 as listed on the NIH website? \n FINAL ANSWER: 90 \n")
+
         return {
-            "messages": [self.chat_with_tools.invoke(state["messages"])],
+            "messages": [self.chat_with_tools.invoke([sys_msg] + state["messages"])],            
+            "input_file": image
         }