Update agent.py

agent.py

@@ -1,297 +1,79 @@
-# -*- coding: utf-8 -*-
-"""
-Initializes and configures a SmolAgents CodeAgent with custom tools
-for file handling and web interaction.
-"""
-
-import importlib.resources
+import importlib
 import os
-import logging # Added for logging errors
-from typing import Type # Added for more specific type hints
 
 import requests
 import yaml
 import pandas as pd
 
-# Assuming 'config.py' exists in the same directory or Python path
-# and contains: DEFAULT_API_URL = "your_api_url_here"
-try:
-    from config import DEFAULT_API_URL
-except ImportError:
-    # Provide a default or raise a more specific error if config is crucial
-    DEFAULT_API_URL = "http://localhost:8000" # Example default, adjust as needed
-    logging.warning("config.py not found or DEFAULT_API_URL not set. Using default: %s", DEFAULT_API_URL)
-
-from smolagents import (
-    CodeAgent,
-    Tool,
-    OpenAIServerModel,
-    # Standard Tools
-    DuckDuckGoSearchTool,
-    VisitWebpageTool,
-    WikipediaSearchTool,
-    SpeechToTextTool,
-)
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# --- Custom Tools ---
+from config import DEFAULT_API_URL
+from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, WikipediaSearchTool, Tool, OpenAIServerModel, SpeechToTextTool
 
 class GetTaskFileTool(Tool):
-    """
-    A tool to download a file associated with a specific task ID from a predefined API endpoint.
-    """
     name = "get_task_file_tool"
-    description = "Downloads the file content associated with the given task_id if it exists. Returns the absolute file path of the downloaded file."
+    description = """This tool downloads the file content associated with the given task_id if exists. Returns absolute file path"""
     inputs = {
-        "task_id": {"type": "string", "description": "The unique identifier for the task."},
-        "file_name": {"type": "string", "description": "The desired local name for the downloaded file."},
+        "task_id": {"type": "string", "description": "Task id"},
+        "file_name": {"type": "string", "description": "File name"},
     }
-    output_type = "string" # Output is the file path or an error message
+    output_type = "string"
 
     def forward(self, task_id: str, file_name: str) -> str:
-        """
-        Executes the file download process.
-
-        Args:
-            task_id: The ID of the task whose file should be downloaded.
-            file_name: The name to save the downloaded file as locally.
-
-        Returns:
-            The absolute path to the downloaded file if successful,
-            otherwise an error message string.
-        """
-        url = f"{DEFAULT_API_URL}/files/{task_id}"
-        logging.info("Attempting to download file from: %s", url)
-        try:
-            response = requests.get(url, timeout=30) # Increased timeout slightly
-            response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
-
-            # Ensure the directory exists if file_name includes a path
-            # For simplicity here, we assume file_name is just a name,
-            # and it's saved in the current working directory.
-            # Consider adding directory creation logic if needed:
-            # os.makedirs(os.path.dirname(file_path), exist_ok=True)
-
-            file_path = os.path.abspath(file_name)
-            with open(file_path, 'wb') as file:
-                file.write(response.content)
-            logging.info("File successfully downloaded and saved to: %s", file_path)
-            return file_path
-
-        except requests.exceptions.RequestException as e:
-            error_msg = f"Error downloading file for task {task_id}: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message for the agent
-        except IOError as e:
-            error_msg = f"Error saving file {file_name}: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message
-
+        response = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=15)
+        response.raise_for_status()
+        with open(file_name, 'wb') as file:
+            file.write(response.content)
+        return os.path.abspath(file_name)
 
 class LoadXlsxFileTool(Tool):
-    """
-    A tool to load data from an XLSX (Excel) file into a pandas DataFrame.
-    """
     name = "load_xlsx_file_tool"
-    description = "Loads data from an XLSX file specified by its path into a pandas DataFrame."
+    description = """This tool loads xlsx file into pandas and returns it"""
     inputs = {
-        "file_path": {"type": "string", "description": "The local path to the XLSX file."}
+        "file_path": {"type": "string", "description": "File path"}
     }
-    # Using object is acceptable here as DataFrames are complex types,
-    # but adding pandas type hint for internal clarity.
     output_type = "object"
 
