Update app.py

app.py

@@ -1,473 +1,99 @@
-import os, re, requests, pandas as pd, gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace
-from langchain_community.tools import DuckDuckGoSearchRun
-from langchain.tools import tool
-from langchain_core.output_parsers import JsonOutputParser
-from langchain.agents import AgentExecutor, create_react_agent, initialize_agent, AgentType
-from youtube_transcript_api import YouTubeTranscriptApi, TranscriptsDisabled, NoTranscriptFound
-import torch
-import traceback # For detailed error logging
-#import timeout_decorator # Optional: for execute_code timeout (pip install timeout-decorator)
-import whisper
-import chess, chess.engine
-from bs4 import BeautifulSoup
-from SPARQLWrapper import SPARQLWrapper, JSON
+import os
+import threading
+import requests
+import pandas as pd
+import gradio as gr
+
+from dotenv import load_dotenv
+
+from text_inspector_tool import TextInspectorTool
+from text_web_browser import (
+    ArchiveSearchTool,
+    FinderTool,
+    FindNextTool,
+    PageDownTool,
+    PageUpTool,
+    SimpleTextBrowser,
+    VisitTool,
+)
+from visual_qa import visualizer
+
+from smolagents import (
+    CodeAgent,
+    GoogleSearchTool,
+    LiteLLMModel,
+    ToolCallingAgent,
+)
+
+load_dotenv(override=True)
 
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 HF_TOKEN = os.getenv("HF_TOKEN", None)
-REQUESTS_TIMEOUT = 15 # Define a standard timeout for requests
-HEADERS = {'User-Agent': 'GAIAgent/1.0 (Langchain Agent; +http://example.com/info)'} # Be a good citizen
-
-@tool
-def web_search(query: str) -> str:
-    """Runs a web search and returns the results."""
-    search = DuckDuckGoSearchRun()
-    return search.run(query)
-
-@tool
-def read_file(file_path: str) -> str:
-    """Reads the content of a text file."""
-    try:
-        with open(file_path, 'r', encoding='utf-8') as f:
-            return f.read()
-    except Exception as e:
-        return f"Error reading file {file_path}: {e}"
-
-@tool
-def transcribe_audio(file_path: str) -> str:
-    """Transcribes audio from a file path."""
-    try:
-        # Load model here or use pre-loaded one
-        model = whisper.load_model("base") # Or tiny, small, medium, large
-        result = model.transcribe(file_path)
-        return result["text"]
-    except Exception as e:
-        return f"Error transcribing audio file {file_path}: {e}"
-
-@tool
-def analyze_sales_data(file_path: str) -> str:
-    """Reads the specific sales data Excel file, calculates total food sales."""
-    try:
-        df = pd.read_excel(file_path)
-        # Assuming columns 'Category' and 'Total Sales' exist
-        food_sales = df[df['Category'] != 'Drink']['Total Sales'].sum()
-        return f"${food_sales:.2f}" # Format as USD
-    except Exception as e:
-        return f"Error processing sales data from {file_path}: {e}"
-
-@tool
-def find_chess_mate_move(fen: str) -> str:
-    """
-    Given a FEN string representing a chess position (Black to move),
-    finds the best move that guarantees a win using Stockfish engine.
-    Requires Stockfish engine installed at engine_path.
-    Returns the move in algebraic notation (e.g., 'Qh4').
-    """
-    try:
-        engine = chess.engine.SimpleEngine.popen_uci("/usr/bin/stockfish")
-        board = chess.Board(fen)
-        if board.turn != chess.BLACK:
-            return "Error: It's not Black's turn in the provided FEN."
-        info = engine.analyse(board, chess.engine.Limit(time=2.0))
-
-        score = info.get("score")
-        if score is not None and score.is_mate():
-             mate_score = score.white().mate()
-             if mate_score < 0:
-                 best_move = info["pv"][0]
-                 engine.quit()
-                 return best_move.uci()
-        elif score is not None and score.relative.score(mate_score=10000) < -500: # Significant advantage for Black (-5 pawns)
-            best_move = info["pv"][0]
-            engine.quit()
-            return best_move.uci()
-
-        result = engine.play(board, chess.engine.Limit(time=1.0)) # Get a move anyway
-        engine.quit()
-        #return f"No guaranteed mate found quickly. Best move found: {result.move.uci()}"
-        return result.move.uci() # Return best move found even if not provably mate
 
