Update final_agent.py

final_agent.py

@@ -1,10 +1,6 @@
 # -*- coding: utf-8 -*-
 """
 GAIA Benchmark Agent using LangChain, Groq, Tavily, and various tools.
-
-This agent is designed to interact with files, search the web, scrape pages,
-execute Python code, read Excel files, and transcribe audio/YouTube videos
-to tackle complex tasks like those found in the GAIA benchmark.
 """
 
 # --- Core Libraries ---
@@ -22,303 +18,199 @@ from dotenv import load_dotenv
 # --- LangChain Imports ---
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.tools import BaseTool, tool
-from langchain.pydantic_v1 import BaseModel, Field # Use Pydantic v1 for Langchain tool compatibility
+# Using Pydantic v2 is recommended if your environment supports it fully
+# from pydantic import BaseModel, Field # Pydantic v2
+from pydantic import BaseModel, Field # Pydantic v1 compatibility shim
 from langchain.memory import ConversationBufferWindowMemory
-from langchain.agents import AgentExecutor, create_structured_chat_agent
+from langchain.agents import AgentExecutor, create_openai_tools_agent # Keep OpenAI Tools Agent
 
 # --- Tool Specific Imports ---
 # Search
-from langchain_community.utilities import TavilySearchResults
+from langchain_community.tools.tavily_search import TavilySearchResults
 # Web Scraping
 import requests
 from bs4 import BeautifulSoup
 # LLM
 from langchain_groq import ChatGroq
 # Audio/Video Transcription (Optional)
-try:
-    import openai
-    OPENAI_AVAILABLE = True
-except ImportError:
-    OPENAI_AVAILABLE = False
+try: import openai; OPENAI_AVAILABLE = True
+except ImportError: OPENAI_AVAILABLE = False
 # Excel Reading (Optional)
-try:
-    import pandas as pd
-    PANDAS_AVAILABLE = True
-except ImportError:
-    PANDAS_AVAILABLE = False
+try: import pandas as pd; PANDAS_AVAILABLE = True
+except ImportError: PANDAS_AVAILABLE = False
 # YouTube Processing (Optional)
-try:
-    from pytube import YouTube
-    from pytube.exceptions import PytubeError
-    PYTUBE_AVAILABLE = True
-except ImportError:
-    PYTUBE_AVAILABLE = False
+try: from pytube import YouTube, PytubeError; PYTUBE_AVAILABLE = True
+except ImportError: PYTUBE_AVAILABLE = False
 
 # ==============================================================================
 # 1. CONFIGURATION
 # ==============================================================================
-load_dotenv() # Load environment variables from .env file if it exists
-
-# --- Agent Settings ---
-AGENT_WORKSPACE = Path("./gaia_agent_workspace")
-AGENT_WORKSPACE.mkdir(exist_ok=True) # Ensure workspace directory exists
-MAX_ITERATIONS = 15
-MEMORY_WINDOW_SIZE = 10
-
-# --- LLM Configuration ---
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-GROQ_MODEL_NAME = os.getenv("GROQ_MODEL_NAME", "llama3-70b-8192") # Default if not set
-
-# --- Tool Configuration ---
-TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
-TAVILY_MAX_RESULTS = 3
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # Needed for Whisper
-WHISPER_MODEL = "whisper-1"
-
-# --- Dependency & API Key Checks ---
-if not GROQ_API_KEY:
-    print("ERROR: GROQ_API_KEY environment variable not set. Agent cannot run.")
-    sys.exit(1)
-if not TAVILY_API_KEY:
-    print("ERROR: TAVILY_API_KEY environment variable not set. Search tool disabled.")
-    # Decide if this is fatal or just disables the tool
-    # sys.exit(1) # Uncomment to make it fatal
-
+load_dotenv()
+AGENT_WORKSPACE = Path("./gaia_agent_workspace"); AGENT_WORKSPACE.mkdir(exist_ok=True)
+MAX_ITERATIONS = 15; MEMORY_WINDOW_SIZE = 10
+GROQ_API_KEY = os.getenv("GROQ_API_KEY"); GROQ_MODEL_NAME = os.getenv("GROQ_MODEL_NAME", "llama3-70b-8192")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY"); TAVILY_MAX_RESULTS = 3
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY"); WHISPER_MODEL = "whisper-1"
+if not GROQ_API_KEY: print("ERROR: GROQ_API_KEY not set."); sys.exit(1)
+if not TAVILY_API_KEY: print("Warning: TAVILY_API_KEY not set.")
 openai_client = None
 if OPENAI_AVAILABLE and OPENAI_API_KEY:
