Update agent.py

agent.py

@@ -24,9 +24,7 @@ class BasicAgent():
         self.chat_with_tools = chat.bind_tools(self.tools)
         self._initialize_graph()
         self._initialize_telemetry()
-        print("BasicAgent initialized.")
 
-    
     def _initialize_graph(self):
         builder = StateGraph(AgentState)
 
@@ -41,6 +39,7 @@ class BasicAgent():
 
         # Compile the graph
         self.agent = builder.compile()
+        print("Agent initialized.")
 
     def _initialize_telemetry(self):
         LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY")
@@ -52,33 +51,32 @@ class BasicAgent():
             secret_key=LANGFUSE_SECRET_KEY,
             host=LANGFUSE_HOST
         )
+        print("Telemetry initialized.")
 
-    
     def __call__(self, question: str, file_name : str) -> str:
-        print(f"Agent received question: {question}. /nProvided file: {file_name}.")
+
+        sys_msg = SystemMessage(content=f"""
+            You are a general AI assistant. I will ask you a question. Reason step by step and search for the information you need using available tools.
+            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 separated 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 (e.g. 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 be put in the list is a number or a string.
+            When providing the final answer, ONLY give [YOUR FINAL ANSWER]. Do not add anything else, no additional motivation or explanation, and do not return 'FINAL ANSWER:'.
+            """)
+        
+        print(f"Agent received question: {question}.")
+        
         if file_name is not None and file_name!='':
-            messages=[HumanMessage(content=f"{question}. The filename you have access to is {file_name}.")]
+            print(f"Provided file: {file_name}.")
+            messages=[[sys_msg] + HumanMessage(content=f"{question}. The file you have access to is {file_name}.")]
         else:
-            messages=[HumanMessage(content=question)]
+            messages=[[sys_msg] + HumanMessage(content=question)]
+            
         response = self.agent.invoke({"messages":messages}, config={"callbacks": [self.langfuse_handler]})
         answer = response['messages'][-1].content
         print(f"Agent returning answer: {answer}")
         return answer
 
-    
     def assistant(self, state: AgentState):
- 
-        sys_msg = SystemMessage(content=f"""
-        You are a general AI assistant. I will ask you a question. Reason step by step and search for the information you need using available tools, one step at a time.
-        If you do not have enough information to answer, use the tools available to search for it. 
-        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 separated 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 (e.g. 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 be put in the list is a number or a string.
-        When providing the final answer, ONLY give [YOUR FINAL ANSWER]. Do not add anything else, no additional motivation or explanation, and do not return 'FINAL ANSWER:'.
-        """)
-        response = self.chat_with_tools.invoke([sys_msg] + state["messages"])
- 
+        response = self.chat_with_tools.invoke(state["messages"])
         return {
             "messages":  state["messages"] + [response],            
         }
-
-