-    except Exception as e:
-        return f"Chess engine error: {e}. Is Stockfish installed at {engine_path} and is the FEN valid?"
-    
-@tool
-def wiki_get_page(title: str) -> str:
-    """
-    Fetch raw wikitext content for a given English Wikipedia page title.
-    Returns the page content as a string or an error message.
-    Note: Raw wikitext can be complex to parse.
-    """
-    API = "https://en.wikipedia.org/w/api.php"
-    params = {
-        "action": "query",
-        "format": "json",
-        "prop": "revisions",
-        "rvprop": "content",
-        "rvslots": "*",
-        "titles": title,
-        "redirects": 1 # Automatically follow redirects
+# Browser config copied verbatim
+user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64)…"
+BROWSER_CONFIG = {
+    "viewport_size": 5120,
+    "downloads_folder": "downloads_folder",
+    "request_kwargs": {
+        "headers": {"User-Agent": user_agent},
+        "timeout": 300,
+    },
+    "serpapi_key": os.getenv("SERPAPI_API_KEY"),
+}
+os.makedirs(BROWSER_CONFIG["downloads_folder"], exist_ok=True)
+
+
+def create_agent(model_id="o1"):
+    custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
+    model_params = {
+        "model_id": model_id,
+        "custom_role_conversions": custom_role_conversions,
+        "max_completion_tokens": 8192,
     }
