Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import os
 import gradio as gr
 import requests
@@ -15,24 +16,6 @@ from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_community.document_loaders import ArxivLoader
 
-from typing import List, Union, Dict, Any, TypedDict # Ensure all types are imported
-
-import torch
-from langchain_core.messages import AIMessage, HumanMessage # Corrected import for message types
-from langchain_core.tools import BaseTool
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document
-# No longer needed: from langchain.chains.Youtubeing import load_qa_chain (as it's unused)
-from langchain_community.llms import HuggingFacePipeline
-from langchain.prompts import ChatPromptTemplate # SystemMessage moved to langchain_core.messages
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langgraph.graph import END, StateGraph
-
-# --- Import for actual YouTube transcription (if you make the tool functional) ---
-# from youtube_transcript_api import YouTubeTranscriptApi
-
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -48,6 +31,8 @@ import json
 from typing import TypedDict, List, Union, Any, Dict, Optional
 
 # LangChain and LangGraph imports
+from langchain.schema import HumanMessage, AIMessage, SystemMessage
+from langchain.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph, END
 from langchain_community.llms import HuggingFacePipeline 
 
@@ -65,540 +50,509 @@ import arxiv
 from transformers import pipeline as hf_pipeline # Renamed to avoid clash with main pipeline
 from youtube_transcript_api import YouTubeTranscriptApi
 
-from typing import List, Literal, TypedDict
-
-
-
 # --- Helper function for python_execution tool ---
 def indent_code(code: str, indent: str = "    ") -> str:
     """Indents multi-line code for execution within a function."""
     return "\n".join(indent + line for line in code.splitlines())
 
 # --- Tool Definitions ---
