Update app.py

app.py

@@ -15,6 +15,7 @@ from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_community.document_loaders import ArxivLoader
 
+
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -22,6 +23,7 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #Load environment variables
 load_dotenv()
 
+
 from langgraph.graph import END, StateGraph
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
@@ -142,99 +144,208 @@ class VideoTranscriptionTool(BaseTool):
     def _arun(self, *args, **kwargs):
         raise NotImplementedError("Async not supported for this tool.")
 
-        
 
-def indent_code(code: str) -> str:
-    return '\n    '.join(code.splitlines())
 
-# ====== Agent State ======
+
+
+import os
+import time
+import json
+from typing import TypedDict, List, Union, Any, Dict
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.schema import HumanMessage, AIMessage, SystemMessage
+from langchain.prompts import ChatPromptTemplate
+from langgraph.graph import StateGraph, END
+from google.api_core.exceptions import ResourceExhausted
+
+# Assume these tools are defined elsewhere and imported
+# Placeholder for your actual tool implementations
+# For example:
+# from your_tools_module import duckduckgo_search, wikipedia_search, arxiv_search, document_qa, python_execution
+# And ensure you have a proper VideoTranscriptionTool
+def duckduckgo_search(query: str) -> str:
+    """Performs a DuckDuckGo search for current events or general facts."""
+    # Placeholder for actual implementation
+    print(f"DEBUG: duckduckgo_search called with: {query}")
+    return f"Search result for '{query}': Example relevant information from web."
+
+def wikipedia_search(query: str) -> str:
+    """Searches Wikipedia for encyclopedic information."""
+    # Placeholder for actual implementation
+    print(f"DEBUG: wikipedia_search called with: {query}")
+    return f"Wikipedia result for '{query}': Found detailed article."
+
+def arxiv_search(query: str) -> str:
+    """Searches ArXiv for scientific preprints and papers."""
+    # Placeholder for actual implementation
+    print(f"DEBUG: arxiv_search called with: {query}")
+    return f"ArXiv result for '{query}': Found relevant research paper."
+
+def document_qa(document_path: str, question: str) -> str:
+    """Answers questions based on the content of a given document file (PDF, DOCX, TXT)."""
+    # Placeholder for actual implementation
+    print(f"DEBUG: document_qa called with: {document_path}, question: {question}")
+    return f"Document QA result for '{question}': Answer extracted from document."
+
+def python_execution(code: str) -> str:
+    """Executes Python code in a sandboxed environment for calculations or data manipulation."""
+    # Placeholder for actual implementation - IMPORTANT: Implement this securely!
+    # Example (UNSAFE for real use without proper sandboxing):
+    try:
+        exec_globals = {}
+        exec_locals = {}
+        exec(code, exec_globals, exec_locals)
+        return str(exec_locals.get('result', 'Code executed, no explicit result assigned to "result" variable.'))
+    except Exception as e:
+        return f"Python execution error: {str(e)}"
+
+class VideoTranscriptionTool:
+    """Transcribes and analyzes video content from a URL or ID."""
+    def __call__(self, video_id_or_url: str) -> str:
+        # Placeholder for actual implementation using youtube-transcript-api etc.
+        print(f"DEBUG: VideoTranscriptionTool called with: {video_id_or_url}")
+        return f"Video transcription/analysis result for '{video_id_or_url}': Summary of video content."
+
+
+# --- Agent State Definition ---
 class AgentState(TypedDict):
     question: str
-    history: Annotated[List[Dict], operator.add]
-    context: str
+    history: List[Union[HumanMessage, AIMessage, Dict[str, Any]]] # Allows for tool calls as dicts
+    context: Dict[str, Any]
     reasoning: str
     iterations: int
+    final_answer: Union[str, float, int, None]
+    current_task: str # Added for more focused reasoning
+    current_thoughts: str # Added for more focused reasoning
 
-# ====== Graph Components ======
-def init_state(question: str) -> AgentState:
-    return {
-        "question": question,
-        "history": [],
-        "context": f"User question: {question}",
-        "reasoning": "",
-        "iterations": 0
-    }
+# --- Utility Functions ---
+def parse_agent_response(response_content: str) -> tuple[str, str, str]:
+    """
+    Parses the LLM's JSON output for reasoning, action, and action input.
+    """
+    try:
+        response_json = json.loads(response_content)
+        reasoning = response_json.get("Reasoning", "").strip()
+        action = response_json.get("Action", "").strip()
+        action_input = response_json.get("Action Input", "").strip()
+        return reasoning, action, action_input
+    except json.JSONDecodeError:
+        # Fallback for when LLM doesn't return perfect JSON (less likely with good prompt)
+        print(f"WARNING: LLM response not perfectly JSON: {response_content[:200]}...")
+        # Attempt heuristic parsing as a last resort
+        reasoning_match = response_content.split("Reasoning:", 1)
+        reasoning = reasoning_match[1].split("Action:", 1)[0].strip() if len(reasoning_match) > 1 else ""
 