-    try:
-        response = requests.get(API, params=params, timeout=REQUESTS_TIMEOUT, headers=HEADERS)
-        response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx)
-        data = response.json()
-        page = next(iter(data["query"]["pages"].values()))
-
-        if "missing" in page:
-            return f"Error: Wikipedia page '{title}' not found."
-        if "invalid" in page:
-            return f"Error: Invalid page title '{title}' requested."
-        if "revisions" not in page or not page["revisions"]:
-             return f"Error: No revisions found for page '{title}' (page might be empty or protected)."
-
-        # Access content safely
-        content = page["revisions"][0].get("slots", {}).get("main", {}).get("*")
-        if content is None:
-            return f"Error: Could not extract main content slot for page '{title}'."
-        return content
-
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching Wikipedia page '{title}': Network error - {e}"
-    except KeyError as e:
-        return f"Error parsing Wikipedia response for '{title}': Unexpected structure - missing key {e}"
-    except Exception as e:
-        return f"An unexpected error occurred fetching Wikipedia page '{title}': {e}"
-
-@tool
-def youtube_transcript(video_id: str) -> str:
-    """
-    Retrieve the English transcript for a given YouTube video ID.
-    Returns the transcript as a single string or an error message.
-    """
-    try:
-        # Fetch available transcripts and prioritize English
-        transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
-        transcript = transcript_list.find_generated_transcript(['en']) # Prefer generated English
-        # You could add fallbacks here for manual 'en' or other languages if needed
-        # transcript = transcript_list.find_manually_created_transcript(['en'])
-        # transcript = transcript_list.find_transcript(['en', 'en-US', ...])
-
-        full_transcript = transcript.fetch()
-        return " ".join(t["text"] for t in full_transcript)
-    except (TranscriptsDisabled, NoTranscriptFound):
-        return f"Error: Transcripts are disabled or no English transcript found for YouTube video ID '{video_id}'."
-    except Exception as e:
-        # Catch other potential errors from the API or network issues
-        return f"An unexpected error occurred fetching transcript for YouTube video ID '{video_id}': {e}"
-
-@tool
-def reverse_text(text: str) -> str:
-    """Reverses the input string character by character."""
-    if not isinstance(text, str):
-        return "Error: Input must be a string."
-    return text[::-1]
-
-@tool
-def find_non_commutative(table: dict) -> str:
-    """
-    Given a dictionary representing a multiplication table (keys are tuples (row_elem, col_elem)),
-    finds all elements involved in non-commutative pairs (where table[(x,y)] != table[(y,x)]).
-    Returns a comma-separated list of these elements in alphabetical order, or an error message.
-    Example input: {('a','a'):'a', ('a','b'):'c', ('b','a'):'b', ...}
-    """
-    try:
-        if not isinstance(table, dict):
-            return "Error: Input must be a dictionary."
-        if not all(isinstance(k, tuple) and len(k) == 2 for k in table.keys()):
-            return "Error: Dictionary keys must be tuples of length 2, e.g., ('a', 'b')."
-
-        elems = sorted(list(set(x for k in table.keys() for x in k))) # Get all unique elements alphabetically
-        bad_elements = set()
-
-        for x in elems:
-            for y in elems:
-                # Check if both pairs exist in the table before comparing
-                pair_xy = (x, y)
-                pair_yx = (y, x)
-                if pair_xy in table and pair_yx in table:
-                    if table[pair_xy] != table[pair_yx]:
-                        bad_elements.add(x)
-                        bad_elements.add(y)
-                # Optional: Handle cases where one pair exists but the other doesn't,
-                # depending on how strictly commutativity should be defined for partial tables.
-                # else:
-                #     # If one exists and the other doesn't, it could be considered non-commutative
-                #     # or simply an incomplete table. Current logic ignores this.
-                #     pass
-
-
-        if not bad_elements:
-             return "Result: The operation defined by the table is commutative for all checked pairs."
-        return ",".join(sorted(list(bad_elements)))
-
-    except Exception as e:
-        return f"An unexpected error occurred processing the table: {e}"
-
-
-@tool
-def libretext_extract(query: str) -> str:
-    """
-    Extracts text content from a web page using a URL and a CSS selector.
-    Input must be a string formatted as 'url||css_selector'.
-    Returns the text of the first matching element or an error message.
-    """
-    try:
-        if "||" not in query:
-            return "Error: Input format must be 'url||css_selector'."
-        url, selector = query.split("||", 1)
-
-        response = requests.get(url, timeout=REQUESTS_TIMEOUT, headers=HEADERS)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, "html.parser")
-        element = soup.select_one(selector)
-
-        if element:
-            return element.get_text(strip=True)
-        else:
-            return f"Error: CSS selector '{selector}' did not find any elements on page {url}."
-
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching URL '{url}': Network error - {e}"
-    except Exception as e:
-        # Catch potential errors from BeautifulSoup or invalid selectors
-        return f"An unexpected error occurred during extraction from {url}: {e}"
-
-@tool
-def classify_vegetables(items: list) -> str:
-    """
-    Filters a list of items, keeping only those considered common culinary vegetables.
-    Returns a comma-separated, alphabetized list of the identified vegetables.
-    Note: This uses a predefined list and may not align perfectly with botanical definitions
-    (e.g., tomatoes, bell peppers are botanically fruits but often treated as vegetables).
-    Input items should be strings.
-    """
-    # Using a case-insensitive comparison by converting known veggies to lowercase
-    # Added more items, still imperfect and culturally dependent.
-    VEGETABLE_SET = {
-        "broccoli", "celery", "green beans", "lettuce", "zucchini", "sweet potato", # original + fixed space
-        "carrot", "spinach", "kale", "onion", "garlic", "potato", "cabbage", "asparagus",
-        "cucumber", # Botanically fruit, culinary vegetable
-        "bell pepper", # Botanically fruit, culinary vegetable
-        "corn", # Botanically fruit/grain, culinary vegetable
-        # Avoid controversial ones like tomato unless explicitly needed
-    }
-    try:
-        if not isinstance(items, list):
-            return "Error: Input must be a list of strings."
-        # Filter using lowercase comparison
-        vegetables = sorted([item for item in items if isinstance(item, str) and item.lower() in VEGETABLE_SET])
-        if not vegetables:
-            return "Result: No items from the list were classified as vegetables based on the predefined set."
-        return ",".join(vegetables)
-    except Exception as e:
-        return f"An unexpected error occurred classifying vegetables: {e}"
-
-@tool
-# Optional: Add timeout to prevent runaway code execution
-#@timeout_decorator.timeout(10, timeout_exception=TimeoutError) # Limit execution to 10 seconds
-def execute_code(code: str) -> str:
-    """
-    Executes a given Python code snippet and returns the value of the 'output' variable.
-    WARNING: Executes arbitrary code. Use with extreme caution in trusted environments only.
-    The code runs in a restricted environment, but vulnerabilities might exist.
-    The code should assign its result to a variable named 'output'.
-    Example: "output = sum([1, 2, 3])"
-    """
-    print(f"[!!!] Executing potentially unsafe code:\n---\n{code}\n---") # Log execution
-    local_ns = {}
-    # Restrict builtins more severely for safety. Allow only necessary ones.
-    # This is still not perfectly safe. Sandboxing is complex.
-    safe_builtins = {
-        'print': print, # Allow print for debugging within the code
-        'range': range, 'len': len, 'list': list, 'dict': dict, 'set': set,
-        'str': str, 'int': int, 'float': float, 'bool': bool, 'sum': sum,
-        'min': min, 'max': max, 'abs': abs, 'pow': pow, 'round': round,
-        'True': True, 'False': False, 'None': None,
-        # Add other safe builtins carefully if absolutely required by expected code snippets
-    }
-    # Also restrict imports if possible, though exec doesn't directly prevent them easily.
-
-    try:
-        # Using exec within a function's local scope
-        exec(code, {"__builtins__": safe_builtins}, local_ns)
-        # Check if 'output' was assigned, otherwise return empty string or error
-        output_val = local_ns.get("output", None)
-        if output_val is None:
-            return "Result: Code executed, but no variable named 'output' was assigned."
-        return str(output_val)
-    except TimeoutError:
-        return "Error: Code execution timed out."
-    except Exception as e:
-        # Capture and return execution errors
-        error_details = traceback.format_exc()
-        print(f"Error during code execution: {e}\n{error_details}") # Log full traceback
-        return f"Error during code execution: {type(e).__name__}: {e}"
-
-
-@tool
-def least_athletes_olympics(year: int) -> str:
-    """
-    Finds the country (IOC code) that sent the fewest athletes to the specified Summer Olympics year.
-    Data is scraped from the English Wikipedia page for that year's Olympics.
-    Returns the IOC code as a string. If there's a tie, returns the first code alphabetically.
-    Returns an error message if data cannot be retrieved or parsed.
-    """
-    try:
-        if not isinstance(year, int):
-             return "Error: Year must be an integer."
-
-        url = f"https://en.wikipedia.org/wiki/{year}_Summer_Olympics"
-        response = requests.get(url, timeout=REQUESTS_TIMEOUT, headers=HEADERS)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, "html.parser")
-
-        # Find the participating NOCs table - this selector might need adjustment over time
-        # Look for tables with captions containing 'Participating National Olympic Committees' or similar
-        tables = soup.find_all("table", class_="wikitable")
-        noc_table = None
-        for table in tables:
-             caption = table.find("caption")
-             # Check caption text or look for characteristic headers like 'NOC', 'Athletes'
-             if caption and "Participating National Olympic" in caption.get_text():
-                 noc_table = table
-                 break
-             # Fallback: check headers if no caption found or caption doesn't match
-             headers = [th.get_text(strip=True).lower() for th in table.find_all("th")]
-             if "noc" in headers and "athletes" in headers:
-                  noc_table = table
-                  break
-
-        if noc_table is None:
-            return f"Error: Could not find the expected NOC table on the Wikipedia page for {year} Summer Olympics."
-
-        rows = noc_table.find_all("tr")[1:] # Skip header row
-        data = []
-        for r in rows:
-            cols = r.find_all("td")
-            # Adapt column indices based on typical table structure (NOC code, Athletes count)
-            # This is fragile and depends on Wikipedia's table layout.
-            try:
-                # Attempt to find columns by text content or relative position
-                # Assuming NOC code is often linked, e.g., inside an <a> tag
-                noc_link = cols[0].find("a")
-                noc_code = noc_link.get_text(strip=True) if noc_link else cols[0].get_text(strip=True)
-                # Clean up potential bracketed numbers like (123) in NOC code cell
-                noc_code = re.sub(r'\s*\(\d+\)\s*$', '', noc_code).strip()
-
-                # Find athletes column - often the next column, check if it's numeric
-                athletes_text = cols[1].get_text(strip=True).replace(',', '') # Remove commas
-                athletes_count = int(athletes_text)
-
-                data.append((noc_code, athletes_count))
-            except (IndexError, ValueError, AttributeError):
-                # Skip rows that don't match the expected format
-                print(f"Skipping malformed row in table for {year}: {r.get_text(strip=True)}")
-                continue
-
-        if not data:
-             return f"Error: No valid NOC/athlete data parsed from the table for {year}."
-
-        min_athletes = min(count for _, count in data)
-        candidates = sorted([code for code, count in data if count == min_athletes])
-
-        if not candidates:
-             return f"Error: Could not determine country with fewest athletes for {year}."
-        return candidates[0]
-
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching Olympics page for {year}: Network error - {e}"
-    except Exception as e:
-        return f"An unexpected error occurred processing Olympics data for {year}: {e}\n{traceback.format_exc()}"
-
-
-@tool
-def get_nasa_award_number(qid: str) -> str:
-    """
-    Retrieves the NASA award number (property P496) associated with a given Wikidata Item QID.
-    Input must be a valid Wikidata QID string (e.g., 'Q42').
-    Returns the award number as a string, or an error message.
-    """
-    if not isinstance(qid, str) or not re.match(r'^Q\d+$', qid):
-        return f"Error: Invalid Wikidata QID format provided: '{qid}'. Must be like 'Q42'."
-
-    sparql = SPARQLWrapper("https://query.wikidata.org/sparql")
-    sparql.setMethod('POST') # Recommended by Wikidata for robustness
-    sparql.agent = HEADERS['User-Agent'] # Set User-Agent for SPARQL queries
-
-    query = f"""
-    SELECT ?award WHERE {{
-      wd:{qid} wdt:P496 ?award .
-    }}
-    LIMIT 1
-    """
-    sparql.setQuery(query)
-    sparql.setReturnFormat(JSON)
+    if model_id == "o1":
+        model_params["reasoning_effort"] = "high"
+    model = LiteLLMModel(**model_params)
+
+    browser = SimpleTextBrowser(**BROWSER_CONFIG)
+    WEB_TOOLS = [
+        GoogleSearchTool(provider="serper"),
+        VisitTool(browser),
+        PageUpTool(browser),
+        PageDownTool(browser),
+        FinderTool(browser),
+        FindNextTool(browser),
+        ArchiveSearchTool(browser),
+        TextInspectorTool(model, text_limit=100000),
+    ]
+    text_webbrowser_agent = ToolCallingAgent(
+        model=model,
+        tools=WEB_TOOLS,
+        max_steps=20,
+        verbosity_level=2,
+        planning_interval=4,
+        name="search_agent",
+        description="""
+A team member that will search the internet to answer your question.
+Ask him for all your questions that require browsing the web.
+Provide him as much context as possible…
+""",
+        provide_run_summary=True,
+    )
 
