Update app.py

app.py

@@ -157,43 +157,37 @@ 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
+from langchain.tools import Tool # Import Tool for consistent tool definitions
 
 # 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 = {}
+        # WARNING: This is a highly insecure way to execute arbitrary Python code.
+        # For production, use a secure, sandboxed environment (e.g., Docker container, dedicated service).
         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:
@@ -202,7 +196,6 @@ def python_execution(code: str) -> str:
 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."
 
@@ -210,13 +203,15 @@ class VideoTranscriptionTool:
 # --- Agent State Definition ---
 class AgentState(TypedDict):
     question: str
-    history: List[Union[HumanMessage, AIMessage, Dict[str, Any]]] # Allows for tool calls as dicts
-    context: Dict[str, Any]
+    history: List[Union[HumanMessage, AIMessage]] # History only contains proper messages
+    context: Dict[str, Any] # Use context for internal agent state
     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
+    current_task: str
+    current_thoughts: str
+    tools: List[Tool] # Pass tools into state
+
 
 # --- Utility Functions ---
 def parse_agent_response(response_content: str) -> tuple[str, str, str]:
@@ -230,9 +225,8 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
         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
+        # Fallback heuristic parsing (less reliable but better than nothing)
         reasoning_match = response_content.split("Reasoning:", 1)
         reasoning = reasoning_match[1].split("Action:", 1)[0].strip() if len(reasoning_match) > 1 else ""
 
@@ -251,29 +245,19 @@ def should_continue(state: AgentState) -> str:
     """
     Determines if the agent should continue reasoning, use a tool, or end.
     """
-    history = state.get("history", [])
-
-    # 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"
+    print(f"DEBUG: Entering should_continue. Current context: {state.get('context', {})}")
     
-    # 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"
-
-    # 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("type") == "action_request":
-            print("DEBUG: should_continue -> ACTION (Action request pending)")
-            return "action"
-
-    # If nothing else, assume we need to reason
+    # End if agent has produced a final answer
+    if state.get("final_answer"):
+        print("DEBUG: should_continue -> END (Final Answer set in state)")
+        return "end"
+
+    # Check if a tool action is pending in context
+    if state.get("context", {}).get("pending_action"):
+        print("DEBUG: should_continue -> ACTION (Pending action in context)")
+        return "action"
+
+    # Otherwise, go back to reasoning (e.g., after initial question, or after tool output)
     print("DEBUG: should_continue -> REASON (Default to reasoning)")
     return "reason"
 
@@ -291,21 +275,20 @@ def reasoning_node(state: AgentState) -> AgentState:
     if not GOOGLE_API_KEY:
         raise ValueError("GOOGLE_API_KEY not set in environment variables.")
 
-    # Ensure history is well-formed for the LLM prompt
-    if "history" not in state or not isinstance(state["history"], list):
-        state["history"] = []
-    
     # 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", "")
+    
+    # Clear any old pending action from context before generating a new one
+    state["context"].pop("pending_action", None) 
 
     # Create Gemini model wrapper
     llm = ChatGoogleGenerativeAI(
-        model="gemini-1.5-flash", # Use a fast model for agentic loops
-        temperature=0.1, # Keep it low for more deterministic reasoning
+        model="gemini-1.5-flash",
+        temperature=0.1,
         google_api_key=GOOGLE_API_KEY
     )
 
@@ -314,7 +297,6 @@ def reasoning_node(state: AgentState) -> AgentState:
         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, "
@@ -326,8 +308,8 @@ def reasoning_node(state: AgentState) -> AgentState:
         "- 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"
+        "- Use **python_execution** for complex calculations, data manipulation, or logical operations that cannot be done with simple reasoning. Always provide the full Python code, ensuring it's valid and executable, and assign the final result to a variable named 'result' if applicable.\n"
+        "- Use **VideoTranscriptionTool** for any question involving video or audio content. Provide the full YouTube URL or video ID.\n\n"
         "**Current Context:**\n{context}\n\n"
         "**Previous Reasoning Steps:**\n{reasoning}\n\n"
         "**Current Task:** {current_task}\n"
