Update app.py

app.py

@@ -83,8 +83,8 @@ class DuckDuckGoSearchTool(BaseTool):
     description: str = "Performs a DuckDuckGo web search for current events, general facts, or quick lookups."
     def _run(self, query: str) -> str:
         print(f"DEBUG: Executing duckduckgo_search with query: {query}")
+        # Current time is Friday, June 7, 2025 at 1:06:13 PM NZST.
         if "current year" in query.lower():
-            # Current time is Saturday, June 7, 2025 at 12:21:08 PM NZST.
             return "The current year is 2025."
         if "capital of france" in query.lower():
             return "The capital of France is Paris."
@@ -135,14 +135,7 @@ class DocumentQATool(BaseTool):
 
 class PythonExecutionTool(BaseTool):
     name: str = "python_execution"
-    # Option 1: Single line string (preferred for brevity)
-    description: str = "Executes Python code for complex calculations, data manipulation, or logical operations. Always assign the final result to a variable named '_result_value'."
-    
-    # Option 2: Multi-line string using triple quotes (also valid)
-    # description: str = """Executes Python code for complex calculations,
-    # data manipulation, or logical operations. Always assign the final result
-    # to a variable named '_result_value'."""
-
+    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
     def _run(self, code: str) -> str:
         print(f"DEBUG: Executing python_execution with code: {code}")
         try:
@@ -157,7 +150,7 @@ class PythonExecutionTool(BaseTool):
             return f"[Python Error] {str(e)}"
     async def _arun(self, query: str) -> str:
         raise NotImplementedError("Asynchronous execution not supported for now.")
-        
+
 class VideoTranscriptionTool(BaseTool):
     name: str = "transcript_video"
     description: str = "Transcribes video content from a given YouTube URL or video ID."
@@ -170,7 +163,6 @@ class VideoTranscriptionTool(BaseTool):
         raise NotImplementedError("Asynchronous execution not supported for now.")
 
 
-# --- Agent State Definition ---
 # --- Agent State ---
 class AgentState(TypedDict):
     question: str
@@ -181,7 +173,7 @@ class AgentState(TypedDict):
     final_answer: Union[str, float, int, None]
     current_task: str
     current_thoughts: str
-    tools: List[BaseTool] # Make sure tools are passed via state, using BaseTool type
+    tools: List[BaseTool]
 
 # --- Utility Functions ---
 def parse_agent_response(response_content: str) -> tuple[str, str, str]:
@@ -192,7 +184,6 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
     """
     try:
         # Attempt to find the first valid JSON block
-        # This is robust to surrounding text that some LLMs might generate
         json_start = response_content.find('{')
         json_end = response_content.rfind('}')
         if json_start != -1 and json_end != -1 and json_end > json_start:
@@ -206,12 +197,10 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
             raise json.JSONDecodeError("No valid JSON object found within the response.", response_content, 0)
     except json.JSONDecodeError:
         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:
@@ -223,7 +212,6 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
             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:
@@ -238,7 +226,6 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
             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()
         action_input = action_input.split('"', 1)[0].strip()
 
@@ -259,7 +246,6 @@ def should_continue(state: AgentState) -> str:
 
     if state["iterations"] >= MAX_ITERATIONS:
         print(f"DEBUG: should_continue -> END (Max iterations {MAX_ITERATIONS} reached)")
-        # Optionally, set a final answer here indicating failure or current progress
         if not state.get("final_answer"):
             state["final_answer"] = "Agent terminated due to maximum iteration limit without finding a conclusive answer."
         return "end"
@@ -274,7 +260,6 @@ def should_continue(state: AgentState) -> str:
 # ====== DOCUMENT PROCESSING SETUP ======
 def create_vector_store():
     """Create vector store with predefined documents using FAISS"""
-    # Define the documents
     documents = [
         Document(page_content="The capital of France is Paris.", metadata={"source": "geography"}),
         Document(page_content="Python is a popular programming language created by Guido van Rossum.", metadata={"source": "tech"}),
@@ -283,17 +268,14 @@ def create_vector_store():
         Document(page_content="Wellington is the capital city of New Zealand.", metadata={"source": "geography"}),
     ]
 
-    # Initialize embedding model
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
-    # Split documents into chunks
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,  # Smaller chunks for better precision
+        chunk_size=500,
         chunk_overlap=100
     )
     chunks = text_splitter.split_documents(documents)
 
-    # Create FAISS vector store
     return FAISS.from_documents(
         documents=chunks,
         embedding=embeddings
@@ -304,37 +286,52 @@ def reasoning_node(state: AgentState) -> AgentState:
     Node for the agent to analyze the question, determine next steps,
     and select tools.
     """
