Update app.py

app.py

@@ -148,6 +148,8 @@ class VideoTranscriptionTool(BaseTool):
 
 
 
+
+
 import os
 import time
 import json
@@ -203,20 +205,22 @@ class VideoTranscriptionTool:
 # --- Agent State Definition ---
 class AgentState(TypedDict):
     question: str
-    history: List[Union[HumanMessage, AIMessage]] # History only contains proper messages
-    context: Dict[str, Any] # Use context for internal agent state
+    history: List[Union[HumanMessage, AIMessage]]
+    context: Dict[str, Any]
     reasoning: str
     iterations: int
     final_answer: Union[str, float, int, None]
     current_task: str
     current_thoughts: str
-    tools: List[Tool] # Pass tools into state
+    tools: List[Tool]
 
 
 # --- 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.
+    Returns (reasoning, action, action_input).
+    If JSON parsing fails, it attempts heuristic parsing.
     """
     try:
         response_json = json.loads(response_content)
@@ -225,17 +229,46 @@ 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:
-        print(f"WARNING: LLM response not perfectly JSON: {response_content[:200]}...")
-        # 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 ""
-
-        action_part_match = response_content.split("Action:", 1)
-        action_part = action_part_match[1].strip() if len(action_part_match) > 1 else ""
+        print(f"WARNING: JSONDecodeError: LLM response was not valid JSON. Attempting heuristic parse: {response_content[:200]}...")
+        # Heuristic parsing for non-JSON or partial JSON responses
+        reasoning = ""
+        action = ""
+        action_input = ""
+
+        # Attempt to find Reasoning
+        reasoning_idx = response_content.find("Reasoning:")
+        action_idx = response_content.find("Action:")
+        if reasoning_idx != -1 and action_idx != -1 and reasoning_idx < action_idx:
+            reasoning = response_content[reasoning_idx + len("Reasoning:"):action_idx].strip()
+            # Clean up leading/trailing quotes if present
+            if reasoning.startswith('"') and reasoning.endswith('"'):
+                reasoning = reasoning[1:-1]
+        elif reasoning_idx != -1: # If only reasoning is found
+             reasoning = response_content[reasoning_idx + len("Reasoning:"):].strip()
+             if reasoning.startswith('"') and reasoning.endswith('"'):
+                reasoning = reasoning[1:-1]
+
+
+        # Attempt to find Action and Action Input
+        if action_idx != -1:
+            action_input_idx = response_content.find("Action Input:", action_idx)
+            if action_input_idx != -1:
+                action_part = response_content[action_idx + len("Action:"):action_input_idx].strip()
+                action = action_part
+                action_input = response_content[action_input_idx + len("Action Input:"):].strip()
+            else:
+                action = response_content[action_idx + len("Action:"):].strip()
+            
+            # Clean up action and action_input
+            if action.startswith('"') and action.endswith('"'):
+                action = action[1:-1]
+            if action_input.startswith('"') and action_input.endswith('"'):
+                action_input = action_input[1:-1]
+
+        # Final cleanup for any trailing JSON artifacts if heuristic grabs too much
+        action = action.split('"', 1)[0].strip() # Stop at first quote for safety if it's "Action": "tool_name",
+        action_input = action_input.split('"', 1)[0].strip() # Similar for input
 
-        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
 
 
@@ -247,17 +280,14 @@ def should_continue(state: AgentState) -> str:
     """
     print(f"DEBUG: Entering should_continue. Current context: {state.get('context', {})}")
     
-    # End if agent has produced a final answer
-    if state.get("final_answer") is not None: # Check for None explicitly
+    if state.get("final_answer") is not None:
         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"
 
@@ -270,11 +300,9 @@ def reasoning_node(state: AgentState) -> AgentState:
     print(f"DEBUG: Entering reasoning_node. Iteration: {state['iterations']}")
     print(f"DEBUG: Current history length: {len(state.get('history', []))}")
 
-    # --- CHANGE: Use HF_TOKEN environment variable ---
     HF_TOKEN = os.getenv("HF_TOKEN")
     if not HF_TOKEN:
         raise ValueError("HF_TOKEN not set in environment variables.")
-    # --- END CHANGE ---
 
     state.setdefault("context", {})
     state.setdefault("reasoning", "")