-    try:
-        openai_client = openai.OpenAI(api_key=OPENAI_API_KEY)
-        print("OpenAI client initialized for Whisper transcription.")
-    except Exception as e:
-        print(f"Warning: Failed to initialize OpenAI client: {e}. Transcription tools disabled.")
-        openai_client = None
-elif OPENAI_AVAILABLE:
-    print("Warning: OpenAI library installed, but OPENAI_API_KEY not set. Transcription tools disabled.")
-else:
-    print("Info: OpenAI library not installed. Transcription tools disabled.")
-
-if not PANDAS_AVAILABLE:
-    print("Info: 'pandas' library not installed. Excel tool disabled. Install with: pip install pandas openpyxl")
-if not PYTUBE_AVAILABLE:
-    print("Info: 'pytube' library not installed. YouTube tool disabled. Install with: pip install pytube")
+    try: openai_client = openai.OpenAI(api_key=OPENAI_API_KEY); print("OpenAI client initialized.")
+    except Exception as e: print(f"Warning: OpenAI client init failed: {e}"); openai_client = None
+if not PANDAS_AVAILABLE: print("Info: 'pandas' not installed. Excel tool disabled.")
+if not PYTUBE_AVAILABLE: print("Info: 'pytube' not installed. YouTube tool disabled.")
 
 # ==============================================================================
 # 2. TOOL DEFINITIONS
 # ==============================================================================
 
 # --- Tool Input Schemas (Pydantic Models) ---
-# Using Pydantic v1 as required by Langchain tools at the time of writing
-
 class FileWriteArgs(BaseModel):
     relative_path: str = Field(description="Relative path within the agent's workspace where the file should be written.")
     content: str = Field(description="The text content to write into the file.")
-
 class FileReadArgs(BaseModel):
     relative_path: str = Field(description="Relative path within the agent's workspace of the file to read.")
-
 class ListDirectoryArgs(BaseModel):
     relative_path: str = Field(default=".", description="Relative path within the agent's workspace to list contents of. Use '.' for the root.")
-
 class RunPythonCodeArgs(BaseModel):
     code: str = Field(description="The Python code to execute. Use 'print()' to output results. Code runs in isolation.")
-
 class WebScrapeArgs(BaseModel):
     url: str = Field(description="The URL of the webpage to scrape.")
     query: Optional[str] = Field(default=None, description="Optional specific question to answer from the page content.")
-
 class ReadExcelArgs(BaseModel):
     relative_path: str = Field(description="Relative path within the agent's workspace of the Excel file (.xlsx or .xls).")
     sheet_name: Optional[str] = Field(default=None, description="Optional name of the specific sheet to read. Reads the first sheet if not specified.")
     max_rows_preview: int = Field(default=20, description="Maximum number of rows to include in the text preview.")
-
 class TranscribeAudioArgs(BaseModel):
     relative_path: str = Field(description="Relative path within the agent's workspace of the audio file (e.g., .mp3, .wav, .m4a). Max 25MB.")
-
 class TranscribeYouTubeArgs(BaseModel):
     youtube_url: str = Field(description="The URL of the YouTube video to transcribe. Audio will be downloaded temporarily.")
 
 # --- Helper Functions ---
-
 def _resolve_path(relative_path: str) -> Optional[Path]:
     """Resolves a relative path against the workspace and checks bounds."""
     try:
-        full_path = (AGENT_WORKSPACE / relative_path).resolve()
-        # Security Check: Ensure the resolved path is within the workspace
-        if not str(full_path).startswith(str(AGENT_WORKSPACE.resolve())):
-            return None # Path is outside the workspace
-        return full_path
-    except Exception: # Handle potential path resolution errors
+        normalized_relative_path = os.path.normpath(relative_path)
+        # Prevent absolute paths or paths trying to escape the workspace
+        if os.path.isabs(normalized_relative_path) or ".." in normalized_relative_path.split(os.sep):
+             print(f"Error: Invalid path characters or attempt to escape workspace in '{relative_path}'.")
+             return None
+        full_path = (AGENT_WORKSPACE / normalized_relative_path).resolve()
+        if AGENT_WORKSPACE.resolve() in full_path.parents or full_path == AGENT_WORKSPACE.resolve():
+             return full_path
+        # Check prefix as a fallback, although resolve should handle canonical paths
+        if str(full_path).startswith(str(AGENT_WORKSPACE.resolve())):
+             print(f"Warning: Path resolution for '{relative_path}' seems complex but within workspace: {full_path}")
+             return full_path
+        print(f"Error: Path '{relative_path}' resolved to '{full_path}' which is outside the allowed workspace '{AGENT_WORKSPACE.resolve()}'.")
+        return None
+    except Exception as e:
+        print(f"Error resolving path '{relative_path}': {e}")
         return None
 
 def _transcribe_audio(file_path: Path, file_description: str) -> str:
     """Helper to transcribe an audio file using OpenAI Whisper."""
-    if not openai_client:
-        return "Error: OpenAI client not available for transcription."
+    if not openai_client: return "Error: OpenAI client not available for transcription."
     if not file_path.is_file():
-        return f"Error: Audio file not found at '{file_path.relative_to(AGENT_WORKSPACE)}'"
-
+        try: rel_path_str = file_path.relative_to(AGENT_WORKSPACE)
+        except ValueError: rel_path_str = file_path
+        return f"Error: Audio file not found at '{rel_path_str}'"
     try:
         file_size_mb = file_path.stat().st_size / (1024 * 1024)
-        if file_size_mb > 25:
-            return f"Error: Audio file '{file_description}' is too large ({file_size_mb:.2f} MB). Max 25 MB."
-
+        if file_size_mb > 25: return f"Error: Audio file '{file_description}' is too large ({file_size_mb:.2f} MB). Max 25 MB."
         print(f"Transcribing audio: {file_description}...")