+        action_part_match = response_content.split("Action:", 1)
+        action_part = action_part_match[1].strip() if len(action_part_match) > 1 else ""
 
+        action_input_match = action_part.split("Action Input:", 1)
+        action = action_input_match[0].strip()
+        action_input = action_input_match[1].strip() if len(action_input_match) > 1 else ""
+        return reasoning, action, action_input
 
-def should_continue(state: AgentState) -> str:
-    history = state.get("history", [])
 
-    if not history:
-        return "reason"  # No history yet, reason first
+# --- Graph Nodes ---
 
-    last_message = history[-1]
+def should_continue(state: AgentState) -> str:
+    """
+    Determines if the agent should continue reasoning, use a tool, or end.
+    """
+    history = state.get("history", [])
 
-    # End if agent has produced a final answer
-    if isinstance(last_message, AIMessage) and "FINAL ANSWER:" in last_message.content:
+    # Check for final answer in the last AIMessage
+    if history and isinstance(history[-1], AIMessage) and "FINAL ANSWER:" in history[-1].content:
+        print("DEBUG: should_continue -> END (Final Answer detected)")
         return "end"
+    
+    # Check if a tool was just executed (its output is in history)
+    # and the next step should be reasoning over that output
+    for msg in reversed(history):
+        if isinstance(msg, AIMessage) and any(f"[{tool.name} output]" in msg.content for tool in state.get("tools", [])):
+            print("DEBUG: should_continue -> REASON (Tool output detected, need to process)")
+            return "reason"
 
-    # If an action_request exists, trigger tool use
+    # Check if there's an action request to be executed
+    # This happens *after* reasoning has determined a tool is needed,
+    # but *before* the tool has run.
     for msg in reversed(history):
-        if isinstance(msg, dict) and msg.get("role") == "action_request":
-            return "continue"
+        if isinstance(msg, dict) and msg.get("type") == "action_request":
+            print("DEBUG: should_continue -> ACTION (Action request pending)")
+            return "action"
 
-    # Otherwise, go back to reasoning
+    # If nothing else, assume we need to reason
+    print("DEBUG: should_continue -> REASON (Default to reasoning)")
     return "reason"
 
 
-
 def reasoning_node(state: AgentState) -> AgentState:
-    import os
-    import time
-    from langchain_google_genai import ChatGoogleGenerativeAI
-    from langchain.schema import HumanMessage, AIMessage
-    from langchain.prompts import ChatPromptTemplate
-    from google.api_core.exceptions import ResourceExhausted
+    """
+    Node for the agent to analyze the question, determine next steps,
+    and select tools.
+    """
+    print(f"DEBUG: Entering reasoning_node. Iteration: {state['iterations']}")
+    print(f"DEBUG: Current history length: {len(state.get('history', []))}")
 
     # Load API key
     GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
     if not GOOGLE_API_KEY:
         raise ValueError("GOOGLE_API_KEY not set in environment variables.")
 
-    # Ensure history is well-formed
+    # Ensure history is well-formed for the LLM prompt
     if "history" not in state or not isinstance(state["history"], list):
         state["history"] = []
-    if not state["history"] or not isinstance(state["history"][-1], HumanMessage):
-        state["history"].append(HumanMessage(content="Continue."))
-
-    # Ensure context and reasoning fields
+    
+    # Initialize/update state fields
     state.setdefault("context", {})
     state.setdefault("reasoning", "")
     state.setdefault("iterations", 0)
+    state.setdefault("current_task", "Understand the question and plan the next step.")
+    state.setdefault("current_thoughts", "")
 
     # Create Gemini model wrapper
     llm = ChatGoogleGenerativeAI(
-        model="gemini-1.5-flash",
-        temperature=0.1,
+        model="gemini-1.5-flash", # Use a fast model for agentic loops
+        temperature=0.1, # Keep it low for more deterministic reasoning
         google_api_key=GOOGLE_API_KEY
     )
 
