Update agent.py

agent.py

@@ -1,79 +1,297 @@
-import importlib
+# -*- coding: utf-8 -*-
+"""
+Initializes and configures a SmolAgents CodeAgent with custom tools
+for file handling and web interaction.
+"""
+
+import importlib.resources
 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
 
-from config import DEFAULT_API_URL
-from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, WikipediaSearchTool, Tool, OpenAIServerModel, SpeechToTextTool
+# 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 ---
 
 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 = """This tool downloads the file content associated with the given task_id if exists. Returns absolute file path"""
+    description = "Downloads the file content associated with the given task_id if it exists. Returns the absolute file path of the downloaded file."
     inputs = {
-        "task_id": {"type": "string", "description": "Task id"},
-        "file_name": {"type": "string", "description": "File name"},
+        "task_id": {"type": "string", "description": "The unique identifier for the task."},
+        "file_name": {"type": "string", "description": "The desired local name for the downloaded file."},
     }
-    output_type = "string"
+    output_type = "string" # Output is the file path or an error message
 
     def forward(self, task_id: str, file_name: str) -> str:
-        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)
+        """
+        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
+
 
 class LoadXlsxFileTool(Tool):
+    """
+    A tool to load data from an XLSX (Excel) file into a pandas DataFrame.
+    """
     name = "load_xlsx_file_tool"
-    description = """This tool loads xlsx file into pandas and returns it"""
+    description = "Loads data from an XLSX file specified by its path into a pandas DataFrame."
     inputs = {
-        "file_path": {"type": "string", "description": "File path"}
+        "file_path": {"type": "string", "description": "The local path to the XLSX file."}
     }
+    # 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) -> object:
-        return pd.read_excel(file_path)
+    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
+
 
 class LoadTextFileTool(Tool):
+    """
+    A tool to load the content of a text file into a single string.
+    """
     name = "load_text_file_tool"
-    description = """This tool loads any text file"""
+    description = "Loads the entire content of any text file specified by its path."
     inputs = {
-        "file_path": {"type": "string", "description": "File path"}
+        "file_path": {"type": "string", "description": "The local path to the text file."}
     }
-    output_type = "string"
+    output_type = "string" # Output is the file content or an error message
 
-    def forward(self, file_path: str) -> object:
-        with open(file_path, 'r', encoding='utf-8') as file:
-            return file.read()
+    def forward(self, file_path: str) -> str:
+        """
+        Executes the text file loading process.
 
+        Args:
+            file_path: The path to the text file.
 
-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
+        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