-        with open(file_path, "rb") as audio_file_handle:
-            # Note: response_format="text" returns a simple string
-            transcript = openai_client.audio.transcriptions.create(
-                model=WHISPER_MODEL,
-                file=audio_file_handle,
-                response_format="text"
-            )
+        with open(file_path, "rb") as audio_file_handle: transcript = openai_client.audio.transcriptions.create(model=WHISPER_MODEL, file=audio_file_handle, response_format="text")
         print("Transcription complete.")
-
-        if isinstance(transcript, str):
-            max_len = 10000 # Limit transcription length in output
-            if len(transcript) > max_len:
-                transcript = transcript[:max_len] + "\n... [Transcription truncated]"
-            return f"Transcription of '{file_description}':\n{transcript}"
-        else:
-            return f"Transcription of '{file_description}' succeeded, but format was unexpected: {type(transcript)}"
-
-    except openai.APIError as e:
-        return f"OpenAI API Error during transcription of '{file_description}': {e}"
-    except Exception as e:
-        return f"Error transcribing '{file_description}': {e}"
+        if isinstance(transcript, str): max_len = 10000; transcript = transcript[:max_len] + ("\n... [Transcription truncated]" if len(transcript) > max_len else ""); return f"Transcription of '{file_description}':\n{transcript}"
+        else: return f"Transcription of '{file_description}' succeeded, but format was unexpected: {type(transcript)}"
+    except openai.APIError as e: return f"OpenAI API Error during transcription of '{file_description}': {e}"
+    except Exception as e: return f"Error transcribing '{file_description}': {e}"
 
 # --- Tool Implementations ---
-
 @tool("write_file", args_schema=FileWriteArgs)
 def write_file(relative_path: str, content: str) -> str:
     """Writes text content to a file within the agent's workspace. Creates parent directories if needed."""
-    full_path = _resolve_path(relative_path)
-    if not full_path:
-        return f"Error: Invalid or disallowed path '{relative_path}'."
-    try:
-        full_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(full_path, 'w', encoding='utf-8') as f:
-            f.write(content)
-        return f"Successfully wrote to file: {relative_path}"
-    except Exception as e:
-        return f"Error writing file '{relative_path}': {e}"
+    full_path = _resolve_path(relative_path);
+    if not full_path: return f"Error: Invalid or disallowed path '{relative_path}'."
+    try: full_path.parent.mkdir(parents=True, exist_ok=True); open(full_path, 'w', encoding='utf-8').write(content); return f"Successfully wrote to file: {relative_path}"
+    except Exception as e: return f"Error writing file '{relative_path}': {e}"
 
 @tool("read_file", args_schema=FileReadArgs)
 def read_file(relative_path: str) -> str:
     """Reads the text content of a file from the agent's workspace. Limited read size."""
-    full_path = _resolve_path(relative_path)
-    if not full_path:
-        return f"Error: Invalid or disallowed path '{relative_path}'."
-    if not full_path.is_file():
-        return f"Error: File not found at '{relative_path}'"
+    full_path = _resolve_path(relative_path);
+    if not full_path: return f"Error: Invalid or disallowed path '{relative_path}'."
+    if not full_path.is_file(): return f"Error: File not found at '{relative_path}'"
     try:
-        with open(full_path, 'r', encoding='utf-8') as f:
-            content = f.read(10000) # Limit read size
-            if len(f.read(1)) > 0:
-                 content += "\n... [File truncated due to length]"
+        with open(full_path, 'r', encoding='utf-8') as f: content = f.read(10000); content += "\n... [File truncated due to length]" if len(f.read(1)) > 0 else ""
         return content
-    except Exception as e:
-        return f"Error reading file '{relative_path}': {e}"
+    except Exception as e: return f"Error reading file '{relative_path}': {e}"
 
 @tool("list_directory", args_schema=ListDirectoryArgs)
 def list_directory(relative_path: str = ".") -> str:
     """Lists the contents (files and directories) of a specified directory within the agent's workspace."""
-    target_path = _resolve_path(relative_path)
-    if not target_path:
-        return f"Error: Invalid or disallowed path '{relative_path}'."
-    if not target_path.is_dir():
-        return f"Error: '{relative_path}' is not a valid directory."
-    try:
-        items = [f.name + ('/' if f.is_dir() else '') for f in target_path.iterdir()]
-        if not items:
-            return f"Directory '{relative_path}' is empty."
-        return f"Contents of '{relative_path}':\n" + "\n".join(items)
-    except Exception as e:
-        return f"Error listing directory '{relative_path}': {e}"
+    target_path = _resolve_path(relative_path);
+    if not target_path: return f"Error: Invalid or disallowed path '{relative_path}'."
+    if not target_path.is_dir(): return f"Error: '{relative_path}' is not a valid directory."
+    try: items = [f.name + ('/' if f.is_dir() else '') for f in target_path.iterdir()]; items.sort(); return f"Contents of '{relative_path}':\n" + "\n".join(items) if items else f"Directory '{relative_path}' is empty."
+    except Exception as e: return f"Error listing directory '{relative_path}': {e}"
 
 @tool("run_python_code", args_schema=RunPythonCodeArgs)
 def run_python_code(code: str) -> str:
     """Executes Python code in a subprocess and returns the stdout/stderr. Use print() for output. WARNING: Executes arbitrary code."""
