Update app.py

app.py

@@ -1,335 +1,195 @@
 import os
 import gradio as gr
 import requests
-import inspect
-import pandas as pd
 import ast
-import operator
-import time
 import json
+import time
+import pandas as pd
 from datetime import datetime
-from typing import List, Dict, Any, Optional, Annotated
+from typing import List, Dict, Any, Annotated
 from langgraph.graph import Graph, StateGraph
-from langgraph.prebuilt import ToolNode
-from tools import simple_search, jina_search_tool
-from openai import OpenAI
 from typing_extensions import TypedDict
+from openai import OpenAI
+
+# -------------------------
+# Utility helpers
+# -------------------------
 
-def override(_, new): return new
+def override(_, new):
+    return new
 
-def merge_dicts(dict1: Dict, dict2: Dict) -> Dict:
-    """Merge two dictionaries, with values from dict2 taking precedence."""
-    return {**dict1, **dict2}
+def merge_dicts(old: Dict, new: Dict) -> Dict:
+    """Merge two dictionaries, with *new* values taking precedence."""
+    return {**old, **new}
+
+# -------------------------
+# Environment & constants
+# -------------------------
 
-print("trial")
-# (Keep Constants as is)
-# --- Constants ---
 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
-JINA_API_KEY = os.getenv("JINA_API_KEY")  # Get Jina API key from environment
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+JINA_API_KEY = os.getenv("JINA_API_KEY")
 
-# Create logs directory if it doesn't exist
 LOGS_DIR = "question_logs"
 os.makedirs(LOGS_DIR, exist_ok=True)
 
+# -------------------------
+# Jina AI search tool (replaces DDG + Reader)
+# -------------------------
+
+def jina_search_tool(query: str, api_key: str, max_results: int = 5) -> List[str]:
+    """Return *max_results* clean markdown snippets for *query* using s.jina.ai."""
+    endpoint = f"https://s.jina.ai/{query.replace(' ', '+')}"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+        "User-Agent": "Mozilla/5.0",
+    }
+    resp = requests.get(endpoint, headers=headers, timeout=15)
+    if resp.status_code != 200:
+        raise RuntimeError(f"Jina search failed with status {resp.status_code}: {resp.text[:200]}")
+    data = resp.json()
+    return [item.get("content", "") for item in data.get("results", [])][:max_results]
+
+# -------------------------
+# Logging helper
+# -------------------------
+
 def log_to_file(task_id: str, question: str, log_data: Dict[str, Any]):
-    """Store logs for a question in a JSON file."""
-    try:
-        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
-        }
-        
-        print(f"\n=== Saving Logs ===")
-        print(f"Task ID: {task_id}")
-        print(f"Question: {question}")
-        print(f"Log Data: {json.dumps(log_data, indent=2)}")
-        
-        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}")
-    except Exception as e:
-        print(f"Error saving logs: {e}")
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = os.path.join(LOGS_DIR, f"question_{task_id}_{ts}.json")
+    with open(filename, "w", encoding="utf-8") as f:
+        json.dump({"task_id": task_id, "question": question, "timestamp": ts, "logs": log_data}, f, indent=2)
+
+# -------------------------
+# State definition
+# -------------------------
 
 class AgentState(TypedDict):
     question: Annotated[str, override]
     current_step: Annotated[str, override]
-    tool_output: Annotated[str, override]
     final_answer: Annotated[str, override]
-    history: Annotated[List[Dict[str, str]], operator.add]
-    needs_more_info: Annotated[bool, override]
+    history: Annotated[List[Dict[str, str]], list.__add__]
+    needs_search: Annotated[bool, override]
     search_query: Annotated[str, override]
     task_id: Annotated[str, override]
-    logs: Annotated[Dict[str, Any], merge_dicts]  # Use merge_dicts instead of operator.add
+    logs: Annotated[Dict[str, Any], merge_dicts]
+
+# -------------------------
+# BasicAgent implementation
+# -------------------------
 
 class BasicAgent:
     def __init__(self):
-        print("Initializing BasicAgent with OpenAI...")
         if not OPENAI_API_KEY:
-            raise ValueError("OPENAI_API_KEY environment variable not set. Please set your OpenAI API key.")
-        
-        # Initialize OpenAI client
+            raise EnvironmentError("OPENAI_API_KEY not set")
+        if not JINA_API_KEY:
+            raise EnvironmentError("JINA_API_KEY not set")
         self.llm = OpenAI(api_key=OPENAI_API_KEY)