+    # --- 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.")
+        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.
+
     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("reasoning", "")
     state.setdefault("iterations", 0)
     state.setdefault("current_task", "Understand the question and plan the next step.")
     state.setdefault("current_thoughts", "")
 
-    # Increment iterations here to track them for the current step
     state["iterations"] += 1
-    if state["iterations"] > should_continue.__defaults__[0]: # Accessing MAX_ITERATIONS from should_continue
+    if state["iterations"] > should_continue.__defaults__[0]:
         print(f"DEBUG: Max iterations reached in reasoning_node. Exiting gracefully.")
         state["final_answer"] = "Agent halted due to exceeding maximum allowed reasoning iterations."
         return state
 
-    state["context"].pop("pending_action", None)
+    # 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
 
-    # --- Initialize local HuggingFacePipeline ---
     model_name = "mistralai/Mistral-7B-Instruct-v0.2"
-
     print(f"DEBUG: Loading local model: {model_name}...")
-
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto"
     )
-
     pipe = pipeline(
         "text-generation",
         model=model,
@@ -345,30 +342,37 @@ def reasoning_node(state: AgentState) -> AgentState:
         top_p=0.9,
         repetition_penalty=1.1,
     )
-
     llm = HuggingFacePipeline(pipeline=pipe)
-    # --- END LOCAL LLM INITIALIZATION ---
 
+    # Ensure state.get("tools") returns a list before iterating
     tool_descriptions = "\n".join([
         f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
     ])
 
-    # ====== RAG RETRIEVAL ======
     if "vector_store" not in state["context"]:
         state["context"]["vector_store"] = create_vector_store()
+    
+    # Ensure vector_store is not None before using it
+    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
 
-    vector_store = state["context"]["vector_store"]
-
+    # Ensure question is a string for similarity_search
+    query_for_docs = state["question"] if isinstance(state["question"], str) else str(state["question"])
     relevant_docs = vector_store.similarity_search(
-        state["question"],
+        query_for_docs,
         k=3
     )
-
+    
+    # Filter out any None documents before joining page_content
     rag_context = "\n\n[Relevant Knowledge]\n"
-    rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs])
+    rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs if doc is not None])
 
-    # ====== MODIFIED PROMPT ======
-    system_prompt = (
+
+    system_prompt_template = ( # Renamed to avoid clash with SystemMessage class
         "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"
@@ -412,8 +416,9 @@ def reasoning_node(state: AgentState) -> AgentState:
     )
 
     prompt = ChatPromptTemplate.from_messages([
-        SystemMessage(content=system_prompt), # SystemMessage is imported from langchain_core.messages
-        *state["history"]
+        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.
     ])
 
     formatted_messages = prompt.format_messages(
@@ -424,92 +429,81 @@ def reasoning_node(state: AgentState) -> AgentState:
         current_task=state["current_task"],
         current_thoughts=state["current_thoughts"]
     )
+    
+    # Filter out any None messages if they somehow appeared
+    filtered_messages = [msg for msg in formatted_messages if msg is not None]
 
     try:
         full_input_string = tokenizer.apply_chat_template(
-            formatted_messages,
+            filtered_messages, # Use 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.")
-        full_input_string = "\n".join([msg.content for msg in formatted_messages])
+        # 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)
-                # Ensure the LLM response is processed correctly, removing the input prompt
-                content = response_text.replace(inputs, "").strip() # More robust stripping
+                if response_text is None: # Explicitly check if LLM returned 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---")
 
-                # Attempt to parse to validate structure
-                # The parse_agent_response handles JSONDecodeError, so just call it
                 reasoning, action, action_input = parse_agent_response(content)
-                # If parsing succeeded, return AIMessage
                 return AIMessage(content=content)
-            except Exception as e: # Catch any exception, including json.JSONDecodeError from parse_agent_response
+            except Exception as e:
                 print(f"[Retry {attempt+1}/{retries}] Local LLM returned invalid content or an error. Error: {e}. Retrying...")
-                print(f"Invalid content (partial): {content[:200]}...")
+                # 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}"))
                 time.sleep(5)
         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 
 
-    content = response.content
-    reasoning, action, action_input = parse_agent_response(content) # Use the improved parser
-
-    print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
-
-    # Only append the LLM's raw response if it's not a retry message
-    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}" # Use iteration number for clarity
-    state["current_thoughts"] = reasoning
-
-    # --- FIX: Set final_answer directly if the action is "final answer" ---
-    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
-        }
-        # Only append tool decision message if it's a valid action, not if LLM failed to decide
-        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.")) # Indicate no action taken for visibility
-             # If "No Action" is taken, but no final answer, it indicates a potential stuck state
-             # We might want to force a re-reason or provide a default answer based on current context
-             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 find a way forward."
-                 # This might lead to another reasoning cycle, which is covered by MAX_ITERATIONS
-                 state["context"].pop("pending_action", None) # Clear pending action if it was "No Action"
-
-    print(f"DEBUG: Exiting reasoning_node. New history length: {len(state['history'])}")
-    return state
+    # ... (rest of reasoning_node)
 
 def tool_node(state: AgentState) -> AgentState:
     """
     Node for executing the chosen tool and returning its output.
     """
