agent.py

from typing import Any, List, Optional

from smolagents import CodeAgent, InferenceClientModel
from utils.logger import get_logger

logger = get_logger(__name__)


class Agent:
    """
    Agent class that wraps a CodeAgent and provides a callable interface for answering questions.
    Args:
        tools (Optional[List[Any]]): List of tools to provide to the agent.
        prompt (Optional[str]): Custom prompt template for the agent.
    """

    def __init__(
        self,
        tools: Optional[List[Any]] = None,
        prompt: Optional[str] = None,
    ):
        logger.info("Initializing Agent")
        # Verify correct parameter name for the inference client (e.g. api_provider)
        self.model = InferenceClientModel(
            model_id="deepseek-ai/DeepSeek-R1",
            api_provider="together"  # Changed parameter to api_provider
        )
        self.tools = tools
        self.imports = [
            "pandas",
            "numpy",
            "os",
            "requests",
            "tempfile",
            "datetime",
            "json",
            "time",
            "re",
            "openpyxl",
            "pathlib",
        ]
        self.agent = CodeAgent(
            model=self.model,
            tools=self.tools,
            add_base_tools=True,
            additional_authorized_imports=self.imports,
        )
        # Improved prompt with clearer context handling instructions
        self.prompt = prompt or (
            """You are an advanced AI assistant specialized in solving complex tasks using tools.

            Key Instructions:
            1. ALWAYS use tools for file paths in the context:
               - Use read_file for text files
               - Use extract_text_from_image for images (.png, .jpg)
            2. Final answer must ONLY contain the requested result in the EXACT format needed.
            
            QUESTION: {question}
            {context}
            ANSWER:"""
        )
        logger.info("Agent initialized")

    def __call__(self, question: str, file_path: Optional[str] = None) -> str:
        # Handle context more cleanly
        context = f"CONTEXT: File available at: {file_path}" if file_path else ""
        answer = self.agent.run(
            self.prompt.format(question=question, context=context)
        )
        return str(answer).strip("'").strip('"').strip()