+    print(f"Executing Python code:\n```python\n{code}\n```")
     try:
-        process = subprocess.run(
-            [sys.executable, "-c", code],
-            capture_output=True, text=True, timeout=30, cwd=AGENT_WORKSPACE, check=False # Don't raise error on non-zero exit
-        )
+        process = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=30, cwd=AGENT_WORKSPACE, check=False)
         output, error = process.stdout, process.stderr
-        result = ""
-        if output:
-            max_output = 2000
-            if len(output) > max_output: output = output[:max_output] + "\n... [Output truncated]"
-            result += f"Output:\n{output}\n"
-        if error:
-             result += f"Error Output:\n{error}\n"
-
-        if process.returncode == 0:
-            return f"Execution successful.\n{result}"
-        else:
-            return f"Execution failed (Return Code: {process.returncode}).\n{result}"
-    except subprocess.TimeoutExpired:
-        return "Error: Code execution timed out after 30 seconds."
-    except Exception as e:
-        return f"Error executing Python code: {e}"
+        result = "Execution successful.\n" if process.returncode == 0 else f"Execution failed (Return Code: {process.returncode}).\n"
+        if output: max_output = 2000; output = output[:max_output] + ("\n... [Output truncated]" if len(output) > max_output else ""); result += f"Output:\n{output}\n"
+        if error: max_error = 1000; error = error[:max_error] + ("\n... [Error truncated]" if len(error) > max_error else ""); result += f"Error Output:\n{error}\n"
+        if not output and not error: result += "No output produced." if process.returncode == 0 else "No output or error message produced despite non-zero exit code."
+        return result.strip()
+    except subprocess.TimeoutExpired: return "Error: Code execution timed out after 30 seconds."
+    except Exception as e: return f"Error executing Python code: {e}"
 
 @tool("scrape_webpage", args_schema=WebScrapeArgs)
 def scrape_webpage(url: str, query: Optional[str] = None) -> str:
     """Scrapes text content from a given URL using BeautifulSoup. If a query is provided, returns content for the agent to answer it."""
+    print(f"Attempting to scrape URL: {url}")
     try:
-        headers = {'User-Agent': 'Mozilla/5.0 (compatible; GAIA-Agent/1.0)'} # Identify the agent
-        response = requests.get(url, headers=headers, timeout=20)
-        response.raise_for_status() # Raise HTTPError for bad responses
-
-        # Check content type - avoid trying to parse images, etc.
+        space_id = os.getenv("SPACE_ID", "YOUR_SPACE_ID")
+        headers = {'User-Agent': f'Mozilla/5.0 (compatible; GAIA-Agent/1.0; +https://huggingface.co/spaces/{space_id})'}
+        response = requests.get(url, headers=headers, timeout=20); response.raise_for_status()
         content_type = response.headers.get('content-type', '').lower()
-        if 'text/html' not in content_type:
-            return f"Error: Content type of URL {url} is '{content_type}', not HTML. Cannot scrape."
-
+        if 'text/html' not in content_type: return f"Error: Content type of URL {url} is '{content_type}', not HTML. Cannot scrape."
         soup = BeautifulSoup(response.text, 'html.parser')
-        for script_or_style in soup(["script", "style", "nav", "footer", "aside"]): # Remove common clutter
-            script_or_style.decompose()
-
-        text_content = soup.get_text(separator='\n', strip=True)
-        if not text_content: return f"Could not extract meaningful text content from {url}."
-
-        max_chars = 10000 # Limit content length
-        if len(text_content) > max_chars:
-            text_content = text_content[:max_chars] + "\n... [Content truncated]"
-
-        if query:
-            return f"Use the following content from {url} to answer the query '{query}':\n\n{text_content}"
-        else:
-            return f"Content scraped from {url}:\n\n{text_content}"
-
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching or reading URL {url}: {e}"
-    except Exception as e:
-        return f"Error scraping URL {url}: {e}"
-
-# --- Optional Tools (Conditionally Available) ---
+        for tag in soup(["script", "style", "nav", "footer", "aside", "header", "form", "button", "iframe", "noscript"]): tag.decompose()
+        text_content = soup.get_text(separator='\n', strip=True); text_content = '\n'.join(line for line in text_content.splitlines() if line.strip())
+        if not text_content: return f"Could not extract meaningful text content from {url} after cleaning."
+        max_chars = 10000; text_content = text_content[:max_chars] + ("\n... [Content truncated]" if len(text_content) > max_chars else "")
+        print(f"Scraping successful for {url}. Content length (approx): {len(text_content)}")
+        if query: return f"Use the following content from {url} to answer the query '{query}':\n\n{text_content}"
+        else: return f"Content scraped from {url}:\n\n{text_content}"
+    except requests.exceptions.Timeout: return f"Error: Timeout occurred while trying to fetch URL {url}"
+    except requests.exceptions.RequestException as e: return f"Error fetching or reading URL {url}: {e}"
+    except Exception as e: return f"Error scraping URL {url}: {e}"
 
 if PANDAS_AVAILABLE:
     @tool("read_excel_file", args_schema=ReadExcelArgs)
     def read_excel_file(relative_path: str, sheet_name: Optional[str] = None, max_rows_preview: int = 20) -> str:
         """Reads data from an Excel file (.xlsx or .xls) within the workspace and returns a text preview."""