-        
-        # Create the agent workflow
-        print("Creating workflow variable")
-        self.workflow = self._create_workflow()
-        print("BasicAgent initialization complete.")
-
-    def _call_llm_api(self, prompt: str) -> str:
-        """Call the model and return the raw text output."""
-        try:
-            print("=== Sending prompt ===")
-            print(prompt[:500])
-            response = self.llm.chat.completions.create(
-                model="gpt-4.1-nano",
-                messages=[
-                    {"role": "system", "content": "You are a helpful AI assistant that provides clear and concise answers."},
-                    {"role": "user", "content": prompt}
-                ],
-                max_tokens=200,
-                temperature=0.7,
-                top_p=0.95,
-                frequency_penalty=0.1
-            )
-            print("=== Received response ===")
-            response_text = response.choices[0].message.content
-            print(response_text)
-            return response_text
-        except Exception as e:
-            print(f"Error calling LLM API: {e}")
-            return f"Error getting response from LLM: {str(e)}"
-
-    def _create_workflow(self) -> Graph:
-        """Create the agent workflow using LangGraph."""
-        # Create the workflow with state schema
-        print("Creating Stategraph : error happens here?")
-        workflow = StateGraph(state_schema=AgentState)
-        print("Stategraph created")
-        # Add nodes
-        workflow.add_node("analyze", self._analyze_question)
-        workflow.add_node("search", self._use_search)
-        workflow.add_node("generate_answer", self._generate_final_answer)
-        
-        # Define edges
-        workflow.add_edge("analyze", "search")
-        workflow.add_edge("analyze", "generate_answer")
-        workflow.add_edge("search", "generate_answer")
-        
-        # Define conditional edges
-        def router(state: AgentState) -> str:
-            if state["current_step"] == 'search':
-                return 'search'
-            elif state["current_step"] == 'final_answer':
-                return 'generate_answer'
-            return 'analyze'
-        
-        workflow.add_conditional_edges(
-            "analyze",
-            router,
-            {
-                "search": "search",
-                "final_answer": "generate_answer"
-            }
+        self.workflow = self._build_workflow()
+
+    # ---- Low‑level LLM call
+    def _call_llm(self, prompt: str, max_tokens: int = 256) -> str:
+        resp = self.llm.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": "You are a careful reasoning assistant."},
+                {"role": "user", "content": prompt},
+            ],
+            temperature=0.3,
+            max_tokens=max_tokens,
         )
-        
-        # Set entry and exit points
-        workflow.set_entry_point("analyze")
-        workflow.set_finish_point("generate_answer")
-        
-        return workflow.compile()
+        return resp.choices[0].message.content.strip()
 
+    # ---- Workflow nodes
     def _analyze_question(self, state: AgentState) -> AgentState:
-        """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 Python dictionary in this exact format, with no other text or explanation:
-{{
-    "needs_search": true/false,
-    "search_query": "query if needed"
-}}"""
-        
+        prompt = (
+            "You will receive a user question. Think step‑by‑step to decide whether external web search is required. "
+            "Respond ONLY with a valid Python dict literal in the following format and NOTHING else:\n"
+            "{\n  'needs_search': bool,\n  'search_query': str\n} \n\n"
+            f"Question: {state['question']}"
+        )
+        raw = self._call_llm(prompt)
         try:
-            llm_response = self._call_llm_api(prompt)
-            print("\n=== Analyze Question LLM Response ===")
-            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', '')
-            
-            if analysis.get('needs_search', False):
-                state["current_step"] = 'search'
-            else:
-                state["current_step"] = 'final_answer'
-        except (ValueError, SyntaxError) as e:
-            print(f"Error parsing LLM response: {e}")
-            # Default to search if we can't parse the response
-            state["needs_more_info"] = True
+            decision = ast.literal_eval(raw)
+            state["needs_search"] = bool(decision.get("needs_search", False))
+            state["search_query"] = decision.get("search_query", state["question"])
+        except Exception:
+            # fallback: assume search needed
+            state["needs_search"] = True
             state["search_query"] = state["question"]
-            state["current_step"] = 'search'
-            
-            # Log the error
-            state["logs"] = {
-                "analyze_error": {
-                    "error": str(e),
-                    "timestamp": datetime.now().isoformat()
-                }
-            }
-        
+            decision = {"parse_error": raw}
+        state["logs"] = {
+            "analyze": {"prompt": prompt, "llm_response": raw, "decision": decision}
+        }
+        state["current_step"] = "search" if state["needs_search"] else "answer"
+        state["history"].append({"step": "analyze", "output": decision})
         return state
 
