debug

app.py

@@ -18,6 +18,12 @@ from langchain_core.tools import tool
 from utilities import get_file
 import time
 from tenacity import retry, stop_after_attempt, wait_exponential
+import json
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Constants
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
@@ -104,6 +110,30 @@ def google_search(query: str) -> str:
     response = llm.invoke(query, tools=[GenAITool(google_search={})])
     return response.content
 
+def log_message(message: BaseMessage, prefix: str = ""):
+    """Helper function to log message details."""
+    if isinstance(message, HumanMessage):
+        logger.info(f"{prefix}Human Message:")
+        if isinstance(message.content, list):
+            for item in message.content:
+                if isinstance(item, dict):
+                    if item.get("type") == "media":
+                        logger.info(f"{prefix}  Media content (type: {item.get('mime_type')})")
+                    else:
+                        logger.info(f"{prefix}  {item.get('type')}: {item.get('text')}")
+                else:
+                    logger.info(f"{prefix}  {item}")
+        else:
+            logger.info(f"{prefix}  {message.content}")
+    elif isinstance(message, AIMessage):
+        logger.info(f"{prefix}AI Message:")
+        logger.info(f"{prefix}  Content: {message.content}")
+        if hasattr(message, 'tool_calls') and message.tool_calls:
+            logger.info(f"{prefix}  Tool Calls: {json.dumps(message.tool_calls, indent=2)}")
+    elif isinstance(message, SystemMessage):
+        logger.info(f"{prefix}System Message:")
+        logger.info(f"{prefix}  {message.content}")
+
 class BasicAgent:
     def __init__(self):
         self.llm = ChatGoogleGenerativeAI(
@@ -144,24 +174,32 @@ class BasicAgent:
         # Compile the graph
         self.app = self.workflow.compile()
         
-        print("BasicAgent initialized.")
+        logger.info("BasicAgent initialized.")
     
     @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=60))
     def call_model(self, state: AgentState) -> AgentState:
         """Call the model to generate a response with retry logic."""
         try:
             messages = state["messages"]
+            logger.info("\n=== Model Input ===")
+            log_message(self.sys_msg, "  ")
+            for msg in messages:
+                log_message(msg, "  ")
+            
             response = self.llm.invoke([self.sys_msg] + messages)
             
+            logger.info("\n=== Model Output ===")
+            log_message(response, "  ")
+            
             if not response or not response.content:
-                print("Empty response from model, retrying...")
+                logger.error("Empty response from model")
                 raise ValueError("Empty response from model")
                 
             return {"messages": [response], "next": "tools"}
         except Exception as e:
-            print(f"Error in call_model: {str(e)}")
+            logger.error(f"Error in call_model: {str(e)}")
             if "429" in str(e):
-                print("Rate limit hit, waiting before retry...")
+                logger.warning("Rate limit hit, waiting before retry...")
                 time.sleep(60)  # Wait for 60 seconds before retry
             raise
     
@@ -171,6 +209,7 @@ class BasicAgent:
             messages = state["messages"]
             last_message = messages[-1]
             
+            logger.info("\n=== Tool Execution ===")
             if isinstance(last_message, AIMessage):
                 # Extract tool calls from the message
                 tool_calls = last_message.tool_calls
@@ -179,20 +218,29 @@ class BasicAgent:
                         try:
                             tool_name = tool_call.name
                             tool_args = tool_call.args
+                            logger.info(f"Executing tool: {tool_name}")
+                            logger.info(f"Tool arguments: {json.dumps(tool_args, indent=2)}")
+                            
                             result = self.tool_executor.invoke(tool_name, tool_args)
+                            logger.info(f"Tool result: {result}")
+                            
                             messages.append(AIMessage(content=f"Tool result: {result}"))
                         except Exception as e:
-                            print(f"Error executing tool {tool_name}: {str(e)}")
+                            logger.error(f"Error executing tool {tool_name}: {str(e)}")
                             messages.append(AIMessage(content=f"Tool error: {str(e)}"))
+                else:
+                    logger.info("No tool calls found in AI message")
             
             return {"messages": messages, "next": "agent"}
         except Exception as e:
-            print(f"Error in call_tools: {str(e)}")
+            logger.error(f"Error in call_tools: {str(e)}")
             return {"messages": messages, "next": "agent"}
     
     async def __call__(self, question: str, task_id: str) -> str:
         """Process a question and return the answer with error handling."""
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        logger.info(f"\n=== Processing Question ===")
+        logger.info(f"Task ID: {task_id}")
+        logger.info(f"Question: {question}")
         
         try:
             # Create initial state
@@ -208,12 +256,16 @@ class BasicAgent:
             
             while retry_count < max_retries:
                 try:
+                    logger.info(f"\n=== Attempt {retry_count + 1}/{max_retries} ===")
                     result = self.app.invoke(initial_state)
                     final_message = result["messages"][-1]
                     
                     if isinstance(final_message, AIMessage) and final_message.content:
+                        logger.info(f"\n=== Final Answer ===")
+                        logger.info(f"Answer: {final_message.content}")
                         return final_message.content
                     else:
+                        logger.error("Empty or invalid response")
                         raise ValueError("Empty or invalid response")
                         
                 except Exception as e:
@@ -221,17 +273,17 @@ class BasicAgent:
                     retry_count += 1
                     if "429" in str(e):
                         wait_time = 60 * retry_count
-                        print(f"Rate limit hit, waiting {wait_time} seconds before retry {retry_count}/{max_retries}")
+                        logger.warning(f"Rate limit hit, waiting {wait_time} seconds before retry {retry_count}/{max_retries}")
                         await asyncio.sleep(wait_time)
                     else:
-                        print(f"Error in processing, retry {retry_count}/{max_retries}: {str(e)}")
+                        logger.error(f"Error in processing, retry {retry_count}/{max_retries}: {str(e)}")
                         await asyncio.sleep(5)
             
-            print(f"All retries failed. Last error: {str(last_error)}")
+            logger.error(f"All retries failed. Last error: {str(last_error)}")
             return "Unable to generate answer after multiple attempts"
             
         except Exception as e:
-            print(f"Fatal error in agent: {str(e)}")
+            logger.error(f"Fatal error in agent: {str(e)}")
             return f"Error: {str(e)}"
 
 def run_and_submit_all(profile):