-    # Create prompt chain
+    # Dynamically generate tool descriptions for the prompt
+    tool_descriptions = "\n".join([
+        f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
+    ])
+
+    # Craft a more robust and explicit system prompt
+    system_prompt = (
+        "You are an expert problem solver, designed to provide concise and accurate answers. "
+        "Your process involves analyzing the question, intelligently selecting and using tools, "
+        "and synthesizing information.\n\n"
+        "**Available Tools:**\n"
+        f"{tool_descriptions}\n\n"
+        "**Tool Usage Guidelines:**\n"
+        "- Use **duckduckgo_search** for current events, general facts, or quick lookups.\n"
+        "- Use **wikipedia_search** for encyclopedic information, historical context, or detailed topics.\n"
+        "- Use **arxiv_search** for scientific papers, research, or cutting-edge technical information.\n"
+        "- Use **document_qa** when the question explicitly refers to a specific document file (e.g., 'Analyze this PDF').\n"
+        "- Use **python_execution** for complex calculations, data manipulation, or logical operations that cannot be done with simple reasoning. Always provide the full Python code.\n"
+        "- Use **VideoTranscriptionTool** for any question involving video or audio content.\n\n"
+        "**Current Context:**\n{context}\n\n"
+        "**Previous Reasoning Steps:**\n{reasoning}\n\n"
+        "**Current Task:** {current_task}\n"
+        "**Current Thoughts:** {current_thoughts}\n\n"
+        "**Your Response MUST be a valid JSON object with the following keys:**\n"
+        "```json\n"
+        "{\n"
+        "  \"Reasoning\": \"Your detailed analysis of the question and why you chose a specific action.\",\n"
+        "  \"Action\": \"[Tool name OR 'Final Answer']\",\n"
+        "  \"Action Input\": \"[Input for the selected tool OR the final response]\"\n"
+        "}\n"
+        "```\n"
+        "Ensure 'Action Input' is appropriate for the chosen 'Action'. If 'Action' is 'Final Answer', provide the complete, concise answer."
+    )
+
     prompt = ChatPromptTemplate.from_messages([
-        ("system", (
-            "You're an expert problem solver. Analyze the question, select the best tool, "
-            "and provide reasoning. Available tools: duckduckgo_search, wikipedia_search, "
-            "arxiv_search, document_qa, python_execution.\n\n"
-            "Important: You must select a tool for questions involving video, audio, or code.\n\n"
-            "Current Context:\n{context}\n\n"
-            "Reasoning Steps:\n{reasoning}\n\n"
-            "Response Format:\n"
-            "Reasoning: [Your analysis]\n"
-            "Action: [Tool name OR 'Final Answer']\n"
-            "Action Input: [Input for tool OR final response]"
-        )),
-        *state["history"]
+        SystemMessage(content=system_prompt),
+        *state["history"] # Include full history for conversational context
     ])
 
     chain = prompt | llm
@@ -243,103 +354,118 @@ def reasoning_node(state: AgentState) -> AgentState:
     def call_with_retry(inputs, retries=3, delay=60):
         for attempt in range(retries):
             try:
-                return chain.invoke(inputs)
+                response = chain.invoke(inputs)
+                # Attempt to parse immediately to catch bad JSON before returning
+                parse_agent_response(response.content) 
+                return response
             except ResourceExhausted as e:
-                print(f"[Retry {attempt+1}] Gemini rate limit hit. Waiting {delay}s...")
+                print(f"[Retry {attempt+1}/{retries}] Gemini rate limit hit. Waiting {delay}s...")
                 time.sleep(delay)
-        raise RuntimeError("Failed after multiple retries due to Gemini quota limit.")
+            except json.JSONDecodeError as e:
+                print(f"[Retry {attempt+1}/{retries}] LLM returned invalid JSON. Retrying...")
+                print(f"Invalid JSON content: {response.content[:200]}...")
+                time.sleep(5) # Shorter delay for parsing errors
+            except Exception as e:
+                print(f"[Retry {attempt+1}/{retries}] An unexpected error occurred during LLM call: {e}. Retrying...")
+                time.sleep(delay)
+        raise RuntimeError("Failed after multiple retries due to Gemini quota limit or invalid JSON.")
 
     # Call model with retry protection
     response = call_with_retry({
         "context": state["context"],
         "reasoning": state["reasoning"],
-        "question": state["question"]
+        "question": state["question"], # Redundant as it's in history, but keeps prompt consistent
+        "current_task": state["current_task"],
+        "current_thoughts": state["current_thoughts"]
     })
 
-    # Parse output
+    # Parse output using the robust JSON parser
     content = response.content
     reasoning, action, action_input = parse_agent_response(content)
+    
+    print(f"DEBUG: LLM Response Content: {content[:200]}...")
+    print(f"DEBUG: Parsed Action: {action}, Action Input: {action_input[:100]}...")
 
     # Update state