-    try:
-        results = sparql.query().convert()
-        bindings = results.get("results", {}).get("bindings", [])
-
-        if bindings:
-            award = bindings[0].get("award", {}).get("value")
-            if award:
-                return award
-            else:
-                return f"Error: Found property P496 for {qid}, but the award value is missing."
-        else:
-            return f"Error: No NASA award number (P496) found for Wikidata item {qid}."
+    manager_agent = CodeAgent(
+        model=model,
+        tools=[visualizer, TextInspectorTool(model, 100000)],
+        max_steps=12,
+        verbosity_level=2,
+        additional_authorized_imports=["*"],
+        planning_interval=4,
+        managed_agents=[text_webbrowser_agent],
+    )
 
-    except Exception as e:
-        # Catch SPARQL query errors, network issues, JSON parsing problems
-        return f"An error occurred querying Wikidata for {qid}: {e}"
-
-TOOLS = [
-    web_search,
-    read_file,
-    transcribe_audio,
-    analyze_sales_data, # Or a more general excel tool
-    find_chess_mate_move, # Needs image-to-FEN first!
-    wiki_get_page,
-    youtube_transcript,
-    reverse_text,
-    find_non_commutative,
-    libretext_extract,
-    classify_vegetables,
-    execute_code,
-    least_athletes_olympics,
-    get_nasa_award_number
-]
-
-SYSTEM_MESSAGE = """You are a concise AI assistant with access to the following tools:
-- web_search(query: string) -> string
-- wiki_get_page(title: string) → string
-- youtube_transcript(video_id: string) → string
-- reverse_text(text: string) → string
-- find_non_commutative(table: dict[tuple[string, string]: string]) -> string
-- libretext_extract(url: string, selector: string) → string
-- classify_vegetables(items: list[string]) → list[string]
-- execute_code(code: string) → string
-- least_athletes_olympics(year: int) → string
-- get_nasa_award_number(qid: string) → string
-- read_file(file_path: string) -> string
-- transcribe_audio(file_path: string) -> string
-- analyze_sales_data(file_path: string) -> string
-- find_chess_mate_move(fen: string) -> string
-When you need to use a tool, respond exactly with:
-Action: <tool_name>(<arg_name>=<value>, ...)
-
-IMPORTANT FORMATTING:
-- For the find_non_commutative tool, the 'table' argument MUST be a valid Python dictionary with tuple keys, like this:
-  Action: find_non_commutative(table={('a','a'):'a', ('a','b'):'c', ('b','a'):'d', ('b','b'):'e'})
-
-Then wait for the tool’s output before continuing.
-If a tool returns an error message starting with 'Error:', treat that as the observation.
-If a tool requires a file path, assume the file is accessible in the current environment.
-If a question involves an image or audio file, state that you need the content extracted first (e.g., text from audio, FEN from chess image) before you can proceed.
-Once you have all the information, provide your final answer in as few words as possible, with no extra commentary or prefixes.
-"""
+    return manager_agent
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -478,91 +104,13 @@ class BasicAgent:
             raise ValueError("HF_TOKEN not set in environment")
             
         # --- Replace with your chosen LLM ---