-        full_path = _resolve_path(relative_path)
+        full_path = _resolve_path(relative_path);
         if not full_path: return f"Error: Invalid or disallowed path '{relative_path}'."
         if not full_path.is_file(): return f"Error: Excel file not found at '{relative_path}'"
+        print(f"Reading Excel file: {relative_path}")
         try:
             excel_file = pd.ExcelFile(full_path)
-            if sheet_name:
-                if sheet_name not in excel_file.sheet_names:
-                    return f"Error: Sheet '{sheet_name}' not found. Available: {excel_file.sheet_names}"
-                sheet_to_read = sheet_name
-            else:
-                sheet_to_read = excel_file.sheet_names[0]
-
+            if not excel_file.sheet_names: return f"Error: Excel file '{relative_path}' contains no sheets."
+            sheet_to_read = sheet_name if sheet_name and sheet_name in excel_file.sheet_names else excel_file.sheet_names[0]
+            if sheet_name and sheet_name not in excel_file.sheet_names: print(f"Warning: Sheet '{sheet_name}' not found, reading first sheet '{sheet_to_read}' instead.")
+            print(f"Reading sheet '{sheet_to_read}' from {relative_path}")
             df = pd.read_excel(full_path, sheet_name=sheet_to_read)
-            output = f"Preview of sheet '{sheet_to_read}' from '{relative_path}' ({df.shape[0]} rows, {df.shape[1]} cols):\n"
-            output += df.to_string(max_rows=max_rows_preview, max_cols=15) # Preview format
-
-            max_output_len = 5000
-            if len(output) > max_output_len:
-                output = output[:max_output_len] + "\n... [Output truncated]"
+            if df.empty: return f"Sheet '{sheet_to_read}' in '{relative_path}' is empty."
+            output = f"Preview of sheet '{sheet_to_read}' from '{relative_path}' ({df.shape[0]} rows, {df.shape[1]} columns):\n"
+            output += df.to_string(max_rows=max_rows_preview, max_cols=15, line_width=120)
+            max_output_len = 5000; output = output[:max_output_len] + ("\n... [Output truncated due to length]" if len(output) > max_output_len else "")
             return output
         except Exception as e: return f"Error reading Excel file '{relative_path}': {e}"
 
@@ -326,7 +218,7 @@ if OPENAI_AVAILABLE and openai_client:
     @tool("transcribe_audio_file", args_schema=TranscribeAudioArgs)
     def transcribe_audio_file(relative_path: str) -> str:
         """Transcribes audio content from a file in the workspace using OpenAI Whisper (max 25MB)."""
-        full_path = _resolve_path(relative_path)
+        full_path = _resolve_path(relative_path);
         if not full_path: return f"Error: Invalid or disallowed path '{relative_path}'."
         return _transcribe_audio(full_path, relative_path)
 
@@ -336,80 +228,57 @@ if PYTUBE_AVAILABLE and OPENAI_AVAILABLE and openai_client:
         """Downloads audio from a YouTube URL, transcribes it using OpenAI Whisper, and returns the text."""
         temp_audio_path = None
         try:
-            print(f"Processing YouTube URL: {youtube_url}")
-            yt = YouTube(youtube_url)
-            audio_stream = yt.streams.filter(only_audio=True, file_extension='mp4').order_by('abr').desc().first()
-            if not audio_stream: audio_stream = yt.streams.filter(only_audio=True).order_by('abr').desc().first() # Fallback
-            if not audio_stream: return f"Error: No suitable audio stream found for {youtube_url}"
-
-            # Download to a unique temporary file in workspace
-            timestamp = int(time.time())
-            temp_filename = f"temp_youtube_{timestamp}.{audio_stream.subtype or 'mp4'}"
+            print(f"Processing YouTube URL: {youtube_url}"); yt = YouTube(youtube_url, use_oauth=False, allow_oauth_cache=False)
+            print("Fetching available streams...")
+            audio_stream = yt.streams.filter(only_audio=True, subtype='webm').order_by('abr').desc().first() or \
+                           yt.streams.filter(only_audio=True, subtype='mp4').order_by('abr').desc().first() or \
+                           yt.streams.get_audio_only()
+            if not audio_stream: return f"Error: No suitable audio stream found for YouTube video: {youtube_url}"
+            print(f"Selected audio stream: Itag {audio_stream.itag}, ABR {audio_stream.abr}")
+            try: video_id = yt.video_id
+            except: video_id = f"vid_{int(time.time())}"
+            temp_filename = f"temp_youtube_{video_id}.{audio_stream.subtype or 'mp4'}"
             temp_audio_path = AGENT_WORKSPACE / temp_filename
             print(f"Downloading audio to: {temp_audio_path}...")