-    def _use_search(self, state: AgentState) -> AgentState:
-        """Use the search tool."""
-        time.sleep(2)  # Sleep before search
-        
+    def _perform_search(self, state: AgentState) -> AgentState:
         try:
-            print("\n=== Search Tool ===")
-            print(f"Search Query: {state['search_query']}")
-            
-            # Try Jina search first, fall back to simple search if it fails
-            search_results = []
-            if JINA_API_KEY:
-                try:
-                    search_results = jina_search_tool(
-                        query=state["search_query"],
-                        api_key=JINA_API_KEY
-                    )
-                    print("Using Jina search results")
-                except Exception as e:
-                    print(f"Jina search failed: {e}, falling back to simple search")
-                    search_results = simple_search(
-                        query=state["search_query"],
-                        max_results=3
-                    )
-            else:
-                print("No Jina API key found, using simple search")
-                search_results = simple_search(
-                    query=state["search_query"],
-                    max_results=3
-                )
-            
-            print("Search Results:")
-            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(),
-                "search_type": "jina" if JINA_API_KEY and search_results else "simple"
-            }
-            
-            state["history"].append({
-                'step': 'search',
-                'query': state["search_query"],
-                'results': search_results
-            })
-            state["needs_more_info"] = False
-            state["current_step"] = 'final_answer'
+            results = jina_search_tool(state["search_query"], JINA_API_KEY)
         except Exception as e:
-            print(f"Search Error: {e}")
-            state["history"].append({
-                'step': 'search_error',
-                'error': str(e)
-            })
-            state["current_step"] = 'final_answer'
-            
-            # Log the error
-            state["logs"]["search_error"] = {
-                "error": str(e),
-                "timestamp": datetime.now().isoformat()
-            }
+            results = [f"SEARCH_ERROR: {e}"]
+        state["history"].append({"step": "search", "results": results})
+        state["logs"]["search"] = {"query": state["search_query"], "results": results}
+        state["current_step"] = "answer"
         return state
 
-    def _generate_final_answer(self, state: AgentState) -> AgentState:
-        """Generate the final answer based on all gathered information."""
-        history_str = "\n".join([f"{h['step']}: {h.get('output', h.get('results', h.get('error', '')))}" 
-                               for h in state["history"]])
-        
-        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']}")
-        print("History:")
-        print(history_str)
-        
-        llm_response = self._call_llm_api(prompt)
-        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
+    def _generate_answer(self, state: AgentState) -> AgentState:
+        history_text = "\n".join(str(item) for item in state["history"])
+        prompt = (
+            f"Answer the user question as directly as possible. If sources were retrieved, incorporate them.\n"
+            f"Question: {state['question']}\n\nContext:\n{history_text}\n\n"
+            "Give ONLY the final answer without extra formatting or explanation."
+        )
+        answer = self._call_llm(prompt, max_tokens=150)
+        state["final_answer"] = answer
+        state["history"].append({"step": "answer", "output": answer})
+        state["logs"]["final_answer"] = {"prompt": prompt, "response": answer}
+        state["current_step"] = "done"
         return state
 
+    # ---- Build LangGraph workflow
+    def _build_workflow(self) -> Graph:
+        sg = StateGraph(state_schema=AgentState)
+        sg.add_node("analyze", self._analyze_question)
+        sg.add_node("search", self._perform_search)
+        sg.add_node("answer", self._generate_answer)
+
+        # transitions
+        sg.add_edge("analyze", "search")
+        sg.add_edge("analyze", "answer")
+        sg.add_edge("search", "answer")
+
+        def router(state: AgentState):
+            return state["current_step"]
+
+        sg.add_conditional_edges("analyze", router, {"search": "search", "answer": "answer"})
+        sg.set_entry_point("analyze")
+        sg.set_finish_point("answer")
+        return sg.compile()
+
+    # ---- Public call
     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,