-        model_id = "microsoft/Phi-3-mini-4k-instruct"
-
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) # Some models need trust_remote_code
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                torch_dtype=torch.float32, # Use float32 for CPU compatibility usually
-                device_map=None, # Explicitly set to None or 'cpu' for CPU
-                trust_remote_code=True
-            )
-            model.to('cpu') # Ensure model is on CPU
-
-            pipe = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=512,
-                do_sample=False,
-                return_full_text=False,
-                # No temperature/top_k needed if do_sample=False
-            )
-            self.llm = HuggingFacePipeline(pipeline=pipe)
-
-        except ImportError as e:
-             raise ImportError(f"Required library not found: {e}. Make sure 'transformers', 'torch', 'accelerate' are installed.")
-        except Exception as e:
-             # Catch potential issues like model download failure, OOM errors
-             raise RuntimeError(f"Failed to initialize HuggingFacePipeline for {model_id}: {e}")
-
-        try:
-            # Construct a simplified test prompt (mimicking agent input)
-            test_prompt = SYSTEM_MESSAGE + "\nHuman: What is the capital of France?\nAssistant:" # A simple question
-            # Or use a prompt closer to the problematic one if you know which one it is
-            # test_prompt = SYSTEM_MESSAGE + "\nHuman: [Insert the non-commutative table question here]\nAssistant:"
-
-            # Use invoke which is standard now
-            test_response = self.llm.invoke(test_prompt)
-            print(f"--- Direct LLM Test Response ---:\n{test_response}\n-----------------------------")
-            if not test_response or len(test_response.strip()) == 0:
-                 print("!!! Direct LLM Test returned empty or whitespace result.")
-
-        except Exception as test_e:
-            print(f"!!! Direct LLM Test FAILED: {test_e}")
-            # print traceback for more details
-            import traceback
-            traceback.print_exc()
+        self.agent = create_agent(model_id="o1")
         print("BasicAgent initialized with LLM.")
 
     def __call__(self, question: str) -> str:
         # Comment out agent call temporarily if testing in __init__
         # Or add the direct test here before calling the agent
