Enhance OracleBot: improve message handling and output formatting in answer_question method

agent/agent.py

@@ -4,9 +4,10 @@ from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import tools_condition
 from agent.nodes import call_model, tool_node
 from langgraph.graph import MessagesState
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage, AIMessageChunk
 from langgraph.checkpoint.memory import InMemorySaver
 from agent.config import create_agent_config
+from termcolor import colored, cprint
 
 class OracleBot:
     def __init__(self):
@@ -21,8 +22,40 @@ class OracleBot:
         Answer a question using the LangGraph agent.
         """
         messages = [HumanMessage(content=question)]
-        for mode, chunk in self.graph.stream({"messages": messages}, config=self.config, stream_mode=["messages", "custom"]): # type: ignore
-            print(chunk)
+
+        for mode, chunk in self.graph.stream({"messages": messages}, config=self.config, stream_mode=["messages", "updates"]): # type: ignore
+            if mode == "messages":
+                if isinstance(chunk, tuple) and len(chunk) > 0:
+                    message = chunk[0]
+                    if isinstance(message, (AIMessageChunk, AIMessage)):
+                        # Only print chunks that have actual content (skip tool call chunks)
+                        if hasattr(message, 'content') and message.content and not (hasattr(message, 'tool_calls') and message.tool_calls):
+                            cprint(message.content, color="light_grey", attrs=["dark"], end="", flush=True)
+                # Handle case where chunk is directly the message
+                elif isinstance(chunk, (AIMessageChunk, AIMessage)):
+                    # Only print chunks that have actual content (skip tool call chunks)
+                    if hasattr(chunk, 'content') and chunk.content and not (hasattr(chunk, 'tool_calls') and chunk.tool_calls):
+                        cprint(chunk.content, color="light_grey", attrs=["dark"], end="", flush=True)
+            elif mode == "updates":
+                # Look for complete tool calls in updates
+                if isinstance(chunk, dict) and 'agent' in chunk:
+                    agent_update = chunk['agent']
+                    if 'messages' in agent_update and agent_update['messages']:
+                        for message in agent_update['messages']:
+                            if hasattr(message, 'tool_calls') and message.tool_calls:
+                                for tool_call in message.tool_calls:
+                                    cprint(f"\n🔧 Using tool: {tool_call['name']} with args: {tool_call['args']}\n", color="yellow")
+                            # Handle final answer messages (no tool calls)
+                            elif hasattr(message, 'content') and message.content:
+                                cprint(f"\n{message.content}\n", color="black", on_color="on_white", attrs=["bold"])
+                
+                # Look for tool outputs in updates
+                elif isinstance(chunk, dict) and 'tools' in chunk:
+                    tools_update = chunk['tools']
+                    if 'messages' in tools_update and tools_update['messages']:
+                        for message in tools_update['messages']:
+                            if hasattr(message, 'content') and message.content:
+                                cprint(f"\n📤 Tool output:\n{message.content}\n", color="green")
 
     def _build_agent(self, name: str):
         """
@@ -56,11 +89,10 @@ if __name__ == "__main__":
     question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
     
     try:
-        from logger.phoenix import start_phoenix
+        from config import start_phoenix
         start_phoenix()
         bot = OracleBot()
         bot.answer_question(question)
 
     except Exception as e:
-        print(f"Error running agent: {e}")
-        print("Make sure your API key is correct and the service is accessible.")
+        print(f"Error running agent: {e}")