-import wikipedia # <--- Make sure you have this installed: pip install wikipedia
-
-# --- Dummy Tools (replace with actual, robust implementations for full functionality) ---
-class DuckDuckGoSearchTool(BaseTool):
-    name: str = "duckduckgo_search"
-    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():
-            return "The current year is 2025."
-        if "capital of france" in query.lower():
-            return "The capital of France is Paris."
-        if "python creator" in query.lower():
-            return "Python was created by Guido van Rossum."
-        return f"Search result for '{query}': Information about {query}."
-    async def _arun(self, query: str) -> str:
-        raise NotImplementedError("Asynchronous execution not supported for now.")
-
-class WikipediaSearchTool(BaseTool):
-    name: str = "wikipedia_search"
-    description: str = "Performs a Wikipedia search for encyclopedic information, historical context, or detailed topics. Returns the first 3 sentences of the summary."
-    def _run(self, query: str) -> str:
-        print(f"DEBUG: wikipedia_search called with: {query}")
-        try:
-            return wikipedia.summary(query, sentences=3)
-        except wikipedia.DisambiguationError as e:
-            return f"Disambiguation options: {', '.join(e.options[:3])}. Please refine your query."
-        except wikipedia.PageError:
-            return "Wikipedia page not found for your query."
-        except Exception as e:
-            return f"Error performing Wikipedia search: {str(e)}"
-    async def _arun(self, query: str) -> str:
-        raise NotImplementedError("Asynchronous execution not supported for now.")
-
-class ArxivSearchTool(BaseTool):
-    name: str = "arxiv_search"
-    description: str = "Searches ArXiv for scientific papers, research, or cutting-edge technical information."
-    def _run(self, query: str) -> str:
-        print(f"DEBUG: Executing arxiv_search with query: {query}")
-        return f"ArXiv result for '{query}': Scientific papers related to {query}."
-    async def _arun(self, query: str) -> str:
-        raise NotImplementedError("Asynchronous execution not supported for now.")
-
-class DocumentQATool(BaseTool):
-    name: str = "document_qa"
-    description: str = "Answers questions based on provided document text. Input format: 'document_text||question'."
-    def _run(self, input_str: str) -> str:
-        print(f"DEBUG: Executing document_qa with input: {input_str}")
-        if "||" not in input_str:
-            return "[Error] Invalid input for document_qa. Expected 'document_text||question'."
-        doc_text, question = input_str.split("||", 1)
-        if "Paris" in doc_text and "capital" in question:
-            return "The capital of France is Paris."
-        return f"Answer to '{question}' from document: '{doc_text[:50]}...' is not directly found."
-    async def _arun(self, query: str) -> str:
-        raise NotImplementedError("Asynchronous execution not supported for now.")
-
-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'."
-    def _run(self, code: str) -> str:
-        print(f"DEBUG: Executing python_execution with code: {code}")
-        try:
-            local_vars = {}
-            exec(code, globals(), local_vars)
-            if '_result_value' in local_vars:
-                return str(local_vars['_result_value'])
-            return "Python code executed successfully, but no _result_value was assigned."
-        except Exception as e:
-            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."
-    def _run(self, query: str) -> str:
-        print(f"DEBUG: Executing transcript_video with query: {query}")
-        if "youtube.com" in query or "youtu.be" in query:
-            return f"Transcription of YouTube video '{query}': This is a sample transcription of the video content."
-        return "[Error] Invalid input for transcript_video. Please provide a valid YouTube URL or video ID."
-    async def _arun(self, query: str) -> str:
-        raise NotImplementedError("Asynchronous execution not supported for now.")
-
-
-# --- Agent State ---
-class AgentState(TypedDict):
-    question: str
-    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[BaseTool]
-
-# --- 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.
-    """
+@tool 
+def duckduckgo_search(query: str) -> str:
+    """Search web using DuckDuckGo. Returns top 3 results."""
+    print(f"DEBUG: duckduckgo_search called with: {query}")
     try:
-        # Attempt to find the first valid JSON block
-        json_start = response_content.find('{')
-        json_end = response_content.rfind('}')
-        if json_start != -1 and json_end != -1 and json_end > json_start:
-            json_str = response_content[json_start : json_end + 1]
-            response_json = json.loads(json_str)
-            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
-        else:
-            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]}...")
-        reasoning = ""
-        action = ""
-        action_input = ""
-
-        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()
-            if reasoning.startswith('"') and reasoning.endswith('"'):
-                reasoning = reasoning[1:-1]
-        elif reasoning_idx != -1:
-            reasoning = response_content[reasoning_idx + len("Reasoning:"):].strip()
-            if reasoning.startswith('"') and reasoning.endswith('"'):
-                reasoning = reasoning[1:-1]
-
-        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()
+        with DDGS() as ddgs:
+            return "\n\n".join(
+                f"Title: {res['title']}\nURL: {res['href']}\nSnippet: {res['body']}"
+                for res in ddgs.text(query, max_results=3)
+            )
+    except Exception as e:
+        return f"Error performing DuckDuckGo search: {str(e)}"
 
-            if action.startswith('"') and action.endswith('"'):
-                action = action[1:-1]
-            if action_input.startswith('"') and action_input.endswith('"'):
-                action_input = action_input[1:-1]
+@tool 
+def wikipedia_search(query: str) -> str:
+    """Get Wikipedia summaries. Returns first 3 sentences."""
+    print(f"DEBUG: wikipedia_search called with: {query}")
+    try:
+        return wikipedia.summary(query, sentences=3)
+    except wikipedia.DisambiguationError as e:
+        return f"Disambiguation options: {', '.join(e.options[:3])}"
+    except wikipedia.PageError:
+        return "Wikipedia page not found."
+    except Exception as e:
+        return f"Error performing Wikipedia search: {str(e)}"
 
-        action = action.split('"', 1)[0].strip()
-        action_input = action_input.split('"', 1)[0].strip()
+@tool 
+def arxiv_search(query: str) -> str:
+    """Search academic papers on arXiv. Returns top 3 results."""
+    print(f"DEBUG: arxiv_search called with: {query}")
+    try:
+        results = arxiv.Search(
+            query=query, 
+            max_results=3,
+            sort_by=arxiv.SortCriterion.Relevance
+        ).results()
+        
+        return "\n\n".join(
+            f"Title: {r.title}\nAuthors: {', '.join(a.name for a in r.authors)}\n"
+            f"Published: {r.published.strftime('%Y-%m-%d')}\nSummary: {r.summary[:250]}..."
+            for r in results
+        )
+    except Exception as e:
+        return f"Error performing ArXiv search: {str(e)}"
 
-        return reasoning, action, action_input
+@tool 
+def document_qa(input_str: str) -> str:
+    """Answer questions from documents. Input format: 'document_text||question'"""
+    print(f"DEBUG: document_qa called with: {input_str}")
+    try:
+        if '||' not in input_str:
+            return "Invalid format. Input must be: 'document_text||question'"
+        
+        context, question = input_str.split('||', 1)
+        # Load QA model on first call or ensure it's loaded once globally.
+        # It's better to load once in __init__ for BasicAgent if possible,
+        # but this lazy loading prevents initial heavy load if tool is not used.
+        qa_model = hf_pipeline('question-answering', model='deepset/roberta-base-squad2')
+        return qa_model(question=question, context=context)['answer']
+    except Exception as e:
+        return f"Error answering question from document: {str(e)}"
 
-# --- Graph Nodes ---
-def should_continue(state: AgentState) -> str:
-    """
-    Determines if the agent should continue reasoning, use a tool, or end.
-    Includes a maximum iteration limit to prevent infinite loops.
+@tool 
+def python_execution(code: str) -> str:
+    """Execute Python code and return output.
+    The code should assign its final result to a variable named '_result_value'.
+    Example: '_result_value = 1 + 1'
     """