-            audio_stream.download(output_path=AGENT_WORKSPACE, filename=temp_filename)
-            print("Download complete.")
-
-            # Transcribe the downloaded file
-            result = _transcribe_audio(temp_audio_path, f"YouTube video '{yt.title}'")
-            return result
-
-        except PytubeError as e: return f"Error processing YouTube video {youtube_url}: {e}"
+            audio_stream.download(output_path=AGENT_WORKSPACE, filename=temp_filename); print("Download complete.")
+            result = _transcribe_audio(temp_audio_path, f"YouTube video '{yt.title}'"); return result
+        except PytubeError as e: return f"Error processing YouTube video {youtube_url} (PytubeError): {e}"
         except Exception as e: return f"Unexpected error during YouTube transcription {youtube_url}: {e}"
         finally:
-            # --- IMPORTANT: Clean up temporary file ---
             if temp_audio_path and temp_audio_path.exists():
-                try: temp_audio_path.unlink(); print(f"Cleaned up: {temp_audio_path}")
+                try: temp_audio_path.unlink(); print(f"Cleaned up temporary file: {temp_audio_path}")
                 except Exception as e: print(f"Warning: Failed to delete temp file {temp_audio_path}: {e}")
 
-
 # ==============================================================================
 # 3. AGENT SETUP
 # ==============================================================================
 
 # --- Initialize LLM ---
 try:
-    llm = ChatGroq(
-        temperature=0,
-        model_name=GROQ_MODEL_NAME,
-        groq_api_key=GROQ_API_KEY
-    )
+    llm = ChatGroq(temperature=0, model_name=GROQ_MODEL_NAME, groq_api_key=GROQ_API_KEY)
     print(f"Using Groq LLM: {GROQ_MODEL_NAME}")
-except Exception as e:
-    print(f"FATAL: Error initializing Groq LLM: {e}")
-    sys.exit(1)
+except Exception as e: print(f"FATAL: Error initializing Groq LLM: {e}"); sys.exit(1)
 
 # --- Assemble Available Tools ---
 available_tools = []
 if TAVILY_API_KEY:
-    available_tools.append(TavilySearchResults(max_results=TAVILY_MAX_RESULTS, api_key=TAVILY_API_KEY))
-else:
-    print("Warning: Tavily Search tool disabled (API key missing).")
-
-# Core tools are always added (they don't have external dependencies checked above)
-available_tools.extend([
-    write_file,
-    read_file,
-    list_directory,
-    run_python_code,
-    scrape_webpage,
-])
-
-# Add optional tools if their dependencies/clients are ready
+    try: available_tools.append(TavilySearchResults(max_results=TAVILY_MAX_RESULTS, api_key=TAVILY_API_KEY))
+    except Exception as e: print(f"Warning: Failed to initialize Tavily Search tool: {e}. Tool disabled.")
+else: print("Warning: Tavily Search tool disabled (API key missing).")
+available_tools.extend([write_file, read_file, list_directory, run_python_code, scrape_webpage])
 if PANDAS_AVAILABLE: available_tools.append(read_excel_file)
 if OPENAI_AVAILABLE and openai_client: available_tools.append(transcribe_audio_file)
 if PYTUBE_AVAILABLE and OPENAI_AVAILABLE and openai_client: available_tools.append(transcribe_youtube_video)
-
 print(f"Agent initialized with tools: {[tool.name for tool in available_tools]}")
 
 # --- Define System Prompt ---