-    state["history"].append(AIMessage(content=content))
+    state["history"].append(AIMessage(content=content)) # Store the raw LLM response
     state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
+    state["current_thoughts"] = reasoning # Update current thoughts for next iteration
 
     if "final answer" in action.lower():
         state["history"].append(AIMessage(content=f"FINAL ANSWER: {action_input}"))
+        state["final_answer"] = action_input # Set final answer directly in state
     else:
-        state["context"]["current_tool"] = {
+        # Store the action request in history for tool_node
+        state["history"].append({
+            "type": "action_request",
             "tool": action,
             "input": action_input
-        }
+        })
 
+    print(f"DEBUG: Exiting reasoning_node. New history length: {len(state['history'])}")
     return state
 
 
-
-
-    
-
 def tool_node(state: AgentState) -> AgentState:
-    from langchain.schema import AIMessage
+    """
+    Node for executing the chosen tool and returning its output.
+    """
+    print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
 
     # Ensure history exists
     if "history" not in state or not isinstance(state["history"], list):
         raise ValueError("Invalid or missing history in state")
 
     # Find the most recent action request in history
-    tool_call = None
+    tool_call_dict = None
     for msg in reversed(state["history"]):
-        if isinstance(msg, dict) and msg.get("role") == "action_request":
-            tool_call = msg
+        if isinstance(msg, dict) and msg.get("type") == "action_request":
+            tool_call_dict = msg
             break
 
-    if not tool_call:
-        raise ValueError("No tool call found in history")
+    if not tool_call_dict:
+        # This shouldn't happen if should_continue logic is correct, but for robustness
+        print("WARNING: No action_request found in history, skipping tool execution.")
+        return state
 
-    tool_name = tool_call.get("tool")
-    tool_input = tool_call.get("input")
+    tool_name = tool_call_dict.get("tool")
+    tool_input = tool_call_dict.get("input")
 
     # Defensive check for missing tool or input
     if not tool_name or tool_input is None:
-        raise ValueError("Tool name or input missing from action request")
+        raise ValueError(f"Tool name '{tool_name}' or input '{tool_input}' missing from action request: {tool_call_dict}")
 
-    # Look up and invoke the tool
-    agent = BasicAgent()  # Create agent to access tools
-    tool_fn = next((t for t in agent.tools if t.__name__ == tool_name), None)
+    # Look up and invoke the tool from the state's tool list
+    available_tools = state.get("tools", [])
+    tool_fn = next((t for t in available_tools if t.name == tool_name), None) # Assuming tools are LangChain Tool objects now
 
     if tool_fn is None:
-        raise ValueError(f"Tool '{tool_name}' not found")
-
-    try:
-        tool_output = tool_fn(tool_input)
-    except Exception as e:
-        tool_output = f"[Tool Error] {str(e)}"
+        # Fallback for unrecognized tool - feedback to LLM
+        tool_output = f"[Tool Error] Tool '{tool_name}' not found or not available. Please choose from: {', '.join([t.name for t in available_tools])}"
+        print(f"ERROR: {tool_output}")
+    else:
+        try:
+            print(f"DEBUG: Invoking tool '{tool_name}' with input: '{tool_input[:100]}...'")
+            tool_output = tool_fn.run(tool_input) # Assuming tool.run() method for LangChain Tools
+            if not tool_output: # Handle empty tool output
+                tool_output = f"[{tool_name} output] No specific result found for '{tool_input}'. The tool might have returned an empty response."
+        except Exception as e:
+            tool_output = f"[Tool Error] An error occurred while running '{tool_name}': {str(e)}"
+            print(f"ERROR: {tool_output}")
 
     # Add output to history as an AIMessage
+    # Ensure the history only contains HumanMessage and AIMessage objects for LangGraph's internal processing.
+    # The action_request dict can be removed or transformed if it's no longer needed for internal state.
+    # For now, we'll just add the tool output.
     state["history"].append(AIMessage(content=f"[{tool_name} output]\n{tool_output}"))
-
+    
+    print(f"DEBUG: Exiting tool_node. Tool output added to history. New history length: {len(state['history'])}")
     return state
 
 
-
-def parse_agent_response(response: str) -> tuple:
-    """Extract reasoning, action, and input from response"""
-    reasoning = response.split("Reasoning:")[1].split("Action:")[0].strip()
-    action_part = response.split("Action:")[1].strip()
-    
-    if "Action Input:" in action_part:
-        action, action_input = action_part.split("Action Input:", 1)
-        action = action.strip()
-        action_input = action_input.strip()
-    else:
-        action = action_part
-        action_input = ""
-    
-    return reasoning, action, action_input
-
 # ====== Agent Graph ======