-    MAX_ITERATIONS = 8 # Set a sensible limit to prevent infinite loops
-    print(f"DEBUG: Entering should_continue. Iteration: {state['iterations']}. Current context: {state.get('context', {})}")
+    print(f"DEBUG: python_execution called with: {code}")
+    try:
+        # Create isolated environment
+        env = {}
+        # Wrap code in a function to isolate scope and capture '_result_value'
+        # The exec function is used carefully here. In a production environment,
+        # consider a more robust and secure sandbox (e.g., Docker, dedicated service).
+        exec(f"def __exec_fn__():\n{indent_code(code)}\n_result_value = __exec_fn__()", globals(), env)
+        return str(env.get('_result_value', 'No explicit result assigned to "_result_value" variable.'))
+    except Exception as e:
+        return f"Python execution error: {str(e)}"
 
-    if state.get("final_answer") is not None:
-        print("DEBUG: should_continue -> END (Final Answer set in state)")
-        return "end"
+class VideoTranscriptionTool(BaseTool):
+    name: str = "transcript_video"
+    # CORRECTED LINE BELOW: Added '=' for assignment
+    description: str = "Fetch text transcript from YouTube videos using URL or ID. Use for any question involving video or audio. Input is the YouTube URL or ID." 
+
+    def _run(self, url_or_id: str) -> str:
+        print(f"DEBUG: transcript_video called with: {url_or_id}")
+        video_id = None
+        # Basic parsing for common YouTube URL formats
+        if "youtube.com/watch?v=" in url_or_id:
+            video_id = url_or_id.split("v=")[1].split("&")[0]
+        elif "youtu.be/" in url_or_id:
+            video_id = url_or_id.split("youtu.be/")[1].split("?")[0]
+        elif len(url_or_id.strip()) == 11 and not ("http://" in url_or_id or "https://" in url_or_id):
+            video_id = url_or_id.strip() # Assume it's just the ID
+
+        if not video_id:
+            return f"Invalid or unsupported YouTube URL/ID: {url_or_id}. Please provide a valid YouTube URL or 11-character ID."
 
-    if state["iterations"] >= MAX_ITERATIONS:
-        print(f"DEBUG: should_continue -> END (Max iterations {MAX_ITERATIONS} reached)")
-        if not state.get("final_answer"):
-            state["final_answer"] = "Agent terminated due to maximum iteration limit without finding a conclusive answer."
-        return "end"
+        try:
+            transcription = YouTubeTranscriptApi.get_transcript(video_id)
+            return " ".join([part['text'] for part in transcription])
 