-        print(f"\n>> Processing Question (Agent): {question}")
-        try:
-            response = self.agent.invoke({"input": question})
-            answer = response.get('output', "Agent did not produce an output.")
-            print(f"<< Agent Answer: {answer}")
-            return str(answer).strip()
-        except Exception as e:
-             print(f"Error during agent execution: {e}")
-             # Also print traceback here to see where the error originates
-             import traceback
-             traceback.print_exc()
-             return f"Agent Error: {e}"
-        
-        # --- Agent Initialization (remains the same) ---
-        self.agent = initialize_agent(
-            tools=TOOLS,
-            llm=self.llm,
-            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            agent_kwargs={'prefix': SYSTEM_MESSAGE},
-            verbose=True,
-            handle_parsing_errors="Check your output and make sure it conforms!",
-            max_iterations=10
-        )
-        print("BasicAgent initialized with LLM.")
-
-    # --- Core dispatcher/fallback ---
-    def __call__(self, question: str) -> str:
-        try:
-             response = self.agent.invoke({"input": question})
-             answer = response.get('output', "Agent did not produce an output.")
-             return str(answer).strip()
-        except Exception as e:
-             print(f"Error during agent execution: {e}")
-             return f"Agent Error: {e}"
+        return self.agent.run(question).strip()
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """