Update app.py

app.py

@@ -6,6 +6,8 @@ import pandas as pd
 import ast
 import operator
 import time
+import json
+from datetime import datetime
 from typing import List, Dict, Any, Optional, Annotated
 from langgraph.graph import Graph, StateGraph
 from langgraph.prebuilt import ToolNode
@@ -21,6 +23,26 @@ print("trial")
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # Make sure to set this environment variable
 
+# Create logs directory if it doesn't exist
+LOGS_DIR = "question_logs"
+os.makedirs(LOGS_DIR, exist_ok=True)
+
+def log_to_file(task_id: str, question: str, log_data: Dict[str, Any]):
+    """Store logs for a question in a JSON file."""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"{LOGS_DIR}/question_{task_id}_{timestamp}.json"
+    
+    log_entry = {
+        "task_id": task_id,
+        "question": question,
+        "timestamp": timestamp,
+        "logs": log_data
+    }
+    
+    with open(filename, 'w', encoding='utf-8') as f:
+        json.dump(log_entry, f, indent=2, ensure_ascii=False)
+    
+    print(f"Logs saved to {filename}")
 
 class AgentState(TypedDict):
     question: Annotated[str, override]
@@ -30,6 +52,8 @@ class AgentState(TypedDict):
     history: Annotated[List[Dict[str, str]], operator.add]
     needs_more_info: Annotated[bool, override]
     search_query: Annotated[str, override]
+    task_id: Annotated[str, override]  # Add task_id to state
+    logs: Annotated[Dict[str, Any], operator.add]  # Add logs to state
 
 class BasicAgent:
     def __init__(self):
@@ -112,14 +136,11 @@ class BasicAgent:
         """Analyze the question and determine the next step."""
         prompt = f"""Analyze this question and determine what needs to be done: {state['question']}
 
-Return only a valid Python dictionary in this exact format:
+Return ONLY a Python dictionary in this exact format, with no other text or explanation:
 {{
     "needs_search": true/false,
     "search_query": "query if needed"
-}}
-
-Do not include any other text or explanation. Only return the dictionary.
-"""
+}}"""
         
         try:
             llm_response = self._call_llm_api(prompt)
@@ -127,6 +148,13 @@ Do not include any other text or explanation. Only return the dictionary.
             print(f"Input: {state['question']}")
             print(f"LLM Response: {llm_response}")
             
+            # Log the analysis step
+            state["logs"]["analyze"] = {
+                "prompt": prompt,
+                "response": llm_response,
+                "timestamp": datetime.now().isoformat()
+            }
+            
             analysis = ast.literal_eval(llm_response)
             state["needs_more_info"] = analysis.get('needs_search', False)
             state["search_query"] = analysis.get('search_query', '')
@@ -141,6 +169,12 @@ Do not include any other text or explanation. Only return the dictionary.
             state["needs_more_info"] = True
             state["search_query"] = state["question"]
             state["current_step"] = 'search'
+            
+            # Log the error
+            state["logs"]["analyze_error"] = {
+                "error": str(e),
+                "timestamp": datetime.now().isoformat()
+            }
         
         return state
 
@@ -162,6 +196,13 @@ Do not include any other text or explanation. Only return the dictionary.
             for i, result in enumerate(search_results, 1):
                 print(f"{i}. {result}")
             
+            # Log the search step
+            state["logs"]["search"] = {
+                "query": state["search_query"],
+                "results": search_results,
+                "timestamp": datetime.now().isoformat()
+            }
+            
             state["history"].append({
                 'step': 'search',
                 'query': state["search_query"],
@@ -176,6 +217,12 @@ Do not include any other text or explanation. Only return the dictionary.
                 'error': str(e)
             })
             state["current_step"] = 'final_answer'
+            
+            # Log the error
+            state["logs"]["search_error"] = {
+                "error": str(e),
+                "timestamp": datetime.now().isoformat()
+            }
         return state
 
     def _generate_final_answer(self, state: AgentState) -> AgentState:
@@ -183,13 +230,12 @@ Do not include any other text or explanation. Only return the dictionary.
         history_str = "\n".join([f"{h['step']}: {h.get('output', h.get('results', h.get('error', '')))}" 
                                for h in state["history"]])
         
-        prompt = f"""Based on the following information and history, provide a final answer to the question: {state['question']}
-        
-        History of steps taken:
-        {history_str}
-        
-        Provide a clear, concise answer that addresses the original question.
-        """
+        prompt = f"""Question: {state['question']}
+
+History of steps taken:
+{history_str}
+
+Return ONLY the direct answer to the question. Do not include any explanations, introductions, or formatting. Just the answer."""
         
         print("\n=== Generate Final Answer ===")
         print(f"Question: {state['question']}")
@@ -200,27 +246,45 @@ Do not include any other text or explanation. Only return the dictionary.
         print("\nFinal Answer:")
         print(llm_response)
         
+        # Log the final answer generation
+        state["logs"]["final_answer"] = {
+            "prompt": prompt,
+            "response": llm_response,
+            "history": history_str,
+            "timestamp": datetime.now().isoformat()
+        }
+        
         state["final_answer"] = llm_response
         return state
 
-    def __call__(self, question: str) -> str:
+    def __call__(self, question: str, task_id: str = "unknown") -> str:
         """Process a question through the agent workflow."""
         print(f"Agent received question: {question[:50]}...")
         
         try:
             # Initialize the state
             initial_state: AgentState = {
-                "question": question,  # Now a string, not a list
+                "question": question,
                 "current_step": "analyze",
                 "tool_output": "",
                 "final_answer": "",
                 "history": [],
                 "needs_more_info": False,
-                "search_query": ""
+                "search_query": "",
+                "task_id": task_id,
+                "logs": {}
             }
             
             # Run the workflow
             final_state = self.workflow.invoke(initial_state)
+            
+            # Save logs to file
+            log_to_file(
+                task_id=final_state["task_id"],
+                question=final_state["question"],
+                log_data=final_state["logs"]
+            )
+            
             return final_state["final_answer"]
             
         except Exception as e:
@@ -295,13 +359,15 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         try:
             # Initialize the state for this question
             initial_state = {
-                "question": question_text,  # Now a string, not a list
+                "question": question_text,
                 "current_step": "analyze",
                 "tool_output": "",
                 "final_answer": "",
                 "history": [],
                 "needs_more_info": False,
-                "search_query": ""
+                "search_query": "",
+                "task_id": task_id,
+                "logs": {}
             }
             
             # Run the workflow for this question