-    if state.get("context", {}).get("pending_action"):
-        print("DEBUG: should_continue -> ACTION (Pending action in context)")
-        return "action"
+        except Exception as e:
+            return f"Error fetching transcript for video ID '{video_id}': {str(e)}. It might not have an English transcript, or the video is unavailable."
 
-    print("DEBUG: should_continue -> REASON (Default to reasoning)")
-    return "reason"
+    def _arun(self, *args, **kwargs):
+        raise NotImplementedError("Async not supported for this tool.")
+
+
+# ====== IMPORTS ======
+import json
+import time
+from typing import Dict, List, Tuple, Any, Optional
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+
+# ====== TYPE DEFINITIONS ======
+AgentState = Dict[str, Any]
 
-# ====== DOCUMENT PROCESSING SETUP ======
-def create_vector_store():
+# ====== DOCUMENT PROCESSING ======
+def create_vector_store() -> Optional[FAISS]:
     """Create vector store with predefined documents using FAISS"""
-    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"}),
-        Document(page_content="The Eiffel Tower is located in Paris, France.", metadata={"source": "landmarks"}),
-        Document(page_content="The highest mountain in New Zealand is Aoraki/Mount Cook.", metadata={"source": "geography"}),
-        Document(page_content="Wellington is the capital city of New Zealand.", metadata={"source": "geography"}),
-    ]
-
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=100
-    )
-    chunks = text_splitter.split_documents(documents)
+    try:
+        # 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"}),
+            Document(page_content="The Eiffel Tower is located in Paris, France.", metadata={"source": "landmarks"}),
+        ]
+        
+        # Initialize embedding model
+        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        
+        # Split documents into chunks
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500,
+            chunk_overlap=100
+        )
+        chunks = text_splitter.split_documents(documents)
+        
+        # Create FAISS vector store
+        return FAISS.from_documents(
+            documents=chunks,
+            embedding=embeddings
+        )
+    except Exception as e:
+        print(f"ERROR creating vector store: {str(e)}")
+        return None
 
-    return FAISS.from_documents(
-        documents=chunks,
-        embedding=embeddings
-    )
+# ====== AGENT HELPER FUNCTIONS ======
+def parse_agent_response(content: str) -> Tuple[str, str, str]:
+    """Parse the agent's JSON response"""
+    try:
+        # Extract JSON from content
+        json_start = content.find('{')
+        json_end = content.rfind('}') + 1
+        json_str = content[json_start:json_end]
+        
+        # Parse JSON
+        response_dict = json.loads(json_str)
+        
+        reasoning = response_dict.get("Reasoning", "No reasoning provided")
+        action = response_dict.get("Action", "No action specified")
+        action_input = response_dict.get("Action Input", "No input provided")
+        
+        return reasoning, action, action_input
+    except json.JSONDecodeError:
+        print(f"WARNING: Failed to parse JSON from response: {content[:200]}...")
+        return "Failed to parse response", "Final Answer", "Error: Could not parse agent response"
+    except Exception as e:
+        print(f"ERROR parsing agent response: {str(e)}")
+        return "Error in parsing", "Final Answer", f"Internal error: {str(e)}"
 
+def should_continue(state: AgentState) -> str:
+    """Determine if we should continue processing or end"""
+    if state.get("final_answer"):
+        return "end"
+    if state.get("context", {}).get("pending_action"):
+        return "action"
+    return "reason"
+
+# ====== REASONING NODE ======
 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.")
-
-    # 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"] = {}
+    """Node for analyzing questions and determining next steps"""
+    print(f"DEBUG: Entering reasoning_node. Iteration: {state.get('iterations', 0)}")
     
-    # 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", {}) # Redundant if check above re-initializes, but harmless
+    # Safely initialize state components
+    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", "")
-
-    state["iterations"] += 1
-    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
-
-    # Now that context is guaranteed a dict, this is safe
+    state.setdefault("history", [])
+    
+    # Safely remove pending_action
     state["context"].pop("pending_action", None) 
 