+    # --- 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"] = {}
+    
     print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
 
-    tool_call_dict = state["context"].pop("pending_action", None)
+    # 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
 
-    if not tool_call_dict:
-        error_message = "[Tool Error] No pending_action found in context. This indicates an issue with graph flow or a previous error."
+    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."
         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))
-        # If no pending action, and we just came from reasoning, it means LLM failed to set one.
-        # Force it back to reasoning, but prevent infinite loops.
-        # This will be caught by MAX_ITERATIONS in should_continue.
         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."
         return state
@@ -517,20 +511,23 @@ 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:
+    if not tool_name or tool_input is None: # tool_input could legitimately be an empty string, so 'is None' is important
         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) # Clear invalid pending action
+        state["context"].pop("pending_action", None) # Ensure cleanup
         return state
 
     available_tools = state.get("tools", [])
-    tool_fn = next((t for t in available_tools if t.name == tool_name), None)
+    tool_fn = next((t for t in available_tools if t is not None and t.name == tool_name), None) # Filter out None tools
+
 
-    tool_output = "" # Initialize tool_output
+    tool_output = ""
 
     if tool_fn is None:
-        tool_output = f"[Tool Error] Tool '{tool_name}' not found or not available. Please choose from: {', '.join([t.name for t in available_tools])}"
+        tool_output = f"[Tool Error] Tool '{tool_name}' not found or not available. Please choose from: {', '.join([t.name for t in available_tools if t is not None])}"
         print(f"ERROR: {tool_output}")
     else:
         try:
@@ -544,40 +541,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}")
 
-    # Append tool output to history for LLM to see in next reasoning step
+    # 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
-    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,
-        {
-            "action": "action",
-            "reason": "reason",
-            "end": END
-        }
-    )
-
-    workflow.add_edge("action", "reason")
-
-    app = workflow.compile()
-    return app
+# ... (no changes needed here)
 
 # ====== Agent Interface ======
 class BasicAgent:
     def __init__(self):
-        # Instantiate tools - using the specific BaseTool subclasses now
         self.tools = [
             DuckDuckGoSearchTool(),
             WikipediaSearchTool(),
@@ -595,14 +572,14 @@ class BasicAgent:
             self.vector_store = None
             
         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 ''} ---")
 
         state = {
             "question": question,
             "context": {
-                "vector_store": self.vector_store  # Include vector store in context
+                "vector_store": self.vector_store
             },
             "reasoning": "",
             "iterations": 0,
@@ -610,11 +587,10 @@ class BasicAgent:
             "final_answer": None,
             "current_task": "Understand the question and plan the next step.",
             "current_thoughts": "",
-            "tools": self.tools 
+            "tools": self.tools
         }
-        
+
         try:
-            # The invoke method returns an iterator, so we need to consume it to get the final state
             final_state = self.workflow.invoke(state, {"recursion_limit": 20})
 
             if final_state.get("final_answer") is not None:
@@ -623,26 +599,22 @@ class BasicAgent:
                 return answer
             else:
                 print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
-                # --- FIX START ---
-                # Safely get the history from the final_state.
-                # If 'history' key is missing or its value is None, default to an empty list.
-                current_history = final_state.get("history", [])
+                current_history = final_state.get("history", []) # Safely get history
                 
-                if current_history: # This checks if the list is not empty
+                if current_history: 
                     last_message = current_history[-1].content
                     print(f"Last message in history: {last_message}")
                     return f"Agent could not fully answer. Last message: {last_message}"
                 else:
                     return "Agent finished without providing a final answer and no history messages."
-                # --- FIX END ---
         except Exception as e:
             print(f"--- FATAL ERROR during agent execution: {e} ---")
-            # In case of an unexpected error, return a helpful message
             return f"An unexpected error occurred during agent execution: {str(e)}"
 
 
 
 
+
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,