-                "current_step": "analyze",
-                "tool_output": "",
-                "final_answer": "",
-                "history": [],
-                "needs_more_info": False,
-                "search_query": "",
-                "task_id": task_id,
-                "logs": {}
-            }
-            
-            # Run the workflow
-            final_state = self.workflow.invoke(initial_state)
-            
-            # Ensure logs directory exists
-            os.makedirs(LOGS_DIR, exist_ok=True)
-            
-            # Save logs to file
-            if final_state["logs"]:  # Only save if we have logs
-                log_to_file(
-                    task_id=final_state["task_id"],
-                    question=final_state["question"],
-                    log_data=final_state["logs"]
-                )
-            else:
-                print("No logs to save in final state")
-            
-            return final_state["final_answer"]
-            
-        except Exception as e:
-            print(f"Error in agent processing: {e}")
-            return f"I encountered an error while processing your question: {str(e)}"
+        state: AgentState = {
+            "question": question,
+            "current_step": "analyze",
+            "final_answer": "",
+            "history": [],
+            "needs_search": False,
+            "search_query": "",
+            "task_id": task_id,
+            "logs": {},
+        }
+        final_state = self.workflow.invoke(state)
+        if final_state["logs"]:
+            log_to_file(task_id, question, final_state["logs"])
+        return final_state["final_answer"]
+
+# ----------------------------------------------------------------------------------
+# Gradio Interface & Submission Routines
+# ----------------------------------------------------------------------------------
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
@@ -337,7 +197,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     and displays the results.
     """
     # --- Determine HF Space Runtime URL and Repo URL ---
-    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code
+    space_id = os.getenv("SPACE_ID")
     print("Space ID: ", space_id)
     if profile:
         username = f"{profile.username}"
@@ -352,9 +212,9 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent
     try:
-        print("Initializing agent: trial ")
+        print("Initializing agent...")
         agent = BasicAgent()
-        print("Agent initialized successfully with workflow.")
+        print("Agent initialized successfully.")
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -397,50 +257,25 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             continue
 
         try:
-            # Initialize the state for this question
-            initial_state = {
-                "question": question_text,
-                "current_step": "analyze",
-                "tool_output": "",
-                "final_answer": "",
-                "history": [],
-                "needs_more_info": False,
-                "search_query": "",
-                "task_id": task_id,
-                "logs": {}
-            }
-            
-            # Run the workflow for this question
             print(f"\nProcessing question {task_id}: {question_text[:50]}...")
-            final_state = agent.workflow.invoke(initial_state)
-            
-            # Log the workflow history
-            workflow_history = "\n".join([
-                f"Step: {h['step']}\n" +
-                f"Input: {h.get('input', h.get('query', ''))}\n" +
-                f"Output: {h.get('output', h.get('results', h.get('error', '')))}"
-                for h in final_state["history"]
-            ])
+            answer = agent(question_text, task_id)
             
             # Add to results
-            submitted_answer = final_state["final_answer"]
-            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
+            answers_payload.append({"task_id": task_id, "submitted_answer": answer})
             results_log.append({
                 "Task ID": task_id,
                 "Question": question_text,
-                "Submitted Answer": submitted_answer,
-                "Workflow History": workflow_history
+                "Submitted Answer": answer
             })
             
-            print(f"Completed question {task_id} with {len(final_state['history'])} workflow steps")
+            print(f"Completed question {task_id}")
             
         except Exception as e:
-            print(f"Error running agent workflow on task {task_id}: {e}")
+            print(f"Error running agent on task {task_id}: {e}")
             results_log.append({
                 "Task ID": task_id,
                 "Question": question_text,
-                "Submitted Answer": f"WORKFLOW ERROR: {e}",
-                "Workflow History": "Error occurred before workflow completion"
+                "Submitted Answer": f"ERROR: {e}"
             })
 
     if not answers_payload:
@@ -523,7 +358,6 @@ with gr.Blocks() as demo:
     run_button = gr.Button("Run Evaluation & Submit All Answers")
 
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
-    # Removed max_rows=10 from DataFrame constructor
     results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
 
     run_button.click(
@@ -535,7 +369,7 @@ if __name__ == "__main__":
     print("\n" + "-"*30 + " App Starting " + "-"*30)
     # Check for SPACE_HOST and SPACE_ID at startup for information
     space_host_startup = os.getenv("SPACE_HOST")
-    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup
+    space_id_startup = os.getenv("SPACE_ID")
 
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
@@ -543,7 +377,7 @@ if __name__ == "__main__":
     else:
         print("ℹ️  SPACE_HOST environment variable not found (running locally?).")
 
-    if space_id_startup: # Print repo URLs if SPACE_ID is found
+    if space_id_startup:
         print(f"✅ SPACE_ID found: {space_id_startup}")
         print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
         print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
@@ -553,4 +387,4 @@ if __name__ == "__main__":
     print("-"*(60 + len(" App Starting ")) + "\n")
 
     print("Launching Gradio Interface for Basic Agent Evaluation...")
-    demo.launch(debug=True, share=False)
+    demo.launch(debug=True, share=False)