+    # 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,
-        tokenizer=tokenizer,
-        max_new_tokens=1024,
-        temperature=0.1,
-        do_sample=True,
-        top_p=0.9,
-        repetition_penalty=1.1,
-    )
-    llm = HuggingFacePipeline(pipeline=pipe)
+    
+    try:
+        print(f"DEBUG: Loading local model: {model_name}...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Determine torch dtype based on available hardware
+        if torch.cuda.is_available():
+            torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        else:
+            torch_dtype = torch.float32
+            
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            torch_dtype=torch_dtype,
+            device_map="auto"
+        )
+
+        # Create transformers pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=1024,
+            temperature=0.1,
+            do_sample=True,
+            top_p=0.9,
+            repetition_penalty=1.1,
+        )
+
+        llm = HuggingFacePipeline(pipeline=pipe)
+    except Exception as e:
+        print(f"ERROR loading model: {str(e)}")
+        state["history"].append(AIMessage(content=f"[ERROR] Failed to load model: {str(e)}"))
+        state["final_answer"] = "Error: Failed to initialize language model"
+        return state
 
-    # Ensure state.get("tools") returns a list before iterating and that items are not None
+    # Prepare tool descriptions
     tool_descriptions = "\n".join([
-        f"- **{t.name}**: {t.description}" for t in state.get("tools", []) if t is not None
+        f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
     ])
 
-    if "vector_store" not in state["context"]:
-        state["context"]["vector_store"] = create_vector_store()
-    
-    # Ensure vector_store is not None before using it
+    # RAG Retrieval
+    rag_context = ""
     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.")
-        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.get("question"), str) else str(state["question"])
-    relevant_docs = vector_store.similarity_search(
-        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 if doc is not None])
-
+    if vector_store:
+        try:
+            # Perform retrieval
+            relevant_docs = vector_store.similarity_search(
+                state.get("question", ""), 
+                k=3
+            )
+            
+            # Format context for LLM
+            rag_context = "\n\n[Relevant Knowledge]\n"
+            rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs])
+        except Exception as e:
+            print(f"WARNING: RAG retrieval failed: {str(e)}")
+            rag_context = "\n\n[Relevant Knowledge] Retrieval failed. Proceeding without additional context."
+    else:
+        print("WARNING: No vector store available for RAG")
+        rag_context = "\n\n[Relevant Knowledge] No knowledge base available."
 
-    system_prompt_template = ( 
+    # Enhanced system prompt with RAG context
+    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. Provide a concise search query. Example: `What is the population of New York?`\n"
-        "- Use **wikipedia_search** for encyclopedic information, historical context, or detailed topics. Provide a concise search term. Example: `Eiffel Tower history`\n"
-        "- Use **arxiv_search** for scientific papers, research, or cutting-edge technical information. Provide a concise search query. Example: `Large Language Models recent advances`\n"
-        "- Use **document_qa** when the question explicitly refers to a specific document or when you have content to query. Input format: 'document_text||question'. Example: `The capital of France is Paris.||What is the capital of France?`\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_value'. Example: `_result_value = 1 + 1`\n"
-        "- Use **transcript_video** for any question involving video or audio content (e.g., YouTube). Provide the full YouTube URL or video ID. Example: `youtube.com`\n\n"
-        "**Crucial Instructions:**\n"
-        "1. **Always aim to provide a definitive answer.** If you have enough information, use the 'final answer' action.\n"
-        "2. **To provide a final answer, use the Action 'final answer' with the complete answer in 'Action Input'.** This is how you tell me you're done. Example:\n"
-        "   ```json\n"
-        "   {\n"
-        "     \"Reasoning\": \"I have found the capital of France.\",\n"
-        "     \"Action\": \"final answer\",\n"
-        "     \"Action Input\": \"The capital of France is Paris.\"\n"
-        "   }\n"
-        "   ```\n"
-        "3. **If you need more information or cannot answer yet, select an appropriate tool and provide a clear, concise query.**\n"
-        "4. **Think step-by-step.** Reflect on previous tool outputs and the question.\n"
-        "5. **Do NOT repeat actions or search queries unless the previous attempt yielded an error.**\n\n"
+        "- Use **duckduckgo_search** for current events, general facts, or quick lookups. Provide a concise search query.\n"
+        "- Use **wikipedia_search** for encyclopedic information, historical context, or detailed topics. Provide a concise search term.\n"
+        "- Use **arxiv_search** for scientific papers, research, or cutting-edge technical information. Provide a concise search query.\n"
+        "- Use **document_qa** when the question explicitly refers to a specific document or when you have content to query. Input format: 'document_text||question'.\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_value' (e.g., '_result_value = 1 + 1').\n"
+        "- Use **transcript_video** for any question involving video or audio content (e.g., YouTube). Provide the full YouTube URL or video ID.\n\n"
         "**Retrieved Context:**\n{rag_context}\n\n"
-        "**Current Context (Tool Outputs/Intermediate Info):**\n{context}\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"
-        "**Question:** {question}\n\n"
-        "**Expected JSON Output Format:**\n"
+        "**Your Response MUST be a valid JSON object with the following keys:**\n"
         "```json\n"
         "{\n"
-        "  \"Reasoning\": \"Your reasoning process to decide the next step, including why a tool is chosen or how an answer is derived.\",\n"
-        "  \"Action\": \"The name of the tool to use (e.g., duckduckgo_search, final answer, No Action), if no tool is needed yet, use 'No Action'.\",\n"
-        "  \"Action Input\": \"The input for the tool (e.g., 'What is the capital of France?', 'The final answer is Paris.').\"\n"
+        "  \"Reasoning\": \"Your detailed analysis of the question and why you chose a specific action. Focus on the logical steps.\",\n"
+        "  \"Action\": \"[Tool name OR 'Final Answer']\",\n"
+        "  \"Action Input\": \"[Input for the selected tool OR the complete final answer]\"\n"
         "}\n"
         "```\n"
-        "Ensure your response is ONLY valid JSON and strictly follows this format. Begin your response with ````json`."
+        "**CRITICAL RULE: 'Action' and 'Action Input' MUST NOT be empty strings, unless 'Action' is 'Final Answer' and 'Action Input' is the conclusive response.**\n"
+        "If you cannot determine a suitable tool or a conclusive final answer after exhausting options, 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 always the complete, valid input for the chosen 'Action'. If 'Action' is 'Final Answer', provide the complete, concise answer."
     )
 
+    # Create prompt
     prompt = ChatPromptTemplate.from_messages([
-        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.
+        SystemMessage(content=system_prompt),
+        *state["history"]
     ])
 
+    # Format messages safely
     formatted_messages = prompt.format_messages(
         rag_context=rag_context,
-        context=state["context"],
-        reasoning=state["reasoning"],
-        question=state["question"],
-        current_task=state["current_task"],
-        current_thoughts=state["current_thoughts"]
+        context=state.get("context", {}),
+        reasoning=state.get("reasoning", ""),
+        question=state.get("question", ""),
+        current_task=state.get("current_task", ""),
+        current_thoughts=state.get("current_thoughts", "")
     )
     
-    # Filter out any None messages if they somehow appeared before tokenization
-    filtered_messages = [msg for msg in formatted_messages if msg is not None]
-
+    # Format full input string
     try:
         full_input_string = tokenizer.apply_chat_template(
-            filtered_messages, 
-            tokenize=False,
+            formatted_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 filtered_messages if msg is not None])
+        print(f"WARNING: Failed to apply chat template: {e}. Using simple join.")
+        full_input_string = "\n".join([msg.content for msg in formatted_messages])
 
-    def call_with_retry_local(inputs, retries=3):
+    # Call LLM with retry
+    def call_with_retry_local(inputs: str, retries: int = 3) -> AIMessage:
         for attempt in range(retries):
             try:
-                response_text = llm.invoke(inputs)
-                if response_text is None: 
-                    raise ValueError("LLM invoke returned None response_text.")
+                response_text = llm.invoke(inputs) 
                 
-                content = response_text.replace(inputs, "").strip() if isinstance(response_text, str) else str(response_text).replace(inputs, "").strip()
+                # Strip the prompt from the generated text
+                if response_text.startswith(inputs):
+                    content = response_text[len(inputs):].strip()
+                else:
+                    content = response_text.strip() 
 
                 print(f"DEBUG: RAW LOCAL LLM Response (Attempt {attempt+1}):\n---\n{content}\n---")
-
-                reasoning, action, action_input = parse_agent_response(content)
+                
+                # Attempt to parse to validate structure
+                json.loads(content) 
                 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...")
-                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)
-    content = response.content 
-
-    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']}")
+            except json.JSONDecodeError as e:
+                print(f"[Retry {attempt+1}/{retries}] Invalid JSON. Error: {e}.")
+                print(f"Invalid content: {content[:200]}...")
+                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(3) 
+            except Exception as e: 
+                print(f"[Retry {attempt+1}/{retries}] Error: {e}.")
+                state["history"].append(AIMessage(content=f"[LLM Error] Failed to get response: {e}. Trying again."))
+                time.sleep(5) 
+        return AIMessage(content='{"Reasoning": "Max retries exceeded", "Action": "Final Answer", "Action Input": "Error: Failed after multiple retries"}')
+
+    response = call_with_retry_local(full_input_string) 
+    content = response.content
+
+    # Parse response
+    reasoning, action, action_input = parse_agent_response(content)
+    
+    print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
+
+    # Update state
+    state["history"].append(AIMessage(content=content)) 
+    state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
+    state["iterations"] += 1
+    state["current_thoughts"] = reasoning 
+
+    if "final answer" in action.lower():
+        state["final_answer"] = action_input 
     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)
+        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
 
+# ====== TOOL 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.")
+    """Node for executing the chosen tool"""
+    print(f"DEBUG: Entering tool_node. Iteration: {state.get('iterations', 0)}")
     
-    # 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"] = []
+    # Safely get pending action
+    tool_call_dict = state.get("context", {}).pop("pending_action", None)
 
-    print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
-
-    # 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. This indicates an issue with graph flow or a previous error."
+    if not tool_call_dict:
+        error_message = "[Tool Error] No pending_action found in context."
         print(f"ERROR: {error_message}")
-        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."
+        state.setdefault("history", []).append(AIMessage(content=error_message))
         return state
 
-    tool_name = tool_call_dict.get("tool")
-    tool_input = tool_call_dict.get("input")
+    tool_name = tool_call_dict.get("tool", "")
+    tool_input = tool_call_dict.get("input", "")
 
-    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 not tool_name or not tool_input: 
+        error_message = f"[Tool Error] Invalid action: Tool name '{tool_name}' or input '{tool_input}' was empty."
+        print(f"ERROR: {error_message}") 
         state["history"].append(AIMessage(content=error_message))
-        state["context"].pop("pending_action", None) 
         return state
 
+    # Find and execute tool
     available_tools = state.get("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 = ""
+    tool_fn = next((t for t in available_tools if t.name == tool_name), None)
 
     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 if t is not None])}"
+        tool_output = f"[Tool Error] Tool '{tool_name}' not found. Available: {', '.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]}...'")
-            raw_tool_output = tool_fn.run(tool_input)
-            if raw_tool_output is None or raw_tool_output is False or raw_tool_output == "":
-                tool_output = f"[{tool_name} output] No specific result found for '{tool_input}'. The tool might have returned an empty response."
-            else:
-                tool_output = f"[{tool_name} output]\n{raw_tool_output}"
+            tool_output = tool_fn.run(tool_input)
+            if tool_output is None:
+                tool_output = f"[{tool_name} output] No result returned for '{tool_input}'."
         except Exception as e:
-            tool_output = f"[Tool Error] An error occurred while running '{tool_name}': {str(e)}"
+            tool_output = f"[Tool Error] Error running '{tool_name}': {str(e)}"
             print(f"ERROR: {tool_output}")
 
-    state["history"].append(AIMessage(content=tool_output))
-
-    print(f"DEBUG: Exiting tool_node. Tool output added to history. New history length: {len(state['history'])}")
+    state["history"].append(AIMessage(content=f"[{tool_name} output]\n{tool_output}"))
+    
+    print(f"DEBUG: Exiting tool_node. Tool output added to history.")
     return state
 
-# ====== Agent Graph ======
-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,
-        {
-            "action": "action",
-            "reason": "reason",
-            "end": END
-        }
-    )
-
-    workflow.add_edge("action", "reason")
+# ====== AGENT GRAPH ======
+class StateGraph:
+    """Simple state graph implementation"""
+    def __init__(self, state: AgentState):
+        self.nodes = {}
+        self.entry_point = None
+        self.edges = {}
+        self.conditional_edges = {}
+        
+    def add_node(self, name: str, func: callable):
+        self.nodes[name] = func
+        
+    def set_entry_point(self, name: str):
+        self.entry_point = name
+        
+    def add_conditional_edges(self, source: str, condition: callable, path_map: Dict[str, str]):
+        self.conditional_edges[source] = (condition, path_map)
+        
+    def add_edge(self, source: str, dest: str):
+        self.edges[source] = dest
+        
+    def compile(self):
+        def app(state: AgentState) -> AgentState:
+            current_node = self.entry_point
+            
+            while current_node != END:
+                if current_node in self.nodes:
+                    state = self.nodes[current_node](state)
+                    
+                if current_node in self.conditional_edges:
+                    condition, path_map = self.conditional_edges[current_node]
+                    next_node_key = condition(state)
+                    current_node = path_map.get(next_node_key, END)
+                elif current_node in self.edges:
+                    current_node = self.edges[current_node]
+                else:
+                    current_node = END
+                    
+            return state
+            
+        return app
 
-    app = workflow.compile()
-    return app
+END = "__END__"
 
-# ====== Agent Interface ======
+# ====== AGENT INTERFACE ======
 class BasicAgent:
     def __init__(self):
+        # Instantiate tools (implementation not shown - should be defined elsewhere)
         self.tools = [
-            DuckDuckGoSearchTool(),
-            WikipediaSearchTool(),
-            ArxivSearchTool(),
-            DocumentQATool(),
-            PythonExecutionTool(),
-            VideoTranscriptionTool()
+            # duckduckgo_search, 
+            # wikipedia_search,
+            # arxiv_search,
+            # document_qa,
+            # python_execution,
+            # VideoTranscriptionTool()
         ]
-
-        self.vector_store = create_vector_store()
+        
+        # Pre-initialize RAG vector store
+        try:
+            self.vector_store = create_vector_store()
+            if not self.vector_store:
+                print("WARNING: Vector store creation failed. Proceeding without RAG.")
+        except Exception as e:
+            print(f"ERROR creating vector store: {str(e)}")
+            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 ''} ---")
+        print(f"\n--- Agent received question: {question[:80]}{'...' if len(question) > 80 else ''} ---")
 
         state = {
             "question": question,
@@ -611,34 +565,48 @@ class BasicAgent:
             "final_answer": None,
             "current_task": "Understand the question and plan the next step.",
             "current_thoughts": "",
-            "tools": self.tools
+            "tools": self.tools 
         }
 
         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."
-
+            final_state = self.workflow.invoke(state)
             if final_state.get("final_answer") is not None:
                 answer = final_state["final_answer"]
                 print(f"--- Agent returning FINAL ANSWER: {answer} ---")
                 return answer
             else:
-                print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
-                current_history = final_state.get("history", []) # Safely get history
-                
-                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."
+                print("--- ERROR: Agent finished without setting 'final_answer' ---")
+                if final_state.get("history"):
+                    last_message = final_state["history"][-1].content
+                    print(f"Last message: {last_message}")
+                    return f"Agent could not answer. Last message: {last_message}"
+                return "Error: Agent failed to provide an answer"
         except Exception as e:
-            print(f"--- FATAL ERROR during agent execution: {e} ---")
-            return f"An unexpected error occurred during agent execution: {str(e)}"
+            print(f"FATAL ERROR during agent execution: {str(e)}")
+            return f"Agent encountered a fatal error: {str(e)}"
+
+def create_agent_workflow(tools: List[Any]):
+    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