Update agent.py

agent.py

@@ -12,7 +12,7 @@ class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
 
-class Agent():
+class BasicAgent():
     def __init__(self, llm):
         chat = ChatHuggingFace(llm=llm, verbose=True)
         
@@ -21,7 +21,9 @@ class Agent():
 
         chat_with_tools = chat.bind_tools([extract_text,search_tool])
         self._initialize_graph()
+        print("BasicAgent initialized.")
 
+    
     def _initialize_graph(self):
         builder = StateGraph(AgentState)
 
@@ -37,14 +39,16 @@ class Agent():
         # Compile the graph
         self.agent = builder.compile()
 
-    def call_agent(self, messages):
-        self.agent.invoke({"messages":messages})
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        response = self.agent.invoke({"messages":question})
+        answer = response['messages'][-1].content)
+        print(f"Agent returning answer: {answer}")
+        return answer
 
     def assistant(state: AgentState):
         return {
             "messages": [chat_with_tools.invoke(state["messages"])],
         }
 
-
-