Update app.py

app.py

@@ -135,13 +135,11 @@ class DocumentQATool(BaseTool):
 
 class PythonExecutionTool(BaseTool):
     name: str = "python_execution"
-    description: str = "Executes Python code for complex calculations, data manipulation, or logical operations. Always assign the final result to a variable named '_result_value'." # Fixed syntax error
+    description: str = "Executes Python code for complex calculations, data manipulation, or logical operations. Always assign the final result to a variable named '_result_value'."
     def _run(self, code: str) -> str:
         print(f"DEBUG: Executing python_execution with code: {code}")
         try:
             local_vars = {}
-            # It's generally unsafe to use `exec` with arbitrary user input due to security risks.
-            # For a real application, consider a sandboxed environment or a more restricted approach.
             exec(code, globals(), local_vars)
             if '_result_value' in local_vars:
                 return str(local_vars['_result_value'])
@@ -289,23 +287,29 @@ def reasoning_node(state: AgentState) -> AgentState:
     # --- Defensive checks at the start of the node ---
     if state is None:
         raise ValueError("reasoning_node received a None state object.")
-    if state.get("history") is None:
-        print("WARNING: 'history' is None on entry to reasoning_node. Re-initializing to empty list.")
-        state["history"] = []
-    if state.get("context") is None:
-        print("WARNING: 'context' is None on entry to reasoning_node. Re-initializing to empty dict.")
+
+    # Ensure context is a dictionary
+    if not isinstance(state.get("context"), dict):
+        print("WARNING: state['context'] is not a dictionary on entry to reasoning_node. Re-initializing to empty dict.")
         state["context"] = {}
-    if state.get("tools") is None:
-        print("WARNING: 'tools' is None on entry to reasoning_node. This might cause issues.")
-        # If tools are None, the tool_descriptions generation below will fail.
-        # It's highly unlikely given the BasicAgent init, but good to check.
+    
+    # Ensure history is a list
+    if not isinstance(state.get("history"), list):
+        print("WARNING: state['history'] is not a list on entry to reasoning_node. Re-initializing to empty list.")
+        state["history"] = []
+
+    # Ensure tools is a list
+    if not isinstance(state.get("tools"), list):
+        print("WARNING: state['tools'] is not a list on entry to reasoning_node. This might cause issues downstream.")
+        # If tools become None or corrupted, the tool_descriptions part will fail.
+        # It's better to log and proceed, assuming agent init sets them correctly.
 
     print(f"DEBUG: Entering reasoning_node. Iteration: {state['iterations']}")
     # Use .get() for safety when printing history length
     print(f"DEBUG: Current history length: {len(state.get('history', []))}")
 
     # Set defaults for state components that might be missing, although TypedDict implies presence
-    state.setdefault("context", {})
+    state.setdefault("context", {}) # Redundant if check above re-initializes, but harmless
     state.setdefault("reasoning", "")
     state.setdefault("iterations", 0)
     state.setdefault("current_task", "Understand the question and plan the next step.")
@@ -317,12 +321,8 @@ def reasoning_node(state: AgentState) -> AgentState:
         state["final_answer"] = "Agent halted due to exceeding maximum allowed reasoning iterations."
         return state
 
-    # Ensure context is a dict before popping
-    if isinstance(state["context"], dict):
-        state["context"].pop("pending_action", None)
-    else:
-        print("WARNING: state['context'] is not a dictionary in reasoning_node. Cannot pop pending_action.")
-        state["context"] = {} # Re-initialize if it's corrupted
+    # Now that context is guaranteed a dict, this is safe
+    state["context"].pop("pending_action", None) 
 
     model_name = "mistralai/Mistral-7B-Instruct-v0.2"
     print(f"DEBUG: Loading local model: {model_name}...")
@@ -344,9 +344,9 @@ def reasoning_node(state: AgentState) -> AgentState:
     )
     llm = HuggingFacePipeline(pipeline=pipe)
 
-    # Ensure state.get("tools") returns a list before iterating
+    # Ensure state.get("tools") returns a list before iterating and that items are not None
     tool_descriptions = "\n".join([
-        f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
+        f"- **{t.name}**: {t.description}" for t in state.get("tools", []) if t is not None
     ])
 
     if "vector_store" not in state["context"]:
@@ -356,12 +356,11 @@ def reasoning_node(state: AgentState) -> AgentState:
     vector_store = state["context"].get("vector_store")
     if vector_store is None:
         print("ERROR: Vector store is None after creation/retrieval in reasoning_node. Cannot perform similarity search.")
-        # Handle this error more gracefully, e.g., return an error state or raise exception
         state["final_answer"] = "Internal error: Vector store not available."
         return state
 
     # Ensure question is a string for similarity_search
-    query_for_docs = state["question"] if isinstance(state["question"], str) else str(state["question"])
+    query_for_docs = state["question"] if isinstance(state.get("question"), str) else str(state["question"])
     relevant_docs = vector_store.similarity_search(
         query_for_docs,
         k=3
@@ -372,7 +371,7 @@ def reasoning_node(state: AgentState) -> AgentState:
     rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs if doc is not None])
 
 