-def create_agent_workflow():
+def create_agent_workflow(tools: List[Any]): # tools are passed in now
     workflow = StateGraph(AgentState)
     
     # Define nodes
@@ -354,58 +480,72 @@ def create_agent_workflow():
         "reason",
         should_continue,
         {
-            "continue": "action",
-            "reason": "reason",
-            "end": END
+            "action": "action", # Go to action node if a tool is requested
+            "reason": "reason", # Loop back to reason if more thinking is needed
+            "end": END          # End if final answer detected
         }
     )
     
-    workflow.add_edge("action", "reason")
-    return workflow.compile()
+    workflow.add_edge("action", "reason") # Always go back to reasoning after a tool action
+    
+    # Compile the graph
+    app = workflow.compile()
+    
+    # Pass tools into the state so nodes can access them.
+    # This is a bit of a hacky way to get them into the state, but works for now.
+    # A cleaner way might be to make `tool_node` receive tools as a closure or directly from agent init.
+    # For this example, we'll modify the initial state for each invocation.
+    return app
+
 
 # ====== Agent Interface ======
 class BasicAgent:
     def __init__(self):
-        self.workflow = create_agent_workflow()
+        # Tools need to be LangChain Tool objects for name and description
+        from langchain.tools import Tool
         self.tools = [
-            duckduckgo_search,
-            wikipedia_search,
-            arxiv_search,
-            document_qa,
-            python_execution,
-            VideoTranscriptionTool() 
+            Tool(name="duckduckgo_search", func=duckduckgo_search, description="Performs a DuckDuckGo search for current events or general facts."),
+            Tool(name="wikipedia_search", func=wikipedia_search, description="Searches Wikipedia for encyclopedic information."),
+            Tool(name="arxiv_search", func=arxiv_search, description="Searches ArXiv for scientific preprints and papers."),
+            Tool(name="document_qa", func=document_qa, description="Answers questions based on the content of a given document file (PDF, DOCX, TXT). Requires 'attachment_path' and 'question' as input."),
+            Tool(name="python_execution", func=python_execution, description="Executes Python code in a sandboxed environment for complex calculations or data manipulation."),
+            Tool(name="VideoTranscriptionTool", func=VideoTranscriptionTool(), description="Transcribes and analyzes video content from a URL or ID. Use for any question involving video or audio.")
         ]
+        self.workflow = create_agent_workflow(self.tools) # Pass tools to workflow creator
 
     def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question[:50]}{'...' if len(question) > 50 else ''}")
+        print(f"\n--- Agent received question: {question[:50]}{'...' if len(question) > 50 else ''} ---")
 
-        # Initialize state with proper structure
+        # Initialize state with proper structure and pass tools
         state = {
             "question": question,
-            "context": {},  # Ensure it's a dict
+            "context": {},
             "reasoning": "",
             "iterations": 0,
-            "history": [HumanMessage(content=question)]
+            "history": [HumanMessage(content=question)],
+            "final_answer": None,
+            "current_task": "Understand the question and plan the next step.",
+            "current_thoughts": "",
+            "tools": self.tools # Pass tools into the state
         }
 
+        # Invoke the workflow
         final_state = self.workflow.invoke(state)
 
-        print(f"Final state keys: {list(final_state.keys())}")
-        if 'history' in final_state:
-            print(f"History length: {len(final_state['history'])}")
-            for i, msg in enumerate(final_state['history']):
-                if isinstance(msg, dict):
-                    print(f"Message {i}: dict - {msg}")
-                else:
-                    print(f"Message {i}: {type(msg).__name__} - {msg.content[:100]}...")
-
         # Extract the FINAL ANSWER from history
+        if final_state.get("final_answer"):
+            answer = final_state["final_answer"]
+            print(f"--- Agent returning FINAL ANSWER: {answer} ---")
+            return answer
+        
+        # Fallback if final_answer wasn't set correctly in state
         for msg in reversed(final_state["history"]):
             if isinstance(msg, AIMessage) and "FINAL ANSWER:" in msg.content:
                 answer = msg.content.split("FINAL ANSWER:")[1].strip()
-                print(f"Agent returning answer: {answer}")
+                print(f"--- Agent returning FINAL ANSWER (from history): {answer} ---")
                 return answer
 
+        print(f"--- ERROR: No FINAL ANSWER found in agent history for question: {question} ---")
         raise ValueError("No FINAL ANSWER found in agent history.")