@@ -340,6 +322,8 @@ def reasoning_node(state: AgentState) -> AgentState:
         "  \"Action Input\": \"[Input for the selected tool OR the final response]\"\n"
         "}\n"
         "```\n"
+        "**CRITICAL RULE: 'Action' and 'Action Input' MUST NOT be empty, unless 'Action' is 'Final Answer' and 'Action Input' is the conclusive response.**\n"
+        "If you cannot determine a suitable tool or a final answer, return Action: 'Final Answer' with a message like 'I cannot answer this question with the available tools.' or 'More information is needed.'\n"
         "Ensure 'Action Input' is appropriate for the chosen 'Action'. If 'Action' is 'Final Answer', provide the complete, concise answer."
     )
 
@@ -350,13 +334,12 @@ def reasoning_node(state: AgentState) -> AgentState:
 
     chain = prompt | llm
 
-    # === Add Retry Logic ===
     def call_with_retry(inputs, retries=3, delay=60):
         for attempt in range(retries):
             try:
                 response = chain.invoke(inputs)
                 # Attempt to parse immediately to catch bad JSON before returning
-                parse_agent_response(response.content) 
+                json.loads(response.content) # Validate JSON structure
                 return response
             except ResourceExhausted as e:
                 print(f"[Retry {attempt+1}/{retries}] Gemini rate limit hit. Waiting {delay}s...")
@@ -370,7 +353,6 @@ def reasoning_node(state: AgentState) -> AgentState:
                 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"],
@@ -379,12 +361,11 @@ def reasoning_node(state: AgentState) -> AgentState:
         "current_thoughts": state["current_thoughts"]
     })
 
-    # 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]}...")
+    print(f"DEBUG: LLM Raw Response Content: {content[:200]}...")
+    print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
 
     # Update state
     state["history"].append(AIMessage(content=content)) # Store the raw LLM response
@@ -393,50 +374,54 @@ def reasoning_node(state: AgentState) -> AgentState:
     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
+        # The should_continue check will handle ending the graph based on final_answer presence
     else:
-        # Store the action request in history for tool_node
-        state["history"].append({
-            "type": "action_request",
+        # Store the action request in context, not in history
+        state["context"]["pending_action"] = {
             "tool": action,
             "input": action_input
-        })
+        }
+        # Add a message to history to indicate the agent's intent for the LLM
+        state["history"].append(AIMessage(content=f"Agent decided to use tool: {action} with input: {action_input}"))
+
 
     print(f"DEBUG: Exiting reasoning_node. New history length: {len(state['history'])}")
     return state
 
 
 def tool_node(state: AgentState) -> AgentState:
-    # ... (previous code)
+    """
+    Node for executing the chosen tool and returning its output.
+    """
+    print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
 
-    tool_call_dict = None
-    for msg in reversed(state["history"]):
-        if isinstance(msg, dict) and msg.get("type") == "action_request":
-            tool_call_dict = msg
-            break
+    # Get the pending action from context
+    tool_call_dict = state["context"].pop("pending_action", None)
 
     if not tool_call_dict:
-        print("WARNING: No action_request found in history, skipping tool execution.")
-        return state # Or raise a more specific error if this truly shouldn't happen
+        # This case should ideally not be reached if should_continue is robust,
+        # but provides a fallback.
+        error_message = "[Tool Error] No pending_action found in context. This indicates an issue with graph flow."
+        print(f"ERROR: {error_message}")
+        state["history"].append(AIMessage(content=error_message))
+        return state
 
     tool_name = tool_call_dict.get("tool")
     tool_input = tool_call_dict.get("input")
 