-    system_prompt_template = ( # Renamed to avoid clash with SystemMessage class
+    system_prompt_template = ( 
         "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"
@@ -416,9 +415,8 @@ def reasoning_node(state: AgentState) -> AgentState:
     )
 
     prompt = ChatPromptTemplate.from_messages([
-        SystemMessage(content=system_prompt_template), # Use the template here
-        *state["history"] # This assumes state["history"] is always an iterable (list).
-                          # The check at the start of the node handles if it's None.
+        SystemMessage(content=system_prompt_template), 
+        *state["history"] # This assumes state["history"] is a list. The check at the start of the node handles if it's None.
     ])
 
     formatted_messages = prompt.format_messages(
@@ -430,28 +428,26 @@ def reasoning_node(state: AgentState) -> AgentState:
         current_thoughts=state["current_thoughts"]
     )
     
-    # Filter out any None messages if they somehow appeared
+    # Filter out any None messages if they somehow appeared before tokenization
     filtered_messages = [msg for msg in formatted_messages if msg is not None]
 
     try:
         full_input_string = tokenizer.apply_chat_template(
-            filtered_messages, # Use filtered messages
+            filtered_messages, 
             tokenize=False,
             add_generation_prompt=True
         )
     except Exception as e:
         print(f"WARNING: Failed to apply chat template: {e}. Falling back to simple string join. Model performance may be affected.")
-        # Filter again just in case, before accessing .content
         full_input_string = "\n".join([msg.content for msg in filtered_messages if msg is not None])
 
     def call_with_retry_local(inputs, retries=3):
         for attempt in range(retries):
             try:
                 response_text = llm.invoke(inputs)
-                if response_text is None: # Explicitly check if LLM returned None
+                if response_text is None: 
                     raise ValueError("LLM invoke returned None response_text.")
                 
-                # Ensure response_text is a string before calling .replace()
                 content = response_text.replace(inputs, "").strip() if isinstance(response_text, str) else str(response_text).replace(inputs, "").strip()
 
                 print(f"DEBUG: RAW LOCAL LLM Response (Attempt {attempt+1}):\n---\n{content}\n---")
@@ -460,7 +456,6 @@ def reasoning_node(state: AgentState) -> AgentState:
                 return AIMessage(content=content)
             except Exception as e:
                 print(f"[Retry {attempt+1}/{retries}] Local LLM returned invalid content or an error. Error: {e}. Retrying...")
-                # Safely preview content for debugging
                 safe_content_preview = content[:200] if isinstance(content, str) else "Content was not a string or is None."
                 print(f"Invalid content (partial): {safe_content_preview}...")
                 state["history"].append(AIMessage(content=f"[Parsing Error] The previous LLM output was not valid. Expected format: ```json{{\"Reasoning\": \"...\", \"Action\": \"...\", \"Action Input\": \"...\"}}```. Please ensure your response is ONLY valid JSON and strictly follows the format. Error: {e}"))
@@ -468,10 +463,33 @@ def reasoning_node(state: AgentState) -> AgentState:
         raise RuntimeError("Failed after multiple retries due to local Hugging Face model issues or invalid JSON.")
 
     response = call_with_retry_local(full_input_string)
-    # If response is None, it would have been caught by the ValueError in call_with_retry_local
     content = response.content 
 
-    # ... (rest of reasoning_node)
+    if not content.startswith("[Parsing Error]") and not content.startswith("[Local LLM Error]"):
+        state["history"].append(AIMessage(content=content))
+
+    state["reasoning"] += f"\nStep {state['iterations']}: {reasoning}"
+    state["current_thoughts"] = reasoning
+
+    if action.lower() == "final answer":
+        state["final_answer"] = action_input
+        print(f"DEBUG: Final answer set in state: {state['final_answer']}")
+    else:
+        state["context"]["pending_action"] = {
+            "tool": action,
+            "input": action_input
+        }
+        if action and action != "No Action":
+            state["history"].append(AIMessage(content=f"Agent decided to use tool: {action} with input: {action_input}"))
+        elif action == "No Action":
+             state["history"].append(AIMessage(content=f"Agent decided to take 'No Action' but needs to proceed."))
+             if not state.get("final_answer"):
+                 state["current_task"] = "Re-evaluate the situation and attempt to find a final answer or a new tool."
+                 state["current_thoughts"] = "The previous step resulted in 'No Action'. I need to re-think my next step."
+                 state["context"].pop("pending_action", None)
+
+    print(f"DEBUG: Exiting reasoning_node. New history length: {len(state['history'])}")
+    return state
 
 def tool_node(state: AgentState) -> AgentState:
     """
@@ -480,29 +498,25 @@ def tool_node(state: AgentState) -> AgentState:
     # --- Defensive checks at the start of the node ---
     if state is None:
         raise ValueError("tool_node received a None state object.")
-    if state.get("history") is None:
-        print("WARNING: 'history' is None on entry to tool_node. Re-initializing to empty list.")
-        state["history"] = []
-    if state.get("context") is None:
-        print("WARNING: 'context' is None on entry to tool_node. Re-initializing to empty dict.")
-        state["context"] = {}
     
+    # Ensure context is a dictionary
+    if not isinstance(state.get("context"), dict):
+        print("WARNING: state['context'] is not a dictionary on entry to tool_node. Re-initializing to empty dict.")
+        state["context"] = {}
+
+    # Ensure history is a list
+    if not isinstance(state.get("history"), list):
+        print("WARNING: state['history'] is not a list on entry to tool_node. Re-initializing to empty list.")
+        state["history"] = []
+
     print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
 
-    # Safely access tool_call_dict. Ensure state["context"] is a dictionary before pop.
-    tool_call_dict = None
-    if isinstance(state["context"], dict):
-        tool_call_dict = state["context"].pop("pending_action", None)
-    else:
-        print("WARNING: state['context'] is not a dictionary in tool_node. Cannot pop pending_action.")
-        state["context"] = {} # Re-initialize if it's corrupted
+    # Safely access tool_call_dict. Context is guaranteed to be a dict here.
+    tool_call_dict = state["context"].pop("pending_action", None) 
 
     if tool_call_dict is None:
-        error_message = "[Tool Error] No pending_action found in context or context was invalid. This indicates an issue with graph flow or a previous error."
+        error_message = "[Tool Error] No pending_action found in context. This indicates an issue with graph flow or a previous error."
         print(f"ERROR: {error_message}")
-        # Ensure state["history"] is a list before appending
-        if state.get("history") is None:
-            state["history"] = []
         state["history"].append(AIMessage(content=error_message))
         state["current_task"] = "Re-evaluate the situation; previous tool selection failed or was missing."
         state["current_thoughts"] = "No tool action was found. I need to re-think my next step."
@@ -511,18 +525,16 @@ def tool_node(state: AgentState) -> AgentState:
     tool_name = tool_call_dict.get("tool")
     tool_input = tool_call_dict.get("input")
 
-    if not tool_name or tool_input is None: # tool_input could legitimately be an empty string, so 'is None' is important
+    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 or None. LLM needs to provide valid 'Action' and 'Action Input'."
         print(f"ERROR: {error_message}")
-        if state.get("history") is None:
-            state["history"] = []
         state["history"].append(AIMessage(content=error_message))
-        state["context"].pop("pending_action", None) # Ensure cleanup
+        state["context"].pop("pending_action", None) 
         return state
 
     available_tools = state.get("tools", [])
-    tool_fn = next((t for t in available_tools if t is not None and t.name == tool_name), None) # Filter out None tools
-
+    # Filter out any None tools before iterating
+    tool_fn = next((t for t in available_tools if t is not None and t.name == tool_name), None) 
 
     tool_output = ""
 
@@ -541,24 +553,20 @@ def tool_node(state: AgentState) -> AgentState:
             tool_output = f"[Tool Error] An error occurred while running '{tool_name}': {str(e)}"
             print(f"ERROR: {tool_output}")
 
-    # Ensure state["history"] is a list before appending
-    if state.get("history") is None:
-        state["history"] = []
     state["history"].append(AIMessage(content=tool_output))
 
     print(f"DEBUG: Exiting tool_node. Tool output added to history. New history length: {len(state['history'])}")
     return state
 
-
 # ====== Agent Graph ======
-def create_agent_workflow(tools: List[BaseTool]): # Use BaseTool for consistency
+def create_agent_workflow(tools: List[BaseTool]):
     workflow = StateGraph(AgentState)
-    
+
     workflow.add_node("reason", reasoning_node)
     workflow.add_node("action", tool_node)
-    
+
     workflow.set_entry_point("reason")
-    
+
     workflow.add_conditional_edges(
         "reason",
         should_continue,
@@ -568,13 +576,12 @@ def create_agent_workflow(tools: List[BaseTool]): # Use BaseTool for consistency
             "end": END
         }
     )
-    
+
     workflow.add_edge("action", "reason")
-    
+
     app = workflow.compile()
     return app
 
-
 # ====== Agent Interface ======
 class BasicAgent:
     def __init__(self):
@@ -587,15 +594,9 @@ class BasicAgent:
             VideoTranscriptionTool()
         ]
 
-        # Pre-initialize RAG vector store
-        try:
-            self.vector_store = create_vector_store()
-        except Exception as e:
-            print(f"ERROR: Failed to create vector store: {str(e)}")
-            self.vector_store = None
-            
+        self.vector_store = create_vector_store()
         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 ''} ---")
 
@@ -616,6 +617,11 @@ class BasicAgent:
         try:
             final_state = self.workflow.invoke(state, {"recursion_limit": 20})
 
+            # It's highly unlikely final_state would be None if invoke completes,
+            # but this check is harmless and covers an extreme edge case.
+            if final_state is None:
+                return "Agent workflow completed but returned a None state. This is unexpected."
+
             if final_state.get("final_answer") is not None:
                 answer = final_state["final_answer"]
                 print(f"--- Agent returning FINAL ANSWER: {answer} ---")
@@ -637,7 +643,6 @@ class BasicAgent:
 
 
 
-
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,