@@ -289,9 +317,9 @@ def reasoning_node(state: AgentState) -> AgentState:
     llm = ChatHuggingFace(
         llm=HuggingFaceEndpoint(
             repo_id=model_id,
-            max_new_tokens=1024, # 512
+            max_new_tokens=1024, # Increased max_new_tokens
             temperature=0.1,
-            huggingfacehub_api_token=HF_TOKEN, # --- CHANGE: Pass HF_TOKEN here ---
+            huggingfacehub_api_token=HF_TOKEN,
         )
     )
 
@@ -340,14 +368,20 @@ def reasoning_node(state: AgentState) -> AgentState:
         for attempt in range(retries):
             try:
                 response = chain.invoke(inputs)
-                json.loads(response.content) 
+                # --- NEW DEBUGGING PRINT ---
+                print(f"DEBUG: RAW LLM Response (Attempt {attempt+1}):\n---\n{response.content}\n---")
+                # --- END NEW DEBUGGING PRINT ---
+                json.loads(response.content) # Attempt to parse to validate structure
                 return response
             except json.JSONDecodeError as e:
-                print(f"[Retry {attempt+1}/{retries}] LLM returned invalid JSON. Retrying...")
-                print(f"Invalid JSON content: {response.content[:200]}...")
+                print(f"[Retry {attempt+1}/{retries}] LLM returned invalid JSON. Error: {e}. Retrying...")
+                print(f"Invalid JSON content (partial): {response.content[:200]}...")
+                # Add specific error message to history to guide LLM
+                state["history"].append(AIMessage(content=f"[Parsing Error] The previous LLM output was not valid JSON. Expected format: ```json{{\"Reasoning\": \"...\", \"Action\": \"...\", \"Action Input\": \"...\"}}```. Please ensure your response is ONLY valid JSON and strictly follows the format. Error: {e}"))
                 time.sleep(5) 
             except Exception as e: 
                 print(f"[Retry {attempt+1}/{retries}] An unexpected error occurred during LLM call: {e}. Waiting {delay}s...")
+                state["history"].append(AIMessage(content=f"[API Error] Failed to get a response from the LLM due to an API error: {e}. Trying again."))
                 time.sleep(delay)
         raise RuntimeError("Failed after multiple retries due to Hugging Face API issues or invalid JSON.")
 
@@ -362,10 +396,13 @@ def reasoning_node(state: AgentState) -> AgentState:
     content = response.content
     reasoning, action, action_input = parse_agent_response(content)
     
-    print(f"DEBUG: LLM Raw Response Content: {content[:200]}...")
     print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
 
-    state["history"].append(AIMessage(content=content)) 
+    # Only append the LLM's raw output if it was valid JSON and processed successfully
+    # Otherwise, the specific error message from the retry loop will already be in history.
+    if isinstance(response, AIMessage) and content == response.content: # Check if it's the original response, not an error message
+        state["history"].append(AIMessage(content=content)) 
+        
     state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
     state["current_thoughts"] = reasoning 
@@ -377,6 +414,8 @@ def reasoning_node(state: AgentState) -> AgentState:
             "tool": action,
             "input": action_input
         }
+        # Add a message to history to indicate the agent's intent for the LLM
+        # This will be shown to the LLM in the next turn.
         state["history"].append(AIMessage(content=f"Agent decided to use tool: {action} with input: {action_input}"))
 
 
@@ -403,7 +442,7 @@ def tool_node(state: AgentState) -> AgentState:
 
     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}")
+        print(f"ERROR: {error_output}") # Print the error message for debugging
         state["history"].append(AIMessage(content=error_message))
         state["context"].pop("pending_action", None)
         return state
@@ -491,8 +530,14 @@ class BasicAgent:
             return answer
         else:
             print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
-            raise ValueError("Agent finished without providing a final answer.")
-    
+            # You might want to return the last message from history if no final answer was set
+            if final_state["history"]:
+                last_message = final_state["history"][-1].content
+                print(f"Last message in history: {last_message}")
+                return f"Agent could not fully answer. Last message: {last_message}"
+            else:
+                raise ValueError("Agent finished without providing a final answer and no history messages.")
+
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):