-    def forward(self, file_path: str) -> pd.DataFrame | str:
-        """
-        Executes the XLSX file loading process.
-
-        Args:
-            file_path: The path to the XLSX file.
-
-        Returns:
-            A pandas DataFrame containing the data from the first sheet
-            if successful, otherwise an error message string.
-        """
-        logging.info("Attempting to load XLSX file: %s", file_path)
-        try:
-            # Ensure the file exists before attempting to read
-            if not os.path.exists(file_path):
-                raise FileNotFoundError(f"No such file or directory: '{file_path}'")
-
-            # Load the excel file. You might want to add options like sheet_name=None
-            # to load all sheets into a dictionary of DataFrames if needed.
-            df = pd.read_excel(file_path)
-            logging.info("Successfully loaded XLSX file into DataFrame.")
-            # Note: Returning the actual DataFrame object for the agent to use.
-            # The agent's Python execution environment needs pandas installed.
-            return df
-
-        except FileNotFoundError as e:
-            error_msg = f"Error loading XLSX: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message
-        except Exception as e:
-            # Catch other potential errors during pandas read_excel (e.g., bad format, permissions)
-            # xlrd might be needed for .xls, openpyxl for .xlsx
-            error_msg = f"Error reading Excel file {file_path}: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message
-
+    def forward(self, file_path: str) -> object:
+        return pd.read_excel(file_path)
 
 class LoadTextFileTool(Tool):
-    """
-    A tool to load the content of a text file into a single string.
-    """
     name = "load_text_file_tool"
-    description = "Loads the entire content of any text file specified by its path."
+    description = """This tool loads any text file"""
     inputs = {
-        "file_path": {"type": "string", "description": "The local path to the text file."}
+        "file_path": {"type": "string", "description": "File path"}
     }
-    output_type = "string" # Output is the file content or an error message
+    output_type = "string"
 