-    # --- ADD THIS DEBUG PRINT ---
-    print(f"DEBUG: tool_node received action_request: tool='{tool_name}', input='{tool_input[:100]}...'")
-    # --- END DEBUG PRINT ---
-
-    if not tool_name or tool_input is None: # tool_input can be empty string for some tools, but not None
-        print(f"ERROR: Invalid tool call in action_request. Tool name: '{tool_name}', Input: '{tool_input}'")
-        # Instead of raising directly, you might want to send this back to reasoning
-        # Or provide a specific error message as tool output
-        state["history"].append(AIMessage(content=f"[Tool Error] Invalid tool call: Tool name '{tool_name}' or input was empty. LLM needs to provide valid action."))
+    # Defensive check for empty tool name or input (still needed as LLM might generate empty strings)
+    if not tool_name or tool_input is None:
+        error_message = f"[Tool Error] Invalid action request from LLM: Tool name '{tool_name}' or input '{tool_input}' was empty. LLM needs to provide valid 'Action' and 'Action Input'."
+        print(f"ERROR: {error_message}")
+        state["history"].append(AIMessage(content=error_message))
+        # Clear any problematic pending action
+        state["context"].pop("pending_action", None)
         return state
 
     # 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
+    tool_fn = next((t for t in available_tools if t.name == tool_name), None)
 
     if tool_fn is None:
         # Fallback for unrecognized tool - feedback to LLM
@@ -445,17 +430,14 @@ def tool_node(state: AgentState) -> AgentState:
     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 = tool_fn.run(tool_input)
+            if not tool_output and tool_output is not False: # Ensure 'False' is not treated as empty
                 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.
+    # Add tool output to history as an AIMessage for the LLM to process next
     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'])}")
@@ -463,7 +445,7 @@ def tool_node(state: AgentState) -> AgentState:
 
 
 # ====== Agent Graph ======
-def create_agent_workflow(tools: List[Any]): # tools are passed in now
+def create_agent_workflow(tools: List[Tool]): # tools are passed in now
     workflow = StateGraph(AgentState)
     
     # Define nodes
@@ -478,43 +460,34 @@ def create_agent_workflow(tools: List[Any]): # tools are passed in now
         "reason",
         should_continue,
         {
-            "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
+            "action": "action",
+            "reason": "reason",
+            "end": END
         }
     )
     
-    workflow.add_edge("action", "reason") # Always go back to reasoning after a tool action
+    workflow.add_edge("action", "reason")
     
-    # 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):
-        # Tools need to be LangChain Tool objects for name and description
-        from langchain.tools import Tool
         self.tools = [
             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="document_qa", func=document_qa, description="Answers questions based on the content of a given document file (PDF, DOCX, TXT). Requires 'document_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
+        self.workflow = create_agent_workflow(self.tools)
 
     def __call__(self, question: str) -> str:
-        print(f"\n--- Agent received question: {question[:50]}{'...' if len(question) > 50 else ''} ---")
+        print(f"\n--- Agent received question: {question[:80]}{'...' if len(question) > 80 else ''} ---")
 
-        # Initialize state with proper structure and pass tools
         state = {
             "question": question,
             "context": {},
@@ -524,28 +497,22 @@ class BasicAgent:
             "final_answer": None,
             "current_task": "Understand the question and plan the next step.",
             "current_thoughts": "",
-            "tools": self.tools # Pass tools into the state
+            "tools": self.tools
         }
 
-        # Invoke the workflow
+        # The invoke method will now return the final state, or raise an error if it hits a dead end
+        # LangGraph runs are synchronous by default here.
         final_state = self.workflow.invoke(state)
 
-        # Extract the FINAL ANSWER from history
-        if final_state.get("final_answer"):
+        if final_state.get("final_answer") is not None:
             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 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.")
-
+        else:
+            # This should ideally not happen if the agent is designed to always provide a final answer
+            # or a specific "cannot answer" message.
+            print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
+            raise ValueError("Agent finished without providing a final answer.")