-# This prompt is formatted later with the *actually available* tools
+# Contains {tools} and {agent_workspace} placeholders.
 SYSTEM_PROMPT_TEMPLATE = """You are a highly capable AI assistant designed to solve complex problems step-by-step, mimicking human-like reasoning and actions. Your goal is to accurately answer the user's request based on the GAIA benchmark philosophy.
 
 **Workspace:** You have access to a local workspace directory: '{agent_workspace}'. You can ONLY interact with files inside this directory using the provided tools. Always use relative paths for file operations.
 
 **Available Tools:** You have access to the following tools:
-{tool_descriptions}
+{tools}
 
 **Reasoning Process:**
 1.  **Understand:** Analyze the request. Identify objectives, constraints, and required information (text, web search, file content, Excel data, audio/video transcription, calculations).
@@ -433,149 +302,112 @@ SYSTEM_PROMPT_TEMPLATE = """You are a highly capable AI assistant designed to so
 """
 
 # --- Create Prompt Template ---
-prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", SYSTEM_PROMPT_TEMPLATE.format(
-            agent_workspace=AGENT_WORKSPACE.resolve(),
-            tool_descriptions="\n".join([f"- {tool.name}: {tool.description}" for tool in available_tools])
-            )
-        ),
-        MessagesPlaceholder(variable_name="chat_history"),
-        ("human", "{input}"),
-        MessagesPlaceholder(variable_name="agent_scratchpad"), # Crucial for agent's intermediate steps
-    ]
-)
+# Pre-format the system prompt string fully before creating the template
+try:
+    # Format the tool descriptions manually using the render_text_description utility
+    from langchain.tools.render import render_text_description
+    tool_descriptions = render_text_description(available_tools)
+
+    # Format the entire system prompt string
+    formatted_system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
+        agent_workspace=str(AGENT_WORKSPACE.resolve()),
+        tools=tool_descriptions
+    )
+
+    # Create the template from the fully formatted string
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", formatted_system_prompt), # Use the pre-formatted string
+            MessagesPlaceholder(variable_name="chat_history"),
+            ("human", "{input}"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"), # Still needed by the agent type
+        ]
+    )
+
+except Exception as e:
+    print(f"FATAL: Error creating ChatPromptTemplate: {e}")
+    sys.exit(1)
+
 
 # --- Setup Memory ---
 memory = ConversationBufferWindowMemory(
     k=MEMORY_WINDOW_SIZE,
     memory_key="chat_history",
-    return_messages=True # Return Message objects for chat models
+    return_messages=True
 )
 
 # --- Create Agent ---
-# Structured Chat Agent is generally good for models supporting tool calling/structured output
-agent = create_structured_chat_agent(llm, available_tools, prompt)
+# Using create_openai_tools_agent
+try:
+    agent = create_openai_tools_agent(llm, available_tools, prompt)
+except Exception as e:
+    print(f"FATAL: Error creating agent with create_openai_tools_agent: {e}")
+    import traceback
+    traceback.print_exc()
+    sys.exit(1)
 
 # --- Create Agent Executor ---
-agent_executor = AgentExecutor(
-    agent=agent,
-    tools=available_tools,
-    memory=memory,
-    verbose=True, # Set to True to see agent's thought process, False for cleaner output
-    max_iterations=MAX_ITERATIONS,
-    handle_parsing_errors="Please check your output format and try again.", # Basic guidance on format errors
-    # return_intermediate_steps=True # Uncomment to get intermediate steps in the result dictionary
-)
+try:
+    agent_executor = AgentExecutor(
+        agent=agent,
+        tools=available_tools,
+        memory=memory,
+        verbose=True,
+        max_iterations=MAX_ITERATIONS,
+        handle_parsing_errors=True,
+    )
+except Exception as e:
+    print(f"FATAL: Error creating AgentExecutor: {e}")
+    sys.exit(1)
 
 # ==============================================================================
-# 4. EXECUTION FUNCTION
+# 4. EXECUTION FUNCTION (Exported for app.py)
 # ==============================================================================
-
 def run_gaia_task(task_description: str):
-    """
-    Runs the GAIA agent on a given task description.
-
-    Args:
-        task_description: The natural language description of the task.
-
-    Returns:
-        The final output string from the agent, or an error message.
-    """
-    print("\n" + "="*50)
-    print(f"🚀 Running GAIA Task")
-    print(f"📝 Task: {task_description}")
-    print(f"📍 Workspace: {AGENT_WORKSPACE.resolve()}")
-    print(f"🛠️ Available Tools: {[tool.name for tool in available_tools]}")
-    print("="*50 + "\n")
-
-    # Reset memory for each new task to avoid context bleed
-    memory.clear()
-
+    """Runs the GAIA agent on a given task description. This is the main entry point."""
+    print("\n" + "="*50 + f"\n🚀 Running GAIA Task\n📝 Task: {task_description[:150]}...\n📍 Workspace: {AGENT_WORKSPACE.resolve()}\n🛠️ Tools: {[tool.name for tool in available_tools]}\n" + "="*50 + "\n")
+    memory.clear() # Reset memory for the task
     try:
-        # Invoke the agent executor
+        if 'agent_executor' not in globals() or agent_executor is None: return "Error: Agent Executor not initialized."
         result = agent_executor.invoke({"input": task_description})
-
-        print("\n" + "="*50)
-        print("✅ Agent Execution Finished")
-        print(f"🏁 Final Output:\n{result.get('output', 'No output found.')}")
-        print("="*50 + "\n")
-        return result.get('output', 'Agent finished but produced no output.')
-
+        final_output = result.get('output', 'Agent finished but produced no output.')
+        print("\n" + "="*50 + f"\n✅ Agent Execution Finished\n🏁 Final Output:\n{final_output}\n" + "="*50 + "\n")
+        return str(final_output)
     except Exception as e:
-        print(f"\n{'='*50}")
-        print(f"❌ Agent Execution Error")
-        print(f"An error occurred: {e}")
-        # Optional: Print traceback for detailed debugging
-        # import traceback
-        # traceback.print_exc()
-        print("="*50 + "\n")
+        print(f"\n{'='*50}\n❌ Agent Execution Error during task run\nAn error occurred: {e}\n{'='*50}\n")
+        import traceback; traceback.print_exc() # Print full traceback for debugging
         return f"Agent failed with error: {e}"
 
 # ==============================================================================
-# 5. EXAMPLE USAGE (Entry Point)
+# 5. EXAMPLE USAGE (Local Testing)
 # ==============================================================================
-
 if __name__ == "__main__":
-    # --- Optional: Setup Example Files ---
+    print("\n" + "*"*30 + " LOCAL TEST RUN " + "*"*30)
     print("--- Setting up example files (if needed) ---")
-    # Dummy Excel
     if PANDAS_AVAILABLE:
         try:
             dummy_excel_path = AGENT_WORKSPACE / "sample_data.xlsx"
-            if not dummy_excel_path.exists():
-                 pd.DataFrame({'ID': [1, 2, 3], 'Product': ['Widget', 'Gadget', 'Thingamajig']}).to_excel(dummy_excel_path, index=False)
-                 print(f"Created dummy Excel: {dummy_excel_path}")
+            if not dummy_excel_path.exists(): pd.DataFrame({'ID': [1, 2, 3], 'Product': ['Widget', 'Gadget', 'Thingamajig']}).to_excel(dummy_excel_path, index=False); print(f"Created dummy Excel: {dummy_excel_path}")
         except Exception as e: print(f"Could not create dummy Excel: {e}")
-    # Dummy Text
     try:
         dummy_text_path = AGENT_WORKSPACE / "numbers.txt"
         if not dummy_text_path.exists():
-            with open(dummy_text_path, "w") as f: f.write("15\n-3\n42.5\n100\n")
-            print(f"Created dummy text file: {dummy_text_path}")
+            with open(dummy_text_path, "w") as f: f.write("15\n-3\n42.5\n100\n"); print(f"Created dummy text file: {dummy_text_path}")
     except Exception as e: print(f"Could not create dummy text file: {e}")
-    # Dummy Audio - User needs to provide this manually
     dummy_audio_path = AGENT_WORKSPACE / "sample_audio.mp3"
-    if not dummy_audio_path.exists() and OPENAI_AVAILABLE and openai_client:
-        print(f"INFO: To test audio transcription, place an MP3 file at: {dummy_audio_path}")
+    if not dummy_audio_path.exists() and OPENAI_AVAILABLE and openai_client: print(f"INFO: To test audio transcription, place an MP3 file at: {dummy_audio_path}")
     print("--- Example setup complete ---")
 
-
-    # --- Define Example Tasks ---
-    task_list = [
-        {
-            "id": "excel_read",
-            "description": "Read the file 'sample_data.xlsx' in the workspace. What is the 'Product' where 'ID' is 2? Final answer should be just the product name."
-        },
-        {
-            "id": "python_sum",
-            "description": "Read the numbers from 'numbers.txt' in the workspace (one per line). Calculate their sum using python code. Write the sum into 'sum_result.txt'. Final answer should be the relative path 'sum_result.txt'."
-        },
-        {
-            "id": "search_scrape_write",
-            "description": "Search the web for the official website of the Python Software Foundation. Scrape the main title from the homepage of that website. Write the title into 'psf_title.txt'. Final answer is 'psf_title.txt'."
-        },
-        # { # Uncomment to run audio task if sample_audio.mp3 exists
-        #     "id": "audio_transcribe",
-        #     "description": "Transcribe the audio file 'sample_audio.mp3' from the workspace. Write the first 50 characters of the transcription into 'audio_snippet.txt'. Final answer is 'audio_snippet.txt'."
-        # },
-        # { # Uncomment to run YouTube task
-        #     "id": "youtube_transcribe",
-        #     "description": "Transcribe the YouTube video 'https://www.youtube.com/watch?v=dQw4w9WgXcQ'. What is the first line of the transcription? Final answer is just the first line."
-        # },
+    example_tasks = [
+        {"id": "local_excel_read", "description": "Read the file 'sample_data.xlsx' in the workspace. What is the 'Product' where 'ID' is 2? Final answer should be just the product name."},
+        {"id": "local_python_sum", "description": "Read the numbers from 'numbers.txt' in the workspace (one per line). Calculate their sum using python code. Write the sum into 'sum_result.txt'. Final answer should be the relative path 'sum_result.txt'."},
+        {"id": "local_search_scrape_write", "description": "Search the web for the official website of the Python Software Foundation. Scrape the main title from the homepage of that website. Write the title into 'psf_title.txt'. Final answer is 'psf_title.txt'."},
     ]
 
-    # --- Run Selected Task ---
-    # Choose which task to run by its index or ID
-    task_to_run = task_list[0] # Run the first task (Excel read)
-
-    print(f"\n>>> Running selected task: {task_to_run['id']} <<<")
-    final_answer = run_gaia_task(task_to_run['description'])
-    print(f">>> Task {task_to_run['id']} completed. Agent Output: {final_answer} <<<")
-
-    # To run all tasks:
-    # for task in task_list:
-    #     print(f"\n>>> Running task: {task['id']} <<<")
-    #     final_answer = run_gaia_task(task['description'])
-    #     print(f">>> Task {task['id']} completed. Agent Output: {final_answer} <<<")
-    #     input("Press Enter to continue to the next task...") # Pause between tasks
+    if example_tasks:
+        task_to_run = example_tasks[0] # Change index to test different tasks
+        print(f"\n>>> Running local test task: {task_to_run['id']} <<<")
+        final_answer = run_gaia_task(task_to_run['description'])
+        print(f">>> Local test task {task_to_run['id']} completed. Agent Output: {final_answer} <<<")
+    else: print("No example tasks defined for local testing.")
+    print("\n" + "*"*30 + " LOCAL TEST RUN COMPLETE " + "*"*30)