-    def forward(self, file_path: str) -> str:
-        """
-        Executes the text file loading process.
+    def forward(self, file_path: str) -> object:
+        with open(file_path, 'r', encoding='utf-8') as file:
+            return file.read()
 
-        Args:
-            file_path: The path to the text file.
 
-        Returns:
-            The content of the text file as a string if successful,
-            otherwise an error message string.
-        """
-        logging.info("Attempting to load text file: %s", file_path)
-        try:
-             # Ensure the file exists before attempting to read
-            if not os.path.exists(file_path):
-                raise FileNotFoundError(f"No such file or directory: '{file_path}'")
-
-            with open(file_path, 'r', encoding='utf-8') as file:
-                content = file.read()
-            logging.info("Successfully loaded text file.")
-            return content
-
-        except FileNotFoundError as e:
-            error_msg = f"Error loading text file: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message
-        except UnicodeDecodeError as e:
-            error_msg = f"Encoding error reading file {file_path} as UTF-8: {e}"
-            logging.error(error_msg)
-            # Consider trying other encodings or returning raw bytes if appropriate
-            return error_msg # Return error message
-        except IOError as e:
-            error_msg = f"Error reading file {file_path}: {e}"
-            logging.error(error_msg)
-            return error_msg # Return error message
-
-
-# --- Agent Configuration ---
-
-# Define the custom prefix for the system prompt clearly
-SYSTEM_PROMPT_PREFIX = """You are a general AI assistant. I will ask you a question. 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 separators (e.g., 1000 instead of 1,000) and avoid units like $ or % unless explicitly requested.
-- If you are asked for a string, use standard capitalization, avoid abbreviations (e.g., Los Angeles instead of LA), and write out digits as words (e.g., five instead of 5) unless numbers are specifically requested. Avoid leading/trailing articles (a, an, the) if possible.
-- If you are asked for a comma-separated list, apply the above rules to each element based on whether it's a number or a string.
-"""
-
-def load_prompt_templates(yaml_path: str = "code_agent.yaml") -> dict:
-    """Loads prompt templates from a YAML file packaged with the library."""
-    try:
-        # Assumes 'smolagents.prompts' is a valid package/directory containing yaml_path
-        prompt_text = importlib.resources.files("smolagents.prompts").joinpath(yaml_path).read_text()
-        return yaml.safe_load(prompt_text)
-    except FileNotFoundError:
-        logging.error("Prompt YAML file not found at expected location: smolagents/prompts/%s", yaml_path)
-        # Return default empty dict or raise error, depending on desired behavior
-        return {}
-    except yaml.YAMLError as e:
-        logging.error("Error parsing YAML file %s: %s", yaml_path, e)
-        return {}
-    except Exception as e: # Catch other potential errors like package not found
-        logging.error("Failed to load prompts: %s", e)
-        return {}
-
-def init_agent(api_key: str | None = None,
-               model_id: str = "gemini-1.5-flash", # Updated model ID example
-               api_base: str = "https://generativelanguage.googleapis.com/v1beta", # Updated base URL
-               temperature: float = 0.7,
-               max_steps: int = 15) -> CodeAgent | None:
-    """
-    Initializes and configures the CodeAgent.
-
-    Args:
-        api_key: The API key for the generative model service. Reads from
-                 "API_KEY" environment variable if not provided.
-        model_id: The identifier of the model to use.
-        api_base: The base URL for the API. Note: The original URL seemed incorrect for Gemini via OpenAI proxy format. Check documentation.
-                  The example here uses the direct Gemini API base URL format. Adjust if using an OpenAI proxy.
-        temperature: The sampling temperature for the model.
-        max_steps: The maximum number of steps the agent can take.
-
-    Returns:
-        An initialized CodeAgent instance, or None if initialization fails.
-    """
-    # Prefer passed API key, fallback to environment variable
-    resolved_api_key = api_key or os.getenv("API_KEY")
-    if not resolved_api_key:
-        logging.error("API Key not provided and 'API_KEY' environment variable not set.")
-        return None
-
-    # Load base prompts
-    prompts = load_prompt_templates()
-    if not prompts or "system_prompt" not in prompts:
-        logging.error("Failed to load or parse base prompts. Cannot initialize agent.")
-        return None
-
-    # Prepend the custom instructions to the loaded system prompt
-    prompts["system_prompt"] = SYSTEM_PROMPT_PREFIX + prompts["system_prompt"]
-
-    # Define the model connection
-    # Note: Ensure OpenAIServerModel is compatible with the Gemini API structure
-    # or use a specific Gemini client library if available/preferred.
-    # The api_base URL format might need adjustment based on how OpenAIServerModel constructs the full URL.
-    try:
-        gemini_model = OpenAIServerModel(
-                model_id=model_id,
-                # Make sure api_base is correct for how OpenAIServerModel uses it.
-                # If it expects an OpenAI-like structure, you might need a proxy or adjust this URL.
-                # Example using direct Gemini API base:
-                api_base=api_base,
-                api_key=resolved_api_key,
-                temperature=temperature
-        )
-    except Exception as e:
-        logging.error("Failed to initialize the language model: %s", e)
-        return None
-
-
-    # Define the list of tools available to the agent
-    tools = [
-        DuckDuckGoSearchTool(),
-        VisitWebpageTool(),
-        WikipediaSearchTool(),
-        GetTaskFileTool(),    # Custom tool
-        SpeechToTextTool(),
-        LoadXlsxFileTool(),   # Custom tool
-        LoadTextFileTool()    # Custom tool
-    ]
-
-    # Create the agent instance
-    try:
-        agent = CodeAgent(
-            tools=tools,
-            model=gemini_model,
-            prompt_templates=prompts,
-            max_steps=max_steps,
-            # Explicitly list authorized imports for the code execution sandbox
-            additional_authorized_imports = ["pandas", "os.path"] # Added os.path for potential use
-        )
-        logging.info("CodeAgent initialized successfully.")
-        return agent
-    except Exception as e:
-        logging.error("Failed to initialize CodeAgent: %s", e)
-        return None
+prompts = yaml.safe_load(
+    importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
+)
+prompts["system_prompt"] = ("You are a general AI assistant. I will ask you a question. 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. "
+                            + prompts["system_prompt"])
+
+def init_agent():
+    gemini_model = OpenAIServerModel(
+            model_id="gemini-2.0-flash",
+            api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
+            api_key=os.getenv("API_KEY"),
+            temperature=0.7
+    )
+    agent = CodeAgent(
+        tools=[
+            DuckDuckGoSearchTool(),
+            VisitWebpageTool(),
+            WikipediaSearchTool(),
+            GetTaskFileTool(),
+            SpeechToTextTool(),
+            LoadXlsxFileTool(),
+            LoadTextFileTool()
+        ],
+        model=gemini_model,
+        prompt_templates=prompts,
+        max_steps=15,
+        additional_authorized_imports = ["pandas"]
+    )
+    return agent