multiple revisions

app.py

@@ -4,7 +4,7 @@ import requests
 import inspect
 import pandas as pd
 import traceback
-from tools import intelligent_agent, get_file_type, write_bytes_to_temp_dir, State
+from tools import create_memory_safe_workflow, get_file_type, write_bytes_to_temp_dir, AgentState, extract_final_answer, run_agent
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -87,24 +87,27 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         tb_str = traceback.format_exc()
         print(f"Error creating new states: {tb_str}")
         return f"Error creating new states: {tb_str}", None
+    
+    agent = create_memory_safe_workflow()    
 
-    # 3. Setup states for questions and run agent
+    # Setup states for questions and run agent
     answers_payload = []
     results_log = []
     for r in range(len(hf_questions)):
-            s = State(question = hf_questions[r]['question'],
-                        input_file = hf_questions[r]['input_file'],
-                        file_type = hf_questions[r]['file_type'],
-                        file_path = hf_questions[r]['file_path'])
-            try:
-                task_id = hf_questions[r]['task_id']
-                question_text = hf_questions[r]['question']
-                submitted_answer = intelligent_agent(s)
-                answers_payload.append({"task_id": task_id, "model_answer": submitted_answer})
-                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
-            except:
-                print(f"Error running agent on task {task_id}: {e}")
-                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
+        s = AgentState(question = hf_questions[r]['question'],
+                    input_file = hf_questions[r]['input_file'],
+                    file_type = hf_questions[r]['file_type'],
+                    file_path = hf_questions[r]['file_path'])
+        try:
+            task_id = hf_questions[r]['task_id']
+            question_text = hf_questions[r]['question']
+            full_answer = run_agent(agent, s)
+            submitted_answer = extract_final_answer(full_answer[-1].content)
+            answers_payload.append({"task_id": task_id, "model_answer": submitted_answer})
+            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
+        except:
+            print(f"Error running agent on task {task_id}: {e}")
+            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
 
     if not answers_payload:
         print("Agent did not produce any answers to submit.")

requirements.txt

@@ -19,3 +19,4 @@ accelerate
 en_core_web_sm @ https://huggingface.co/spacy/en_core_web_sm/resolve/main/en_core_web_sm-any-py3-none-any.whl
 transformers==4.40.0
 datasets==2.19.0
+beautifulsoup4

tools.py

@@ -1,45 +1,85 @@
-import numpy as np
-import spacy
+# Standard Library
+import os
+import re
 import tempfile
+import string 
 import glob
-import yt_dlp
 import shutil
+import gc
+import uuid
+import signal
+from datetime import datetime
+from io import BytesIO
+from contextlib import contextmanager
+from langchain_huggingface import HuggingFacePipeline
+from typing import TypedDict, List, Optional, Dict, Any, Annotated, Literal, Union, Tuple, Set
+import time
+from collections import Counter
+
+# Third-Party Packages
 import cv2
-import librosa
+import requests
 import wikipedia
+import spacy
+import yt_dlp
+import librosa
+from PIL import Image
+from bs4 import BeautifulSoup
+from duckduckgo_search import DDGS
+from sentence_transformers import SentenceTransformer
+from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline
 
-from typing import TypedDict, List, Optional, Dict, Any
+# LangChain Ecosystem
 from langchain.docstore.document import Document
 from langchain.prompts import PromptTemplate
 from langchain_community.document_loaders import WikipediaLoader
-from langgraph.graph import START, END, StateGraph
-from langchain_core.messages import AnyMessage, HumanMessage, AIMessage # If you are using it
-from langchain_community.retrievers import BM25Retriever # If you are using it
-from langgraph.prebuilt import ToolNode, tools_condition # If you are using it
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_community.retrievers import BM25Retriever
 from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema import Document
-from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline
-from io import BytesIO
-from sentence_transformers import SentenceTransformer
-from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
-from transformers import AutoTokenizer, AutoModelWithLMHead
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, BaseMessage, SystemMessage, ToolMessage
+from langchain_core.tools import BaseTool, StructuredTool, tool, render_text_description
+from langchain_core.documents import Document
 
+# LangGraph
+from langgraph.graph import START, END, StateGraph
+from langgraph.prebuilt import ToolNode, tools_condition
 
-import os
-import re
-from PIL import Image  # This is correctly imported, but was being used incorrectly
-import numpy as np
-from collections import Counter
+# PyTorch
 import torch
-from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline
-from typing import TypedDict, List, Optional, Dict, Any, Literal, Tuple
-from langgraph.graph import StateGraph, START, END
-from langchain.docstore.document import Document
+from functools import partial
+from transformers import pipeline
+
+# Additional Utilities
+from datetime import datetime
+
+from urllib.parse import urljoin, urlparse
+import logging
 
 
 nlp = spacy.load("en_core_web_sm")
 
+logger = logging.getLogger(__name__)
+
+# --- Model Configuration ---
+def create_llm_pipeline():
+    #model_id = "meta-llama/Llama-2-13b-chat-hf"
+    #model_id = "meta-llama/Llama-3.3-70B-Instruct"
+    #model_id = "mistralai/Mistral-Small-24B-Base-2501"
+    model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+    #model_id = "Qwen/Qwen2-7B-Instruct"
+    return pipeline(
+        "text-generation",
+        model=model_id,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        max_new_tokens=1024,
+        temperature=0.1
+    )
+
 # Define file extension sets for each category
 PICTURE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
 AUDIO_EXTENSIONS = {'.mp3', '.wav', '.aac', '.flac', '.ogg', '.m4a', '.wma'}
@@ -82,8 +122,196 @@ def write_bytes_to_temp_dir(file_bytes: bytes, file_name: str) -> str:
     print(f"File written to: {file_path}")
     return file_path
 
-# 1. Define the State type
-class State(TypedDict, total=False):
+
+def extract_final_answer(text: str) -> str:
+  """
+  Returns the substring starting from the last occurrence of 'FINAL ANSWER:' (case-insensitive)
+  to the end of the string, with any trailing punctuation removed.
+  If not found, returns an empty string.
+  """
+  marker = "FINAL ANSWER:"
+  idx = text.lower().rfind(marker.lower())
+  if idx == -1:
+    return ""
+  result = text[idx:].strip()
+  # Remove trailing punctuation
+  return result.rstrip(string.punctuation + " ")
+
+
+class EnhancedDuckDuckGoSearchTool(BaseTool):
+    name: str = "enhanced_search"
+    description: str = (
+        "Performs a DuckDuckGo web search and retrieves actual content from the top web results. "
+        "Input should be a search query string. "
+        "Returns search results with extracted content from web pages, making it much more useful for answering questions. "
+        "Use this tool when you need up-to-date information, details about current events, or when other tools do not provide sufficient or recent answers. "
+        "Ideal for topics that require the latest news, recent developments, or information not covered in static sources."
+    )
+    max_results: int = 3
+    max_chars_per_page: int = 3000
+    session: Any = None  # Now it's optional and defaults to None
+
+
+    # Use model_post_init for initialization logic in Pydantic v2+
+    def model_post_init(self, __context: Any) -> None:
+        super().model_post_init(__context)
+        # Initialize HTTP session here
+        self.session = requests.Session()
+        self.session.headers.update({
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
+            'Accept-Language': 'en-US,en;q=0.5',
+            'Accept-Encoding': 'gzip, deflate',
+            'Connection': 'keep-alive',
+            'Upgrade-Insecure-Requests': '1',
+        })
+    
+    def _search_duckduckgo(self, query: str) -> List[Dict]:
+        """Perform DuckDuckGo search and return results."""
+        try:
+            with DDGS() as ddgs:
+                results = list(ddgs.text(query, max_results=self.max_results))
+                return results
+        except Exception as e:
+            logger.error(f"DuckDuckGo search failed: {e}")
+            return []
+    
+    def _extract_content_from_url(self, url: str, timeout: int = 10) -> Optional[str]:
+        """Extract clean text content from a web page."""
+        try:
+            # Skip certain file types
+            if any(url.lower().endswith(ext) for ext in ['.pdf', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx']):
+                return "Content type not supported for extraction"
+            
+            response = self.session.get(url, timeout=timeout, allow_redirects=True)
+            response.raise_for_status()
+            
+            # Check content type
+            content_type = response.headers.get('content-type', '').lower()
+            if 'text/html' not in content_type:
+                return "Non-HTML content detected"
+            
+            soup = BeautifulSoup(response.content, 'html.parser')
+            
+            # Remove script and style elements
+            for script in soup(["script", "style", "nav", "header", "footer", "aside", "form"]):
+                script.decompose()
+            
+            # Try to find main content areas
+            main_content = None
+            for selector in ['main', 'article', '.content', '#content', '.post', '.entry']:
+                main_content = soup.select_one(selector)
+                if main_content:
+                    break
+            
+            if not main_content:
+                main_content = soup.find('body') or soup
+            
+            # Extract text
+            text = main_content.get_text(separator='\n', strip=True)
+            
+            # Clean up the text
+            lines = [line.strip() for line in text.split('\n') if line.strip()]
+            text = '\n'.join(lines)
+            
+            # Remove excessive whitespace
+            text = re.sub(r'\n{3,}', '\n\n', text)
+            text = re.sub(r' {2,}', ' ', text)
+            
+            # Truncate if too long
+            if len(text) > self.max_chars_per_page:
+                text = text[:self.max_chars_per_page] + "\n[Content truncated...]"
+            
+            return text
+            
+        except requests.exceptions.Timeout:
+            return "Page loading timed out"
+        except requests.exceptions.RequestException as e:
+            return f"Failed to retrieve page: {str(e)}"
+        except Exception as e:
+            logger.error(f"Content extraction failed for {url}: {e}")
+            return "Failed to extract content from page"
+    
+    def _format_search_result(self, result: Dict, content: str) -> str:
+        """Format a single search result with its content."""
+        title = result.get('title', 'No title')
+        url = result.get('href', 'No URL')
+        snippet = result.get('body', 'No snippet')
+        
+        formatted = f"""
+🔍 **{title}**
+URL: {url}
+Snippet: {snippet}
+
+📄 **Page Content:**
+{content}
+---
+"""
+        return formatted
+    
+    def run(self, query: str) -> str:
+
+        """Execute the enhanced search."""
+        if not query or not query.strip():
+            return "Please provide a search query."
+        
+        query = query.strip()
+        logger.info(f"Searching for: {query}")
+        
+        # Perform DuckDuckGo search
+        search_results = self._search_duckduckgo(query)
+        
+        if not search_results:
+            return f"No search results found for query: {query}"
+        
+        # Process each result and extract content
+        enhanced_results = []
+        processed_count = 0
+        
+        for i, result in enumerate(search_results[:self.max_results]):
+            url = result.get('href', '')
+            if not url:
+                continue
+                
+            logger.info(f"Processing result {i+1}: {url}")
+            
+            # Extract content from the page
+            content = self._extract_content_from_url(url)
+            
+            if content and len(content.strip()) > 50:  # Only include results with substantial content
+                formatted_result = self._format_search_result(result, content)
+                enhanced_results.append(formatted_result)
+                processed_count += 1
+            
+            # Small delay to be respectful to servers
+            time.sleep(0.5)
+        
+        if not enhanced_results:
+            return f"Search completed but no content could be extracted from the pages for query: {query}"
+        
+        # Compile final response
+        response = f"""🔍 **Enhanced Search Results for: "{query}"**
+Found {len(search_results)} results, successfully processed {processed_count} pages with content.
+
+{''.join(enhanced_results)}
+
+💡 **Summary:** Retrieved and processed content from {processed_count} web pages to provide comprehensive information about your search query.
+"""
+        
+        # Ensure the response isn't too long
+        if len(response) > 8000:
+            response = response[:8000] + "\n[Response truncated to prevent memory issues]"
+        
+        return response
+
+    def _run(self, query: str) -> str:
+        """Required by BaseTool interface."""
+        return self.run(query)
+    
+# --- Agent State Definition ---
+class AgentState(TypedDict):
+    messages: Annotated[List[AnyMessage], lambda x, y: x + y]
+    done: bool = False  # Default value of False
     question: str
     task_id: str
     input_file: Optional[bytes]
@@ -93,144 +321,628 @@ class State(TypedDict, total=False):
     youtube_url: Optional[str]
     answer: Optional[str]
     frame_answers: Optional[list]
-    next: Optional[str]  # Added to track the next node
-
-# --- LLM pipeline for general questions ---
-llm_pipe = pipeline(
-    "text-generation",
-    model="microsoft/Phi-3-mini-4k-instruct",
-    device_map="auto",
-    torch_dtype="auto",
-    max_new_tokens=256,
-    trust_remote_code=True
-)
-
-# Initialize RAG components
-tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base", trust_remote_code=True)
-retriever = RagRetriever.from_pretrained(
-    "facebook/rag-token-base",
-    index_name="exact",           # or "legacy" for legacy FAISS index
-    use_dummy_dataset=False,       # set to False and download the full index for real Wikipedia retrieval
-    trust_remote_code=True,       # Trust remote code for dataset loading
-    dataset_revision="main",      # Specify a fixed revision
-    dataset="wiki_dpr",           # Explicitly specify dataset name
-)
-rag_model = RagSequenceForGeneration.from_pretrained(
-    "facebook/rag-token-base", 
-    retriever=retriever,
-    trust_remote_code=True
-)
-# Speech-to-text pipeline
-asr_pipe = pipeline(
-    "automatic-speech-recognition",
-    model="openai/whisper-small",
-    device="auto"
-)
-
-# --- BLIP VQA setup ---
-device = "cuda" if torch.cuda.is_available() else "cpu"
-vqa_model_name = "Salesforce/blip-vqa-base"
-processor_vqa = BlipProcessor.from_pretrained(vqa_model_name)
-
-# Attempt to load model to GPU; fall back to CPU if OOM
-try:
-    model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to(device)
-except torch.cuda.OutOfMemoryError:
-    print("WARNING: Loading model to CPU due to insufficient GPU memory.")
-    device = "cpu"  # Switch device to CPU
-    model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to(device)
-
-# --- Helper functions ---
-def ensure_final_answer_format(answer_text: str) -> str:
-    """Ensure the answer ends with FINAL ANSWER: format"""
-    # Check if the answer already contains a FINAL ANSWER section
-    if "FINAL ANSWER:" in answer_text:
-        # Extract everything after FINAL ANSWER:
-        final_answer_part = answer_text.split("FINAL ANSWER:", 1)[1].strip()
-        return f"FINAL ANSWER: {final_answer_part}"
-    else:
-        # If no FINAL ANSWER section exists, wrap the entire answer
-        return f"FINAL ANSWER: {answer_text.strip()}"
-
-def extract_entities(text: str) -> List[str]:
-    """Extract key entities from text using spaCy if available, or regex fallback"""
-    if nlp:
-        # Using spaCy for better entity extraction
-        doc = nlp(text)
-        entities = [ent.text for ent in doc.ents]
-        keywords = [token.text for token in doc if token.pos_ in ("PROPN", "NOUN")]
-        return entities if entities else keywords
-    else:
-        # Simple fallback using regex to extract potential keywords
-        words = text.lower().split()
-        stopwords = ["what", "who", "when", "where", "why", "how", "is", "are", "the", "a", "an", "of", "in", "on", "at"]
-        keywords = [word for word in words if word not in stopwords and len(word) > 2]
-        return keywords
 
-def answer_question_on_frame(image_path, question):
-    """Answer a question about a single video frame using BLIP"""
-    try:
-        image = Image.open(image_path).convert('RGB')
-        inputs = processor_vqa(image, question, return_tensors="pt").to(device)
-        out = model_vqa.generate(**inputs)
-        answer = processor_vqa.decode(out[0], skip_special_tokens=True)
-        return answer
-    except Exception as e:
-        print(f"Error processing frame {image_path}: {str(e)}")
-        return "Error processing this frame"
 
-def answer_video_question(frames_dir, question):
-    """Answer a question about a video by analyzing extracted frames"""
-    valid_exts = ('.jpg', '.jpeg', '.png')
+def fetch_page_with_tables(page_title):
+    """
+    Fetches Wikipedia page content and extracts all tables as readable text.
+    Returns a tuple: (main_text, [table_texts])
+    """
+    # Fetch the page object
+    page = wikipedia.page(page_title)
+    main_text = page.content
+
+    # Get the HTML for table extraction
+    html = page.html()
+    soup = BeautifulSoup(html, 'html.parser')
+    tables = soup.find_all('table')
+
+    table_texts = []
+    for table in tables:
+        rows = table.find_all('tr')
+        table_lines = []
+        for row in rows:
+            cells = row.find_all(['th', 'td'])
+            cell_texts = [cell.get_text(strip=True) for cell in cells]
+            if cell_texts:
+                # Format as Markdown table row
+                table_lines.append(" | ".join(cell_texts))
+        if table_lines:
+            table_text = "\n".join(table_lines)
+            table_texts.append(table_text)
+
+    return main_text, table_texts
+
+class WikipediaSearchToolWithFAISS(BaseTool):
+    name: str = "wikipedia_semantic_search_all_candidates_strong_entity_priority_list_retrieval"
+    description: str = (
+        "Fetches content from multiple Wikipedia pages based on intelligent NLP query processing "
+        "of various search candidates, with strong prioritization of query entities. It then performs "
+        "entity-focused semantic search across all fetched content to find the most relevant information, "
+        "with improved retrieval for lists like discographies. Uses spaCy for named entity "
+        "recognition and query enhancement. Input should be a search query or topic. "
+        "Note: Uses the current live version of Wikipedia."
+    )
+    embedding_model_name: str = "all-MiniLM-L6-v2"
+    chunk_size: int = 4000
+    chunk_overlap: int = 250 # Maintained moderate overlap
+    top_k_results: int = 3
+    spacy_model: str = "en_core_web_sm"
+    # Increased multiplier to fetch more candidates per semantic query variant
+    semantic_search_candidate_multiplier: int = 1 # Was 2, increased to 3, consider 4 if still problematic
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        try:
+            self._nlp = spacy.load(self.spacy_model)
+            print(f"Loaded spaCy model: {self.spacy_model}")
+            self._embedding_model = HuggingFaceEmbeddings(model_name=self.embedding_model_name)
+            # Refined separators for better handling of Wikipedia lists and sections
+            self._text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=self.chunk_size,
+                chunk_overlap=self.chunk_overlap,
+                separators=[
+                    "\n\n== ", "\n\n=== ", "\n\n==== ",  # Section headers (keep with following content)
+                    "\n\n\n", "\n\n",  # Multiple newlines (paragraph breaks)
+                    "\n* ", "\n- ", "\n# ", # List items
+                    "\n", ". ", "! ", "? ", # Sentence breaks after newline, common punctuation
+                    " ", "" # Word and character level
+                ]
+            )
+        except OSError as e:
+            print(f"Error loading spaCy model '{self.spacy_model}': {e}")
+            print("Try running: python -m spacy download en_core_web_sm")
+            self._nlp = None
+            self._embedding_model = None
+            self._text_splitter = None
+        except Exception as e:
+            print(f"Error initializing WikipediaSearchToolWithFAISS components: {e}")
+            self._nlp = None
+            self._embedding_model = None
+            self._text_splitter = None
+
+    def _extract_entities_and_keywords(self, query: str) -> Tuple[List[str], List[str], str]:
+        if not self._nlp:
+            return [], [], query
+        doc = self._nlp(query)
+        main_entities = [ent.text for ent in doc.ents if ent.label_ in ["PERSON", "ORG", "GPE", "EVENT", "WORK_OF_ART"]]
+        keywords = [token.lemma_.lower() for token in doc if token.pos_ in ["NOUN", "PROPN", "ADJ"] and not token.is_stop and not token.is_punct and len(token.text) > 2]
+        main_entities = list(dict.fromkeys(main_entities))
+        keywords = list(dict.fromkeys(keywords))
+        processed_tokens = [token.lemma_ for token in doc if not token.is_stop and not token.is_punct and token.text.strip()]
+        processed_query = " ".join(processed_tokens)
+        return main_entities, keywords, processed_query
+
+    def _generate_search_candidates(self, query: str, main_entities: List[str], keywords: List[str], processed_query: str) -> List[str]:
+        candidates_set = set()
+        entity_prefix = main_entities[0] if main_entities else None
+
+        for me in main_entities:
+            candidates_set.add(me)
+        candidates_set.add(query)
+        if processed_query and processed_query != query:
+            candidates_set.add(processed_query)
+        
+        if entity_prefix and keywords:
+            first_entity_lower = entity_prefix.lower()
+            for kw in keywords[:3]:
+                if kw not in first_entity_lower and len(kw) > 2:
+                    candidates_set.add(f"{entity_prefix} {kw}")
+            keyword_combo_short = " ".join(k for k in keywords[:2] if k not in first_entity_lower and len(k)>2)
+            if keyword_combo_short: candidates_set.add(f"{entity_prefix} {keyword_combo_short}")
+        
+        if len(main_entities) > 1:
+            candidates_set.add(" ".join(main_entities[:2]))
+        
+        if keywords:
+            keyword_combo = " ".join(keywords[:2])
+            if entity_prefix:
+                candidate_to_add = f"{entity_prefix} {keyword_combo}"
+                if not any(c.lower() == candidate_to_add.lower() for c in candidates_set):
+                     candidates_set.add(candidate_to_add)
+            elif not main_entities:
+                candidates_set.add(keyword_combo)
+        
+        ordered_candidates = []
+        for me in main_entities:
+            if me not in ordered_candidates: ordered_candidates.append(me)
+        for c in list(candidates_set):
+            if c and c.strip() and c not in ordered_candidates: ordered_candidates.append(c)
+        
+        print(f"Generated {len(ordered_candidates)} search candidates for Wikipedia page lookup (entity-prioritized): {ordered_candidates}")
+        return ordered_candidates
 
-    # Check if directory exists
-    if not os.path.exists(frames_dir):
-        return {
-            "most_common_answer": "No frames found to analyze.",
-            "all_answers": [],
-            "answer_counts": Counter()
+    def _smart_wikipedia_search(self, query_text: str, main_entities_from_query: List[str], keywords_from_query: List[str], processed_query_text: str) -> List[Tuple[str, str]]:
+        candidates = self._generate_search_candidates(query_text, main_entities_from_query, keywords_from_query, processed_query_text)
+        found_pages_data: List[Tuple[str, str]] = []
+        processed_page_titles: Set[str] = set()
+
+        for i, candidate_query in enumerate(candidates):
+            print(f"\nProcessing candidate {i+1}/{len(candidates)} for page: '{candidate_query}'")
+            page_object = None
+            final_page_title = None
+            is_candidate_entity_focused = any(me.lower() in candidate_query.lower() for me in main_entities_from_query) if main_entities_from_query else False
+
+            try:
+                try:
+                    page_to_load = candidate_query
+                    suggest_mode = True # Default to auto_suggest=True
+                    if is_candidate_entity_focused and main_entities_from_query:
+                        try: # Attempt precise match first for entity-focused candidates
+                            temp_page = wikipedia.page(page_to_load, auto_suggest=False, redirect=True)
+                            suggest_mode = False # Flag that precise match worked
+                        except (wikipedia.exceptions.PageError, wikipedia.exceptions.DisambiguationError):
+                            print(f"  - auto_suggest=False failed for entity-focused '{page_to_load}', trying with auto_suggest=True.")
+                            # Fallthrough to auto_suggest=True below if this fails
+                    
+                    if suggest_mode: # If not attempted or failed with auto_suggest=False
+                         temp_page = wikipedia.page(page_to_load, auto_suggest=True, redirect=True)
+                    
+                    final_page_title = temp_page.title
+                    
+                    if is_candidate_entity_focused and main_entities_from_query:
+                        title_matches_main_entity = any(me.lower() in final_page_title.lower() for me in main_entities_from_query)
+                        if not title_matches_main_entity:
+                            print(f"  ! Page title '{final_page_title}' (from entity-focused candidate '{candidate_query}') "
+                                  f"does not strongly match main query entities: {main_entities_from_query}. Skipping.")
+                            continue 
+                    if final_page_title in processed_page_titles:
+                        print(f"  ~ Already processed '{final_page_title}'")
+                        continue
+                    page_object = temp_page
+                    print(f"  ✓ Direct hit/suggestion for '{candidate_query}' -> '{final_page_title}'")
+                
+                except wikipedia.exceptions.PageError:
+                    if i < max(2, len(candidates) // 3) : # Try Wikipedia search for a smaller, more promising subset of candidates
+                        print(f"  - Direct access failed for '{candidate_query}'. Trying Wikipedia search...")
+                        search_results = wikipedia.search(candidate_query, results=1) 
+                        if not search_results:
+                            print(f"  - No Wikipedia search results for '{candidate_query}'.")
+                            continue
+                        search_result_title = search_results[0]
+                        try:
+                            temp_page = wikipedia.page(search_result_title, auto_suggest=False, redirect=True) # Search results are usually canonical
+                            final_page_title = temp_page.title
+                            if is_candidate_entity_focused and main_entities_from_query: # Still check against original intent
+                                title_matches_main_entity = any(me.lower() in final_page_title.lower() for me in main_entities_from_query)
+                                if not title_matches_main_entity:
+                                    print(f"    ! Page title '{final_page_title}' (from search for '{candidate_query}' -> '{search_result_title}') "
+                                          f"does not strongly match main query entities: {main_entities_from_query}. Skipping.")
+                                    continue
+                            if final_page_title in processed_page_titles:
+                                print(f"    ~ Already processed '{final_page_title}'")
+                                continue 
+                            page_object = temp_page
+                            print(f"    ✓ Found via search '{candidate_query}' -> '{search_result_title}' -> '{final_page_title}'")
+                        except (wikipedia.exceptions.PageError, wikipedia.exceptions.DisambiguationError) as e_sr:
+                            print(f"    ! Error/Disambiguation for search result '{search_result_title}': {e_sr}")
+                    else:
+                        print(f"  - Direct access failed for '{candidate_query}'. Skipping further search for this lower priority candidate.")
+                except wikipedia.exceptions.DisambiguationError as de:
+                    print(f"  ! Disambiguation for '{candidate_query}'. Options: {de.options[:1]}")
+                    if de.options:
+                        option_title = de.options[0]
+                        try:
+                            temp_page = wikipedia.page(option_title, auto_suggest=False, redirect=True)
+                            final_page_title = temp_page.title
+                            if is_candidate_entity_focused and main_entities_from_query: # Check against original intent
+                                title_matches_main_entity = any(me.lower() in final_page_title.lower() for me in main_entities_from_query)
+                                if not title_matches_main_entity:
+                                    print(f"    ! Page title '{final_page_title}' (from disamb. of '{candidate_query}' -> '{option_title}') "
+                                          f"does not strongly match main query entities: {main_entities_from_query}. Skipping.")
+                                    continue
+                            if final_page_title in processed_page_titles:
+                                print(f"    ~ Already processed '{final_page_title}'")
+                                continue
+                            page_object = temp_page
+                            print(f"    ✓ Resolved disambiguation '{candidate_query}' -> '{option_title}' -> '{final_page_title}'")
+                        except Exception as e_dis_opt:
+                            print(f"    ! Could not load disambiguation option '{option_title}': {e_dis_opt}")
+                
+                if page_object and final_page_title and (final_page_title not in processed_page_titles):
+                    # Extract main text
+                    main_text = page_object.content
+
+                    # Extract tables using BeautifulSoup
+                    try:
+                        html = page_object.html()
+                        soup = BeautifulSoup(html, 'html.parser')
+                        tables = soup.find_all('table')
+                        table_texts = []
+                        for table in tables:
+                            rows = table.find_all('tr')
+                            table_lines = []
+                            for row in rows:
+                                cells = row.find_all(['th', 'td'])
+                                cell_texts = [cell.get_text(strip=True) for cell in cells]
+                                if cell_texts:
+                                    table_lines.append(" | ".join(cell_texts))
+                            if table_lines:
+                                table_text = "\n".join(table_lines)
+                                table_texts.append(table_text)
+                    except Exception as e:
+                        print(f"  !! Error extracting tables for '{final_page_title}': {e}")
+                        table_texts = []
+
+                    # Combine main text and all table texts as separate chunks
+                    all_text_chunks = [main_text] + table_texts
+
+                    for chunk in all_text_chunks:
+                        found_pages_data.append((chunk, final_page_title))
+                    processed_page_titles.add(final_page_title)
+                    print(f"  -> Added page '{final_page_title}'. Main text length: {len(main_text)} | Tables extracted: {len(table_texts)}")
+            except Exception as e:
+                print(f"  !! Unexpected error processing candidate '{candidate_query}': {e}")
+        
+        if not found_pages_data: print(f"\nCould not find any new, unique, entity-validated Wikipedia pages for query '{query_text}'.")
+        else: print(f"\nFound {len(found_pages_data)} unique, validated page(s) for processing.")
+        return found_pages_data
+
+    def _enhance_semantic_search(self, query: str, vector_store, main_entities: List[str], keywords: List[str], processed_query: str) -> List[Document]:
+        core_query_parts = set()
+        core_query_parts.add(query)
+        if processed_query != query: core_query_parts.add(processed_query)
+        if keywords: core_query_parts.add(" ".join(keywords[:2]))
+
+        section_phrases_templates = []
+        lower_query_terms = set(query.lower().split()) | set(k.lower() for k in keywords)
+        
+        section_keywords_map = {
+            "discography": ["discography", "list of studio albums", "studio album titles and years", "albums by year", "album release dates", "official albums", "complete album list", "albums published"],
+            "biography": ["biography", "life story", "career details", "background history"],
+            "filmography": ["filmography", "list of films", "movie appearances", "acting roles"],
         }
+        for section_term_key, specific_phrases_list in section_keywords_map.items():
+            # Check if the key (e.g., "discography") or any of its specific phrases (e.g. "list of studio albums")
+            # are mentioned or implied by the query terms.
+            if section_term_key in lower_query_terms or any(phrase_part in lower_query_terms for phrase_part in section_term_key.split()):
+                section_phrases_templates.extend(specific_phrases_list)
+            # Also check if phrases themselves are in query terms (e.g. query "list of albums by X")
+            for phrase in specific_phrases_list:
+                if phrase in query.lower(): # Check against original query for direct phrase matches
+                    section_phrases_templates.extend(specific_phrases_list) # Add all related if one specific is hit
+                    break 
+        section_phrases_templates = list(dict.fromkeys(section_phrases_templates)) # Deduplicate
+
+        final_search_queries = set()
+        if main_entities:
+            entity_prefix = main_entities[0]
+            final_search_queries.add(entity_prefix) 
+            for part in core_query_parts:
+                final_search_queries.add(f"{entity_prefix} {part}" if entity_prefix.lower() not in part.lower() else part)
+            for phrase_template in section_phrases_templates:
+                final_search_queries.add(f"{entity_prefix} {phrase_template}")
+                if "list of" in phrase_template or "history of" in phrase_template :
+                     final_search_queries.add(f"{phrase_template} of {entity_prefix}")
+        else: 
+            final_search_queries.update(core_query_parts)
+            final_search_queries.update(section_phrases_templates)
+
+        deduplicated_queries = list(dict.fromkeys(sq for sq in final_search_queries if sq and sq.strip()))
+        print(f"Generated {len(deduplicated_queries)} semantic search query variants (list-retrieval focused): {deduplicated_queries}")
+        
+        all_results_docs: List[Document] = []
+        seen_content_hashes: Set[int] = set()
+        k_to_fetch = self.top_k_results * self.semantic_search_candidate_multiplier
+        
+        for search_query_variant in deduplicated_queries:
+            try:
+                results = vector_store.similarity_search_with_score(search_query_variant, k=k_to_fetch) 
+                print(f"  Semantic search variant '{search_query_variant}' (k={k_to_fetch}) -> {len(results)} raw chunk(s) with scores.")
+                for doc, score in results: # Assuming similarity_search_with_score returns (doc, score)
+                    content_hash = hash(doc.page_content[:250]) # Slightly more for hash uniqueness
+                    if content_hash not in seen_content_hashes:
+                        seen_content_hashes.add(content_hash)
+                        doc.metadata['retrieved_by_variant'] = search_query_variant
+                        doc.metadata['retrieval_score'] = float(score) # Store score
+                        all_results_docs.append(doc)
+            except Exception as e:
+                print(f"  Error in semantic search for variant '{search_query_variant}': {e}")
+        
+        # Sort all collected unique results by score (FAISS L2 distance is lower is better)
+        all_results_docs.sort(key=lambda x: x.metadata.get('retrieval_score', float('inf')))
+        print(f"Collected and re-sorted {len(all_results_docs)} unique chunks from all semantic query variants.")
+        
+        return all_results_docs[:self.top_k_results]
 
-    frame_files = [os.path.join(frames_dir, f) for f in os.listdir(frames_dir)
-                  if f.lower().endswith(valid_exts)]
+    def _run(self, query: str) -> str:
+        if not self._nlp or not self._embedding_model or not self._text_splitter:
+            print("ERROR: WikipediaSearchToolWithFAISS components not initialized properly.")
+            return "Error: Wikipedia tool components not initialized properly. Please check server logs."
 
-    # Sort frames properly by number
-    def get_frame_number(filename):
-        match = re.search(r'(\d+)', os.path.basename(filename))
-        return int(match.group(1)) if match else 0
+        try:
+            print(f"\n--- Running {self.name} for query: '{query}' ---")
+            main_entities, keywords, processed_query = self._extract_entities_and_keywords(query)
+            print(f"Initial NLP Analysis - Main Entities: {main_entities}, Keywords: {keywords}, Processed Query: '{processed_query}'")
+            
+            fetched_pages_data = self._smart_wikipedia_search(query, main_entities, keywords, processed_query)
+
+            if not fetched_pages_data:
+                return (f"Could not find any relevant, entity-validated Wikipedia pages for the query '{query}'. "
+                        f"Main entities sought: {main_entities}")
+
+            all_page_titles = [title for _, title in fetched_pages_data]
+            print(f"\nSuccessfully fetched content for {len(fetched_pages_data)} Wikipedia page(s): {', '.join(all_page_titles)}")
+
+            all_documents: List[Document] = []
+            for page_content, page_title in fetched_pages_data:
+                chunks = self._text_splitter.split_text(page_content)
+                if not chunks:
+                    print(f"Warning: Could not split content from Wikipedia page '{page_title}' into chunks.")
+                    continue
+                for i, chunk_text in enumerate(chunks):
+                    all_documents.append(Document(page_content=chunk_text, metadata={
+                        "source_page_title": page_title, 
+                        "original_query": query,
+                        "chunk_index": i # Add chunk index for potential debugging or ordering
+                    }))
+                print(f"Split content from '{page_title}' into {len(chunks)} chunks.")
+
+            if not all_documents:
+                return (f"Could not process content into searchable chunks from the fetched Wikipedia pages "
+                        f"({', '.join(all_page_titles)}) for query '{query}'.")
+            
+            print(f"\nTotal document chunks from all pages: {len(all_documents)}")
 
-    frame_files = sorted(frame_files, key=get_frame_number)
+            print("Creating FAISS index from content of all fetched pages...")
+            try:
+                vector_store = FAISS.from_documents(all_documents, self._embedding_model)
+                print("FAISS index created successfully.")
+            except Exception as e:
+                return f"Error creating FAISS vector store: {e}"
 
-    if not frame_files:
-        return {
-            "most_common_answer": "No valid image frames found.",
-            "all_answers": [],
-            "answer_counts": Counter()
-        }
+            print(f"\nPerforming enhanced semantic search across all collected content...")
+            try:
+                relevant_docs = self._enhance_semantic_search(query, vector_store, main_entities, keywords, processed_query)
+            except Exception as e:
+                return f"Error during semantic search: {e}"
+
+            if not relevant_docs:
+                return (f"No relevant information found within Wikipedia page(s) '{', '.join(list(dict.fromkeys(all_page_titles)))}' "
+                        f"for your query '{query}' using entity-focused semantic search with list retrieval.")
+            
+            unique_sources_in_results = list(dict.fromkeys([doc.metadata.get('source_page_title', 'Unknown Source') for doc in relevant_docs]))
+            result_header = (f"Found {len(relevant_docs)} relevant piece(s) of information from Wikipedia page(s) "
+                             f"'{', '.join(unique_sources_in_results)}' for your query '{query}':\n")
+            nlp_summary = (f"[Original Query NLP: Main Entities: {', '.join(main_entities) if main_entities else 'None'}, "
+                           f"Keywords: {', '.join(keywords[:5]) if keywords else 'None'}]\n\n")
+            result_details = []
+            for i, doc in enumerate(relevant_docs):
+                source_info = doc.metadata.get('source_page_title', 'Unknown Source')
+                variant_info = doc.metadata.get('retrieved_by_variant', 'N/A')
+                score_info = doc.metadata.get('retrieval_score', 'N/A')
+                detail = (f"Result {i+1} (source: '{source_info}', score: {score_info:.4f})\n"
+                          f"(Retrieved by: '{variant_info}')\n{doc.page_content}")
+                result_details.append(detail)
+            
+            final_result = result_header + nlp_summary + "\n\n---\n\n".join(result_details)
+            print(f"\nReturning {len(relevant_docs)} relevant chunks from {len(set(all_page_titles))} source page(s).")
+            return final_result.strip()
 
-    answers = []
-    for frame_path in frame_files:
-        try:
-            ans = answer_question_on_frame(frame_path, question)
-            answers.append(ans)
-            print(f"Processed frame: {os.path.basename(frame_path)}, Answer: {ans}")
         except Exception as e:
-            print(f"Error processing frame {frame_path}: {str(e)}")
+            import traceback
+            print(f"Unexpected error in {self.name}: {traceback.format_exc()}")
+            return f"An unexpected error occurred: {str(e)}"
+
+
+# Example of creating the tool instance:
+# wikipedia_tool_faiss = WikipediaSearchToolWithFAISS()
+
+# To use this new tool in your agent, you would replace the old
+# `wikipedia_tool` instance with `wikipedia_tool_faiss` in your `tools` list.
+# For example:
+# tools = [wikipedia_tool_faiss, search_tool]
+# Create tool instances
+#wikipedia_tool = WikipediaSearchTool()
+
+# --- Define Call LLM function ---
+
+# 3. Improved LLM call with memory management
+
+def call_llm_with_memory_management(state: AgentState, llm_model) -> AgentState: # Added llm_model parameter
+    """Call LLM with memory management, context truncation, and process response."""
+    print("Running call_llm with memory management...")
+
+    # It's crucial to work with a copy of messages for modification within this step
+    # The final state["messages"] should reflect the full history + new response.
+    original_messages = list(state["messages"])
+    messages_for_llm_processing = list(state["messages"]) # Use this for truncation logic
+
+    #ipdb.set_trace()
+
+    # --- Context Truncation Logic ---
+    system_message_content = None
+    # Check if the first message is a system message and preserve it
+    if messages_for_llm_processing and isinstance(messages_for_llm_processing[0], SystemMessage):
+        system_message_content = messages_for_llm_processing[0]
+        # Process only non-system messages for truncation count
+        regular_messages = messages_for_llm_processing[1:]
+    else:
+        regular_messages = messages_for_llm_processing
+
+    # Truncate context if too many messages (e.g., keep system + X most recent)
+    # Max 10 messages total (e.g. 1 system + 9 others)
+    max_regular_messages = 9
+    if len(regular_messages) > max_regular_messages:
+        print(f"🔄 Truncating message count: {len(messages_for_llm_processing)} -> ~{max_regular_messages + (1 if system_message_content else 0)} messages")
+        regular_messages = regular_messages[- (max_regular_messages -1):] # Keep X-1 most recent, to add user input later
+
+    # Reconstruct messages for LLM call
+    messages_for_llm = []
+    if system_message_content:
+        messages_for_llm.append(system_message_content)
+    messages_for_llm.extend(regular_messages)
+
+    # Further truncate based on character count (rough proxy for tokens)
+    total_chars = sum(len(str(msg.content)) for msg in messages_for_llm)
+    # Example character limit, adjust based on your model (e.g. 8k chars for ~4k tokens)
+    char_limit = 8000
+    if total_chars > char_limit:
+        print(f"📏 Context too long ({total_chars} chars > {char_limit}), further truncation needed")
+        # More aggressive truncation of regular messages
+        chars_to_remove = total_chars - char_limit
+        temp_regular_messages = list(regular_messages) # copy
+        while sum(len(str(m.content)) for m in temp_regular_messages) > char_limit and temp_regular_messages:
+            if system_message_content and sum(len(str(m.content)) for m in temp_regular_messages) + len(str(system_message_content.content)) <= char_limit :
+                 break # if removing one more makes it too small with system message
+            print(f"Removing message: {temp_regular_messages[0].type} - {temp_regular_messages[0].content[:50]}...")
+            temp_regular_messages.pop(0)
+
+        regular_messages = temp_regular_messages
+        messages_for_llm = [] # Rebuild
+        if system_message_content:
+            messages_for_llm.append(system_message_content)
+        messages_for_llm.extend(regular_messages)
+        print(f"Context truncated to {sum(len(str(m.content)) for m in messages_for_llm)} chars.")
+
+    new_state = state.copy()  # Start with a copy of the input state
+
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            print(f"🧹 Pre-LLM CUDA cache cleared. Memory: {torch.cuda.memory_allocated()/1024**2:.1f}MB")
+
+        print(f"Invoking LLM with {len(messages_for_llm)} messages.")
+        # This is where you call your actual LLM
+        formatted_input = "\n".join([f"[{msg.type.upper()}] {msg.content}" for msg in messages_for_llm])
+        print(f"\n\nFormatted input for LLM:\n\n{formatted_input}")
+
+        llm_response_object = llm_model.invoke(formatted_input)
+
+        #ipdb.set_trace()
+        
+        # The response_object is typically a BaseMessage subclass (e.g., AIMessage)
+        # or a string for simpler LLMs. Adapt as needed.
+        if isinstance(llm_response_object, BaseMessage):
+            ai_message_response = llm_response_object # It's already a message object
+            if not ai_message_response.content: # Ensure content is not empty
+                 ai_message_response.content = ""
+        elif hasattr(llm_response_object, 'content'): # Some models might return a custom object with a content attribute
+            ai_message_response = AIMessage(content=str(llm_response_object.content) if llm_response_object.content is not None else "")
+        else: # Assuming it's a string for basic LLMs
+            ai_message_response = AIMessage(content=str(llm_response_object) if llm_response_object is not None else "")
+
+        print(f"LLM Response: {ai_message_response.content[:300]}...") # Print a snippet
+
+        # Append the LLM's response to the original full list of messages
+        final_messages = original_messages + [ai_message_response]
+        new_state["messages"] = final_messages
+        new_state.pop("done", None)  # LLM responded, so not 'done' by default
+
+    except Exception as e:
+        print(f"LLM call failed: {e}")
+        error_message_content = f"LLM call failed with error: {str(e)}. Input consisted of {len(messages_for_llm)} messages."
+
+        if "out of memory" in str(e).lower():
+            print("🚨 CUDA OOM detected during LLM call! Implementing emergency cleanup...")
+            error_message_content = f"LLM failed due to Out of Memory: {str(e)}."
+            try:
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                gc.collect()
+            except Exception as cleanup_e:
+                print(f"Emergency OOM cleanup failed: {cleanup_e}")
+
+        # Append an error message to the original message history
+        error_ai_message = AIMessage(content=error_message_content)
+        final_messages_on_error = original_messages + [error_ai_message]
+        new_state["messages"] = final_messages_on_error
+        new_state["done"] = True  # Mark as done to prevent loops on LLM failure
+    finally:
+        try:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                print(f"🧹 Post-LLM CUDA cache cleared. Memory: {torch.cuda.memory_allocated()/1024**2:.1f}MB")
+        except Exception:
+            pass # Avoid error in cleanup hiding the main error
+
+    return new_state
+import re
+import uuid
+
+def parse_react_output(state: AgentState) -> AgentState:
+    print("Running parse_react_output (Action prioritized)...")
+    messages = state["messages"]
+    last_message = messages[-1]
+    new_state = state.copy()
+
+    # Only process AI messages (not system/user)
+    if not isinstance(last_message, AIMessage):
+        return new_state
+
+    content = last_message.content
+
+    # Remove any system prompt/instructions (if present in content)
+    # Assume that the actual AI output is after the last occurrence of "You are a general AI assistant" or similar system prompt marker
+    sys_prompt_pattern = r"(You are a general AI assistant.*?)(?=\n\n|$)"
+    content_wo_sys_prompt = re.sub(sys_prompt_pattern, '', content, flags=re.DOTALL | re.IGNORECASE).strip()
+
+    # Find the last occurrence of FINAL ANSWER or Action Input
+    final_answer_match = list(re.finditer(r"FINAL ANSWER:", content_wo_sys_prompt, re.IGNORECASE))
+    action_input_match = list(re.finditer(r"Action Input:", content_wo_sys_prompt, re.IGNORECASE))
+
+    # Helper: get the last match position and which it was
+    last_marker = None
+    last_pos = -1
+    if final_answer_match:
+        last_fa = final_answer_match[-1]
+        last_marker = 'FINAL ANSWER'
+        last_pos = last_fa.start()
+    if action_input_match:
+        last_ai = action_input_match[-1]
+        if last_ai.start() > last_pos:
+            last_marker = 'Action Input'
+            last_pos = last_ai.start()
+
+    # If neither marker found, mark as done
+    if not last_marker:
+        print("No FINAL ANSWER or Action Input found in last AI output.")
+        new_state["done"] = True
+        return new_state
+
+    # Get the substring from the last marker to the end
+    last_section = content_wo_sys_prompt[last_pos:].strip()
+
+    # 2. If FINAL ANSWER is in the last part, end the process
+    if last_marker == 'FINAL ANSWER':
+        # Extract the answer after FINAL ANSWER:
+        answer = re.search(r"FINAL ANSWER:\s*(.+)", last_section, re.IGNORECASE)
+        final_answer_text = answer.group(1).strip() if answer else ""
+        updated_ai_message = AIMessage(content=f"FINAL ANSWER: {final_answer_text}", tool_calls=[])
+        new_state["messages"] = messages[:-1] + [updated_ai_message]
+        new_state["done"] = True
+        print(f"FINAL ANSWER found at end: '{final_answer_text}'")
+        return new_state
+
+    # 3. If Action Input is in the last part, launch tool
+    if last_marker == 'Action Input':
+        # Try to extract the Action and Action Input for the last occurrence
+        action_match = list(re.finditer(r"Action:\s*([^\n]+)", last_section))
+        action_input_match = list(re.finditer(r"Action Input:\s*([^\n]+)", last_section))
+        if action_match and action_input_match:
+            tool_name = action_match[-1].group(1).strip()
+            tool_input_raw = action_input_match[-1].group(1).strip()
+            print(f"ReAct: Found Action: {tool_name}, Input: '{tool_input_raw}'")
+            # Format tool_args as in your original code (simplified here)
+            tool_args = {"query": tool_input_raw}
+            tool_call_id = str(uuid.uuid4())
+            parsed_tool_calls = [{"name": tool_name, "args": tool_args, "id": tool_call_id}]
+            updated_ai_message = AIMessage(content=content, tool_calls=parsed_tool_calls)
+            new_state["messages"] = messages[:-1] + [updated_ai_message]
+            new_state.pop("done", None)
+            print(f"AIMessage updated with tool_calls: {parsed_tool_calls}")
+            return new_state
+        else:
+            print("Action Input found at end, but could not parse Action or Action Input.")
+            new_state["done"] = True
+            return new_state
+
+    # Fallback: mark as done
+    print("No actionable marker found at end of last AI output. Marking as done.")
+    new_state["done"] = True
+    return new_state
 
-    if not answers:
-        return {
-            "most_common_answer": "Could not analyze any frames successfully.",
-            "all_answers": [],
-            "answer_counts": Counter()
-        }
 
-    counted = Counter(answers)
-    most_common_answer, freq = counted.most_common(1)[0]
-    return {
-        "most_common_answer": most_common_answer,
-        "all_answers": answers,
-        "answer_counts": counted
-    }
 
 def download_youtube_video(url, output_dir='/tmp/video/', output_filename='downloaded_video.mp4'):
     """Download a YouTube video using yt-dlp"""
@@ -307,419 +1019,660 @@ def extract_frames(video_path, output_dir, frame_interval_seconds=10):
         print(f"Exception during frame extraction: {e}")
         return False
 
-def image_qa(image_path: str, question: str) -> str:
-    """Answer questions about an image using the BLIP model"""
+def answer_question_on_frame(image_path, question):
+    """Answer a question about a single video frame using BLIP"""
     try:
+        vqa_model_name = "Salesforce/blip-vqa-base" # Not used in the provided graph logic directly
+        processor_vqa = BlipProcessor.from_pretrained(vqa_model_name) # Not used
+        model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to('cpu') # Not used
+        device = "cpu"
+
         image = Image.open(image_path).convert('RGB')
         inputs = processor_vqa(image, question, return_tensors="pt").to(device)
         out = model_vqa.generate(**inputs)
         answer = processor_vqa.decode(out[0], skip_special_tokens=True)
         return answer
     except Exception as e:
-        print(f"Error in image_qa: {str(e)}")
-        return f"Error processing image: {str(e)}"
-
-# --- Node functions ---
-def router(state: Dict[str, Any]) -> str:
-    """Determine the next node based on question content and file type"""
-    question = state.get('question', '')
-
-    # Pattern for Wikipedia and similar sources
-    wiki_pattern = r"(wikipedia\.org|wiki|encyclopedia|britannica\.com|encyclop[a|æ]dia)"
-    has_wiki = re.search(wiki_pattern, question, re.IGNORECASE) is not None
-
-    # Pattern for YouTube
-    yt_pattern = r"(https?://)?(www\.)?(youtube\.com|youtu\.be)/[^\s]+"
-    has_youtube = re.search(yt_pattern, question) is not None
-
-    # Check for image
-    has_image = state.get('file_type') == 'picture'
+        print(f"Error processing frame {image_path}: {str(e)}")
+        return "Error processing this frame"
 
-    # Check for audio
-    has_audio = state.get('file_type') == 'audio'
+def answer_video_question(frames_dir, question):
+    """Answer a question about a video by analyzing extracted frames"""
+    valid_exts = ('.jpg', '.jpeg', '.png')
 
-    print(f"Has Wikipedia reference: {has_wiki}")
-    print(f"Has YouTube link: {has_youtube}")
-    print(f"Has picture file: {has_image}")
-    print(f"Has audio file: {has_audio}")
+    # Check if directory exists
+    if not os.path.exists(frames_dir):
+        return {
+            "most_common_answer": "No frames found to analyze.",
+            "all_answers": [],
+            "answer_counts": Counter()
+        }
 
-    if has_wiki:
-        return "retrieve"
-    elif has_youtube:
-        # Store the extracted YouTube URL in the state
-        url_match = re.search(r"(https?://[^\s]+)", question)
-        if url_match:
-            state['youtube_url'] = url_match.group(0)
-        return "video"
-    elif has_image:
-        return "image"
-    elif has_audio:
-        return "audio"
-    else:
-        return "llm"
-
-def node_decide(state: Dict[str, Any]) -> Dict[str, Any]:
-    """Router node that decides which node to go to next"""
-    print("Running node_decide")
-    # Initialize context list if not present
-    if 'context' not in state:
-        state['context'] = []
-    # Add the next state to the state dict
-    state["next"] = router(state)
-    print(f"Routing to: {state['next']}")
-    return state
+    frame_files = [os.path.join(frames_dir, f) for f in os.listdir(frames_dir)
+                  if f.lower().endswith(valid_exts)]
 
-def node_image(state: Dict[str, Any]) -> Dict[str, Any]:
-    """Process image-based questions"""
-    print("Running node_image")
-    try:
-        # Make sure the image file exists
-        if not os.path.exists(state['file_path']):
-            state['answer'] = ensure_final_answer_format("Image file not found.")
-            return state
-            
-        # Get answer from image QA model
-        answer = image_qa(state['file_path'], state['question'])
-        
-        # Format the final answer
-        state['answer'] = ensure_final_answer_format(answer)
-        
-        # Add document to state for traceability
-        image_doc = Document(
-            page_content=f"Image analysis result: {answer}",
-            metadata={"source": "image_analysis", "file_path": state['file_path']}
-        )
-        state['context'].append(image_doc)
-        
-    except Exception as e:
-        error_msg = f"Error processing image: {str(e)}"
-        print(error_msg)
-        state['answer'] = ensure_final_answer_format(error_msg)
-    
-    return state
+    # Sort frames properly by number
+    def get_frame_number(filename):
+        match = re.search(r'(\d+)', os.path.basename(filename))
+        return int(match.group(1)) if match else 0
 
-def node_video(state: Dict[str, Any]) -> Dict[str, Any]:
-    """Process video-based questions"""
-    print("Running node_video")
-    youtube_url = state.get('youtube_url')
-    if not youtube_url:
-        state['answer'] = ensure_final_answer_format("No YouTube URL found in the question.")
-        return state
+    frame_files = sorted(frame_files, key=get_frame_number)
 
-    question = state['question']
-    # Extract the actual question part (remove the URL)
-    question_text = re.sub(r'https?://[^\s]+', '', question).strip()
-    if not question_text.endswith('?'):
-        question_text += '?'
+    if not frame_files:
+        return {
+            "most_common_answer": "No valid image frames found.",
+            "all_answers": [],
+            "answer_counts": Counter()
+        }
 
-    video_file = download_youtube_video(youtube_url)
-    if not video_file or not os.path.exists(video_file):
-        state['answer'] = ensure_final_answer_format("Failed to download the video.")
-        return state
+    answers = []
+    for frame_path in frame_files:
+        try:
+            ans = answer_question_on_frame(frame_path, question)
+            answers.append(ans)
+            print(f"Processed frame: {os.path.basename(frame_path)}, Answer: {ans}")
+        except Exception as e:
+            print(f"Error processing frame {frame_path}: {str(e)}")
 
-    frames_dir = "/tmp/frames"
-    os.makedirs(frames_dir, exist_ok=True)
+    if not answers:
+        return {
+            "most_common_answer": "Could not analyze any frames successfully.",
+            "all_answers": [],
+            "answer_counts": Counter()
+        }
 
-    success = extract_frames(video_path=video_file, output_dir=frames_dir, frame_interval_seconds=10)
-    if not success:
-        state['answer'] = ensure_final_answer_format("Failed to extract frames from the video.")
-        return state
+    counted = Counter(answers)
+    most_common_answer, freq = counted.most_common(1)[0]
+    return {
+        "most_common_answer": most_common_answer,
+        "all_answers": answers,
+        "answer_counts": counted
+    }
 
-    result = answer_video_question(frames_dir, question_text)
-    final_answer = result['most_common_answer']
-    state['frame_answers'] = result['all_answers']
 
-    # Create Document objects for each frame analysis
-    frame_documents = []
-    for i, ans in enumerate(result['all_answers']):
-        doc = Document(
-            page_content=f"Frame {i}: {ans}",
-            metadata={"frame_number": i, "source": "video_analysis"}
+class YoutubeScreenshotQA(BaseTool):
+    name: str = "youtube_screenshot_qa"
+    description: str = (
+        "Downloads a YouTube video, extracts screenshots at intervals, "
+        "and answers a question about the video based on the screenshots. "
+        "Input should be a dict with keys: 'youtube_url' and 'question'."
+        "Example input: {'youtube_url': 'https://www.youtube.com/watch?v=L1vXCYZAYYM', 'question': 'What is the highest number of bird species on camera simultaneously?'}"
+    )
+    frame_interval_seconds: int = 10  # Can be parameterized if needed
+
+    def _run(self, input_data: Dict[str, Any]) -> str:
+        youtube_url = input_data.get("youtube_url")
+        question = input_data.get("question")
+
+        if not youtube_url or not question:
+            return "Error: Input must include 'youtube_url' and 'question'."
+
+        # Step 1: Download the video
+        video_dir = '/tmp/video/'
+        video_filename = 'downloaded_video.mp4'
+        print(f"Downloading YouTube video from {youtube_url}...")
+        video_path = download_youtube_video(youtube_url, output_dir=video_dir, output_filename=video_filename)
+        if not video_path or not os.path.exists(video_path):
+            return "Error: Failed to download the YouTube video."
+
+        # Step 2: Extract frames
+        frames_dir = '/tmp/video_frames/'
+        print(f"Extracting frames from {video_path} every {self.frame_interval_seconds} seconds...")
+        success = extract_frames(video_path, frames_dir, frame_interval_seconds=self.frame_interval_seconds)
+        if not success:
+            return "Error: Failed to extract frames from the video."
+
+        # Step 3: Analyze frames and answer question
+        print(f"Answering question about the video frames...")
+        answer_result = answer_video_question(frames_dir, question)
+        if not answer_result or not answer_result.get("most_common_answer"):
+            return "Error: Could not analyze video frames to answer the question."
+
+        # Format the result
+        most_common = answer_result["most_common_answer"]
+        all_answers = answer_result["all_answers"]
+        counts = answer_result["answer_counts"]
+
+        result = (
+            f"Most common answer: {most_common}\n"
+            f"All answers: {all_answers}\n"
+            f"Answer counts: {dict(counts)}"
         )
-        frame_documents.append(doc)
-
-    # Add documents to state
-    state['context'].extend(frame_documents)
-    state['answer'] = ensure_final_answer_format(final_answer)
-
-    print(f"Video answer: {state['answer']}")
-    return state
-
-def node_audio_rag(state: Dict[str, Any]) -> Dict[str, Any]:
-    """Process audio-based questions"""
-    print(f"Processing audio file: {state['file_path']}")
+        return result
 
-    try:
-        # Step 1: Transcribe audio
-        audio, sr = librosa.load(state['file_path'], sr=16000)
-        asr_result = asr_pipe({"raw": audio, "sampling_rate": sr})
-        audio_transcript = asr_result['text']
-        print(f"Audio transcript: {audio_transcript}")
-
-        # Step 2: Store transcript in vector store
-        transcript_doc = [Document(page_content=audio_transcript)]
-        embeddings = HuggingFaceEmbeddings(model_name='BAAI/bge-large-en-v1.5')
-        vector_db = FAISS.from_documents(transcript_doc, embedding=embeddings)
-
-        # Step 3: Retrieve relevant docs for the user's question
-        question = state['question']
-        similar_docs = vector_db.similarity_search(question, k=1)
-        retrieved_context = "\n".join([doc.page_content for doc in similar_docs])
-
-        # Step 4: Generate answer
-        prompt = (
-            f"You are an AI assistant that answers questions about audio content.\n\n"
-            f"Audio transcript: {retrieved_context}\n\n"
-            f"Question: {question}\n\n"
-            f"Based only on the provided audio transcript, answer the question. "
-            f"If the transcript does not contain relevant information, state that clearly.\n\n"
-            f"End your response with 'FINAL ANSWER: ' followed by a concise answer."
-        )
-        
-        llm_response = llm_pipe(prompt)
-        answer_text = llm_response[0]['generated_text']
-        
-        # Add documents to state
-        state['context'].extend(transcript_doc)
-        state['context'].append(Document(
-            page_content=prompt,
-            metadata={"source": "audio_analysis_prompt"}
-        ))
+def tools_condition_with_logging(state: AgentState):
+    """
+    Custom tools condition function that checks if the last message contains tool calls
+    in the Thought/Action/Action Input format and logs the transition decision.
+    
+    Args:
+        state (AgentState): The current state containing messages
         
-        # Ensure final answer format
-        state['answer'] = ensure_final_answer_format(answer_text)
-
-    except Exception as e:
-        error_msg = f"Audio processing error: {str(e)}"
-        print(error_msg)
-        state['answer'] = ensure_final_answer_format(error_msg)
-
-    return state
+    Returns:
+        str: "tools" if tool calls are present, "__end__" otherwise
+    """
+    
+    import re
 
-def node_llm(state: Dict[str, Any]) -> Dict[str, Any]:
-    """Process general knowledge questions with LLM"""
-    print("Running node_llm")
-    question = state['question']
-
-    # Compose a detailed prompt
-    prompt = (
-        "You are an AI assistant that answers questions using your general knowledge. "
-        "Follow these steps:\n\n"
-        "1. If the question appears to be scrambled or jumbled, first try to unscramble or reconstruct the intended meaning.\n"
-        "2. Analyze the question (unscrambled if needed) and use your own knowledge to answer it.\n"
-        "3. If the question can't be answered with certainty, provide your best estimate and clearly explain any assumptions.\n"
-        "4. Format your answer using these rules:\n"
-        "   - Numbers: Plain digits without commas/units (e.g. 1234567)\n"
-        "   - Strings: Minimal words, no articles/abbreviations\n"
-        "   - Lists: comma-separated values without extra formatting\n\n"
-        "5. Always conclude with:\n"
-        "FINAL ANSWER: [your answer] (replace bracketed text)\n\n"
-        f"Current question: {question}"
+    # Ensure we have messages in the state
+    if not state.get("messages") or len(state["messages"]) == 0:
+        print("❌ No messages found in state, ending conversation")
+        return "__end__"
+    
+    # Get the last message
+    last_message = state["messages"][-1]
+    
+    # Get message content
+    content = ""
+    if hasattr(last_message, 'content'):
+        content = str(last_message.content)
+    elif isinstance(last_message, dict) and 'content' in last_message:
+        content = str(last_message['content'])
+    else:
+        print("❌ No content found in last message, ending conversation")
+        return "__end__"
+    
+    print(f"🔍 Analyzing message content: {content[:200]}...")
+    
+    # Check for Thought/Action/Action Input format
+    has_tool_calls = False
+    
+    # Pattern to match the format:
+    # Thought: <thought>
+    # Action: <tool_name>
+    # Action Input: <input>
+    thought_action_pattern = re.compile(
+        r'Thought:\s*(.*?)\n\s*Action:\s*(.*?)\n\s*Action Input:\s*(.*?)(?:\n|$)',
+        re.DOTALL | re.IGNORECASE
     )
-
-    # Add document to state for traceability
-    query_doc = Document(
-        page_content=prompt,
-        metadata={"source": "llm_prompt"}
+    
+    # Also check for just Action/Action Input without Thought
+    action_only_pattern = re.compile(
+        r'Action:\s*(.*?)\n\s*Action Input:\s*(.*?)(?:\n|$)',
+        re.DOTALL | re.IGNORECASE
     )
-    state['context'].append(query_doc)
+    
+    # Look for the complete format first
+    match = thought_action_pattern.search(content)
+    if not match:
+        # Try the action-only format
+        match = action_only_pattern.search(content)
+        if match:
+            thought = "No thought provided"
+            action = match.group(1).strip()
+            action_input = match.group(2).strip()
+        else:
+            action = None
+            action_input = None
+            thought = None
+    else:
+        thought = match.group(1).strip()
+        action = match.group(2).strip()
+        action_input = match.group(3).strip()
+    
+    if match and action:
+        has_tool_calls = True
+        print(f"🔧 Found tool call format:")
+        print(f"   Thought: {thought}")
+        print(f"   Action: {action}")
+        print(f"   Action Input: {action_input}")
+        
+        # Map common tool names to your actual tools
+        tool_mappings = {
+            'wikipedia_semantic_search': 'wikipedia_tool',
+            'wikipedia': 'wikipedia_tool',
+            'search': 'search_tool',
+            'duckduckgo_search': 'search_tool',
+            'web_search': 'search_tool',
+            'youtube_screenshot_qa_tool': 'youtube_tool',
+            'youtube': 'youtube_tool',
+        }
+        
+        # Normalize the action name
+        normalized_action = action.lower().strip()
+        
+        # Store the parsed tool call information in the state for the tools node to use
+        if 'parsed_tool_calls' not in state:
+            state['parsed_tool_calls'] = []
+        
+        tool_call_info = {
+            'thought': thought,
+            'action': action,
+            'action_input': action_input,
+            'normalized_action': normalized_action,
+            'tool_mapping': tool_mappings.get(normalized_action, normalized_action)
+        }
+        
+        state['parsed_tool_calls'].append(tool_call_info)
+        print(f"🚀 Added tool call to state: {tool_call_info}")
+        
+        # Don't execute tools here - let call_tool handle execution
+        # Just store the parsed information for call_tool to use
+    
+    # Also check for standalone tool mentions (fallback)
+    if not has_tool_calls:
+        # Check for tool names mentioned in content
+        tool_keywords = [
+            'wikipedia_semantic_search', 'wikipedia', 'search', 'duckduckgo', 
+            'youtube_screenshot_qa_tool', 'youtube', 'web search'
+        ]
+        
+        content_lower = content.lower()
+        for keyword in tool_keywords:
+            if keyword in content_lower:
+                print(f"🔧 Found tool keyword '{keyword}' in content (fallback detection)")
+                has_tool_calls = True
+                break
+    
+    if has_tool_calls:
+        print("🔧 Tool calls detected, transitioning to tools...")
+        return "tools"
+    else:
+        print("✅ No tool calls found, ending conversation")
+        return "__end__"
 
-    try:
-        result = llm_pipe(prompt)
-        answer_text = result[0]['generated_text']
-        state['answer'] = ensure_final_answer_format(answer_text)
-    except Exception as e:
-        print(f"Error in LLM processing: {str(e)}")
-        error_msg = f"An error occurred while processing your question: {str(e)}"
-        state['answer'] = ensure_final_answer_format(error_msg)
 
-    print(f"LLM answer: {state['answer']}")
-    return state
-def retrieve(state: State) -> State:
-    """Retrieve relevant documents using RAG"""
-    print("Running retrieve")
-    question = state["question"]
+# 2. Improved call_tool with memory management
+def call_tool_with_memory_management(state: AgentState) -> AgentState:
+    """Process tool calls with memory management."""
+    print("Running call_tool with memory management...")
     
+    # Clear CUDA cache before processing
     try:
-        # Tokenize the question
-        inputs = tokenizer(question, return_tensors="pt")
-        
-        # Get doc_ids by using the retriever directly
-        question_hidden_states = rag_model.question_encoder(inputs["input_ids"])[0]
-        docs_dict = retriever(
-            inputs["input_ids"].numpy(),
-            question_hidden_states.detach().numpy(), 
-            return_tensors="pt"
-        )
-        
-        # Extract the retrieved passages
-        all_chunks = []
+        import torch
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            print(f"🧹 Cleared CUDA cache. Memory: {torch.cuda.memory_allocated()/1024**2:.1f}MB")
+    except:
+        pass
+    
+    # Check if we have parsed tool calls from the condition function
+    if 'parsed_tool_calls' in state and state['parsed_tool_calls']:
+        return execute_parsed_tool_calls(state)
+    
+    # Fallback to original OpenAI-style tool calls handling
+    messages = state["messages"]
+    last_message = messages[-1]
+    
+    if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
+        print("No tool calls found in last message")
+        return state
+    
+    # Copy the messages to avoid mutating the original list
+    new_messages = list(messages)
+    
+    print(f"Processing {len(last_message.tool_calls)} tool calls")
+    
+    for i, tool_call in enumerate(last_message.tool_calls):
+        print(f"Processing tool call {i+1}: {tool_call['name'] if isinstance(tool_call, dict) else tool_call.name}")
         
-        # Debug print to see what's in docs_dict
-        print(f"docs_dict keys: {docs_dict.keys()}")
+        # Handle both dict and object-style tool calls
+        if isinstance(tool_call, dict):
+            tool_name = tool_call.get("name", "")
+            args = tool_call.get("args", {})
+            tool_call_id = tool_call.get("id", str(uuid.uuid4()))
+        else:
+            tool_name = getattr(tool_call, "name", "")
+            args = getattr(tool_call, "args", {})
+            tool_call_id = getattr(tool_call, "id", str(uuid.uuid4()))
         
-        # Check for different possible keys that might contain the documents
-        doc_text_key = None
-        for possible_key in ['retrieved_doc_text', 'doc_text', 'texts', 'documents']:
-            if possible_key in docs_dict:
-                doc_text_key = possible_key
+        # Find the matching tool
+        selected_tool = None
+        for tool in tools:
+            if tool.name.lower() == tool_name.lower():
+                selected_tool = tool
                 break
         
-        if doc_text_key:
-            # Access the retrieved document texts from the docs_dict
-            for i in range(len(docs_dict["doc_ids"][0])):
-                doc_text = docs_dict[doc_text_key][0][i]
-                all_chunks.append(Document(page_content=doc_text))
-                
-            print(f"Retrieved {len(all_chunks)} documents")
+        if not selected_tool:
+            tool_result = f"Error: Tool '{tool_name}' not found. Available tools: {', '.join(t.name for t in tools)}"
+            print(f"Tool not found: {tool_name}")
         else:
-            # Fallback: Try to extract document text from doc_ids
-            doc_ids = docs_dict.get("doc_ids", [[]])[0]
-            print(f"Retrieved doc_ids: {doc_ids}")
-            
-            # Create minimal document stubs from IDs
-            for doc_id in doc_ids:
-                stub_text = f"Information related to document ID: {doc_id}"
-                all_chunks.append(Document(page_content=stub_text))
-            
-            print(f"Created {len(all_chunks)} document stubs from IDs")
-        
-        # Add documents to state context
-        if not state.get('context'):
-            state['context'] = []
-        state['context'].extend(all_chunks)
+            try:
+                # Extract query
+                if isinstance(args, dict) and "query" in args:
+                    query = args["query"]
+                else:
+                    query = str(args) if args else ""
+                
+                print(f"Executing {tool_name} with query: {query[:100]}...")
+                tool_result = selected_tool.run(query)
+                
+                # Aggressive truncation to prevent memory issues
+                max_length = 3000 if "wikipedia" in tool_name.lower() else 2000
+                if len(tool_result) > max_length:
+                    tool_result = tool_result[:max_length] + f"... [Result truncated from {len(tool_result)} to {max_length} chars to prevent memory issues]"
+                    print(f"📄 Truncated result to {max_length} characters")
+                
+                print(f"Tool result length: {len(tool_result)} characters")
+                
+            except Exception as e:
+                tool_result = f"Error executing tool '{tool_name}': {str(e)}"
+                print(f"Tool execution error: {e}")
         
-    except Exception as e:
-        print(f"Error in retrieval: {str(e)}")
-        # Create an error document
-        error_doc = Document(
-            page_content=f"Error during retrieval: {str(e)}",
-            metadata={"source": "retrieval_error"}
+        # Create tool message
+        tool_message = ToolMessage(
+            content=tool_result,
+            name=tool_name,
+            tool_call_id=tool_call_id
         )
-        if not state.get('context'):
-            state['context'] = []
-        state['context'].append(error_doc)
+        new_messages.append(tool_message)
+        print(f"Added tool message for {tool_name}")
     
-    return state
-
-def generate(state: State) -> State:
-    """Generate an answer based on retrieved documents"""
-    print("Running generate")
+    # Update the state
+    new_state = state.copy()
+    new_state["messages"] = new_messages
     
+    # Clear CUDA cache after processing
     try:
-        # Check if context exists
-        if not state.get('context') or len(state['context']) == 0:
-            state['answer'] = ensure_final_answer_format("No relevant information found to answer your question.")
-            return state
-            
-        # Concatenate all context documents into a single string
-        docs_content = "\n\n".join(doc.page_content for doc in state["context"])
-        
-        # Format the prompt for the LLM
-        prompt_str = PromptTemplate(
-            input_variables=["question", "context"],
-            template=(
-                "You are an AI assistant that answers questions using retrieved context. "
-                "Follow these steps:\n\n"
-                "1. Analyze the provided context:\n{context}\n\n"
-                "2. If the context contains scrambled text, first attempt to reconstruct meaningful information\n"
-                "3. If the question can't be answered from context alone, combine context with general knowledge "
-                "but clearly state this limitation\n"
-                "4. Format your answer using these rules:\n"
-                "   - Numbers: Plain digits without commas/units (e.g. 1234567)\n"
-                "   - Strings: Minimal words, no articles/abbreviations\n"
-                "   - Lists: comma-separated values without extra formatting\n\n"
-                "5. Always conclude with:\n"
-                "FINAL ANSWER: [your answer] (replace bracketed text)\n\n"
-                "Current question: {question}"
-            )
-        ).format(question=state["question"], context=docs_content)
+        import torch
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except:
+        pass
+    
+    return new_state
+
+
+def execute_parsed_tool_calls(state: AgentState):
+    """
+    Execute tool calls that were parsed from the Thought/Action/Action Input format.
+    This is called by call_tool when parsed_tool_calls are present in state.
+    
+    Args:
+        state (AgentState): The current state containing parsed tool calls
         
-        # Generate answer using the LLM pipeline
-        response = llm_pipe(prompt_str)
-        answer_text = response[0]["generated_text"]
+    Returns:
+        AgentState: Updated state with tool results
+    """
+    
+    # Use the same tools list that's available globally
+    # Map tool names to the actual tool instances
+    tool_name_mappings = {
+        'wikipedia_semantic_search': 'wikipedia_tool',
+        'wikipedia': 'wikipedia_tool', 
+        'search': 'enhanced_search',  # Updated mapping
+        'duckduckgo_search': 'enhanced_search',  # Updated mapping  
+        'web_search': 'enhanced_search',  # Updated mapping
+        'enhanced_search': 'enhanced_search',  # Direct mapping
+        'youtube_screenshot_qa_tool': 'youtube_tool',
+        'youtube': 'youtube_tool',
+    }
+    
+    
+    # Create a lookup by tool names for your existing tools list
+    tools_by_name = {}
+    for tool in tools:
+        tools_by_name[tool.name.lower()] = tool
+        # Also map by class name for flexibility
+        class_name = tool.__class__.__name__.lower()
+        if 'wikipedia' in class_name:
+            tools_by_name['wikipedia_tool'] = tool
+        elif 'search' in class_name or 'duck' in class_name:
+            tools_by_name['search_tool'] = tool
+        elif 'youtube' in class_name:
+            tools_by_name['youtube_tool'] = tool
+    
+    # Copy messages to avoid mutation during iteration
+    new_messages = list(state["messages"])
+    
+    for tool_call in state['parsed_tool_calls']:
+        action = tool_call['action']
+        action_input = tool_call['action_input']
+        thought = tool_call['thought']
+        normalized_action = tool_call['normalized_action']
         
-        # Ensure answer has the FINAL ANSWER format
-        state['answer'] = ensure_final_answer_format(answer_text)
+        print(f"🚀 Executing tool: {action} with input: {action_input}")
         
-    except Exception as e:
-        print(f"Error in generate node: {str(e)}")
-        error_msg = f"Error generating answer: {str(e)}"
-        state['answer'] = ensure_final_answer_format(error_msg)
+        # Find the tool instance
+        tool_instance = None
         
+        # Try direct name match first
+        if normalized_action in tools_by_name:
+            tool_instance = tools_by_name[normalized_action]
+        # Try mapped name
+        elif normalized_action in tool_name_mappings:
+            mapped_name = tool_name_mappings[normalized_action]
+            if mapped_name in tools_by_name:
+                tool_instance = tools_by_name[mapped_name]
+        
+        if tool_instance:
+            try:
+                result = tool_instance.run(action_input)
+                if len(result) > 6000:
+                    result = result[:6000] + "... [Result truncated due to length]"
+                
+                # Create observation message in the format your agent expects
+                from langchain_core.messages import AIMessage
+                observation = f"Observation: {result}"
+                observation_message = AIMessage(content=observation)
+                new_messages.append(observation_message)
+                
+                print(f"✅ Tool '{action}' executed successfully")
+                
+            except Exception as e:
+                print(f"❌ Error executing tool '{action}': {e}")
+                from langchain_core.messages import AIMessage
+                error_msg = f"Observation: Error executing '{action}': {str(e)}"
+                error_message = AIMessage(content=error_msg)
+                new_messages.append(error_message)
+        else:
+            print(f"❌ Tool '{action}' not found in available tools")
+            available_tool_names = list(tools_by_name.keys())
+            from langchain_core.messages import AIMessage
+            error_msg = f"Observation: Tool '{action}' not found. Available tools: {', '.join(available_tool_names)}"
+            error_message = AIMessage(content=error_msg)
+            new_messages.append(error_message)
+    
+    # Update state with new messages and clear parsed tool calls
+    state["messages"] = new_messages
+    state['parsed_tool_calls'] = []
+    
     return state
 
-# --- Define the edge condition function ---
-def get_next_node(state: Dict[str, Any]) -> str:
-    """Get the next node from the state"""
-    return state["next"]
-
-# Create the StateGraph
-graph = StateGraph(State)
-
-# Add nodes
-graph.add_node("decide", node_decide)
-graph.add_node("video", node_video)
-graph.add_node("llm", node_llm)
-graph.add_node("retrieve", retrieve)
-graph.add_node("generate", generate)
-graph.add_node("image", node_image)
-graph.add_node("audio", node_audio_rag)
-
-# Add edge from START to decide
-graph.add_edge(START, "decide")
-graph.add_edge("retrieve", "generate")
-
-# Add conditional edges from decide to other nodes based on question
-graph.add_conditional_edges(
-    "decide",
-    get_next_node,
-    {
-        "video": "video",
-        "llm": "llm",
-        "retrieve": "retrieve",
-        "image": "image",
-        "audio": "audio"
-    }
-)
+# 1. Add loop detection to your AgentState
+def should_continue(state: AgentState) -> str:
+    """Determine if the agent should continue or end."""
+    print("Running should_continue....")
+    messages = state["messages"]
+    
+    #ipdb.set_trace()
 
-# Add edges from all terminal nodes to END
-graph.add_edge("video", END)
-graph.add_edge("llm", END)
-graph.add_edge("generate", END)
-graph.add_edge("image", END)
-graph.add_edge("audio", END)
+    # Check if we're done
+    if state.get("done", False):
+        return "end"
+    
+    # Prevent infinite loops - limit tool calls
+    tool_call_count = sum(1 for msg in messages if hasattr(msg, 'tool_calls') and msg.tool_calls)
+    if tool_call_count >= 3:  # Max 3 tool calls per conversation
+        print(f"⚠️ Stopping: Too many tool calls ({tool_call_count})")
+        return "end"
+    
+    # Check for repeated tool calls with same query
+    recent_tool_calls = []
+    for msg in messages[-6:]:  # Check last 6 messages
+        if hasattr(msg, 'tool_calls') and msg.tool_calls:
+            for tool_call in msg.tool_calls:
+                if isinstance(tool_call, dict):
+                    recent_tool_calls.append((tool_call.get('name'), str(tool_call.get('args', {}))))
+    
+    if len(recent_tool_calls) >= 2 and recent_tool_calls[-1] == recent_tool_calls[-2]:
+        print("⚠️ Stopping: Repeated tool call detected")
+        return "end"
+    
+    # Check message count to prevent runaway conversations
+    if len(messages) > 15:
+        print(f"⚠️ Stopping: Too many messages ({len(messages)})")
+        return "end"
+    
+    return "continue"
+
+def route_after_parse_react(state: AgentState) -> str:
+    """Determines the next step after parsing LLM output, prioritizing end state."""
+    if state.get("done", False):  # Check if parse_react_output decided we are done
+        return "end_processing"
+
+    # Original logic: check for tool calls in the last message
+    # Ensure messages list and last message exist before checking tool_calls
+    messages = state.get("messages", [])
+    if messages:
+        last_message = messages[-1]
+        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
+            return "call_tool"
+    return "call_llm"
+
+#wikipedia_tool = WikipediaSearchToolWithFAISS()
+#search_tool = DuckDuckGoSearchRun()
+#youtube_screenshot_qa_tool = YoutubeScreenshotQA()
+
+# Combine all tools
+#tools = [wikipedia_tool, search_tool, youtube_screenshot_qa_tool]
+
+# Update your tools list to use the global instances
+#
+
+# --- Graph Construction ---
+# --- Graph Construction ---
+def create_memory_safe_workflow():
+    """Create a workflow with memory management and loop prevention."""
+    # These models are initialized here but might be better managed if they need to be released/reinitialized
+    # like you attempt in run_agent. Consider passing them or managing their lifecycle carefully.
+    hf_pipe = create_llm_pipeline()
+    llm = HuggingFacePipeline(pipeline=hf_pipe)
+    # vqa_model_name = "Salesforce/blip-vqa-base" # Not used in the provided graph logic directly
+    # processor_vqa = BlipProcessor.from_pretrained(vqa_model_name) # Not used
+    # model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to('cpu') # Not used
+
+    workflow = StateGraph(AgentState)
+
+    # Bind the llm_model to the call_llm_with_memory_management function
+    bound_call_llm = partial(call_llm_with_memory_management, llm_model=llm)
+
+    # Add nodes with memory-safe versions
+    workflow.add_node("call_llm", bound_call_llm)  # Use the bound version here
+    workflow.add_node("parse_react_output", parse_react_output)
+    workflow.add_node("call_tool", call_tool_with_memory_management) # Ensure this doesn't also need llm if it calls back directly
+
+    # Set entry point
+    workflow.set_entry_point("call_llm")
+
+    # Add conditional edges
+    workflow.add_conditional_edges(
+        "call_llm",
+        should_continue,
+        {
+            "continue": "parse_react_output",
+            "end": END
+        }
+    )
 
-# Compile the graph
-agent = graph.compile()
+    workflow.add_conditional_edges(
+        "parse_react_output",
+        route_after_parse_react,
+        {
+            "call_tool": "call_tool",
+            "call_llm": "call_llm",
+            "end_processing": END
+        }
+    )
 
-# --- Intelligent Agent Function ---
-def intelligent_agent(state: State) -> str:
-    """Process a question using the appropriate pipeline based on content."""
-    try:
-        # Ensure state has proper structure
-        if not isinstance(state, dict):
-            return "FINAL ANSWER: Error - input must be a valid State dictionary"
-            
-        # Make sure question exists
-        if 'question' not in state:
-            return "FINAL ANSWER: Error - question is required"
-            
-        # Initialize context if not present
-        if 'context' not in state:
-            state['context'] = []
-            
-        print(f"Processing question: {state['question']}")
-        
-        # Invoke the agent with the state
-        final_state = agent.invoke(state)
-        
-        # Ensure answer has FINAL ANSWER format
-        answer = final_state.get('answer', "No answer found.")
-        formatted_answer = ensure_final_answer_format(answer)
-        
-        return formatted_answer
+    workflow.add_edge("call_tool", "call_llm")
+
+    return workflow.compile()
+
+# --- Run the Agent ---
+def run_agent(myagent, state: AgentState):
+    """
+    Initialize agent with proper system message and formatted query.
+    """
+    #global llm, hf_pipe, model_vqa, processor_vqa
+    global WIKIPEDIA_TOOL, SEARCH_TOOL, YOUTUBE_TOOL, tools
+
+    #ipdb.set_trace()
+
+    # At the module level, create instances once
+    WIKIPEDIA_TOOL = WikipediaSearchToolWithFAISS()
+    SEARCH_TOOL = EnhancedDuckDuckGoSearchTool(max_results=3, max_chars_per_page=3000)
+    YOUTUBE_TOOL = YoutubeScreenshotQA()
+
+    tools = [WIKIPEDIA_TOOL, SEARCH_TOOL, YOUTUBE_TOOL]
+
+    # Create a fresh system message each time
+    formatted_tools_description = render_text_description(tools)
+    current_date_str = datetime.now().strftime("%Y-%m-%d")
+    
+    system_content = f"""You are a general AI assistant. with access to these tools:
+
+    {formatted_tools_description}
+
+    If you need the most current information as of 2025, use enhanced_search
+    If you need to do in-depth research, use wikipedia_semantic_search_all_candidates_strong_entity_priority_list_retrieval
+    If you can answer the question confidently, do so directly.
+    If the question seems like gibberish (not English), try flipping the entire question and re-read the question.
+    If you need more information, use a tool.
+    (Think through the problem step by step)
+
+    When using a tool, follow this format:
+    Thought: <your thought>
+    Action: <tool_name>
+    Action Input: <tool_input>
+
+    Only use the tools listed above for the Action: step. Do not invent new tool names or actions. If you need to reason, do so in the Thought: step. After using a tool, process its output in your Thought: step, not as an Action.
+
+    Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
+    YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
+    If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
+    If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
+    If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string
+    Do not provide disclaimers. 
+    Do not provide supporting details.
+
+    """
+
+    # Get user question from AgentState
+    query = state['question']
+
+    # Pattern for YouTube
+    yt_pattern = r"(https?://)?(www\.)?(youtube\.com|youtu\.be)/[^\s]+"
+    has_youtube = re.search(yt_pattern, query) is not None
+
+    if has_youtube:
+        # Store the extracted YouTube URL in the state
+        url_match = re.search(r"(https?://[^\s]+)", query)
+        if url_match:
+            state['youtube_url'] = url_match.group(0)
         
-    except Exception as e:
-        print(f"Error in agent execution: {str(e)}")
-        return f"FINAL ANSWER: An error occurred - {str(e)}"
+    # Format the user query to guide the model better
+    formatted_query = f"""{query}"""
+    
+    # Initialize agent state with proper message types
+    system_message = SystemMessage(content=system_content)
+    human_message = HumanMessage(content=formatted_query)
+    
+    # Initialize state with properly typed messages and done=False
+    # state = {"messages": [system_message, human_message], "done": False}
+    state['messages'] = [system_message, human_message]
+    state["done"] = False
+
+    # Use the new method to run the graph
+    result = myagent.invoke(state)
+
+    # Check if FINAL ANSWER was given (i.e., workflow ended)
+    if result.get("done"):
+        #del llm
+        #del hf_pipe
+        #del model_vqa
+        #del processor_vqa
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+        gc.collect()
+        print("Released GPU memory after FINAL ANSWER.")
+        # Re-initialize for the next run
+        #hf_pipe = create_llm_pipeline()
+        #llm = HuggingFacePipeline(pipeline=hf_pipe)
+        #print("Re-initilized llm...")
+    
+    # Extract and return just the messages for cleaner output
+    return result["messages"]
+
+
+