agent.py

from pathlib import Path
import litellm
from smolagents import (
    CodeAgent,
    WebSearchTool,
    VisitWebpageTool,
    WikipediaSearchTool,
    PythonInterpreterTool,
    FinalAnswerTool,
    LiteLLMModel,
)
from smolagents.monitoring import LogLevel
from smolagents import tool

import os

class SubAgent:
    def __init__(self):
        """
        Initialize the SubAgent with a given agent.
        """

        self._model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash",  # you can see other model names here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models. It is important to prefix the name with "gemini/"
            api_key=os.getenv("GEMINI_API_KEY"),
            max_tokens=8192,
            temperature=0.2,
        )

class ReadingAgent(SubAgent):
    def __init__(
        self,
    ):
        super().__init__()
        self._r_agent = CodeAgent(
            model=self._model,
            tools=[read_excel_file, read_python_file],
            name="reading_agent",
            description="Reads the contents of files and extracts the contents from them. This agent should not be used to execute code or make requests to URLs.",
            verbosity_level=LogLevel.INFO,
            max_steps=1,
        )

    @property
    def agent(self):
        """
        Property that returns the CodeAgent instance
        """
        return self._r_agent


class DescribeAgent(SubAgent):
    def __init__(
        self,
    ):
        super().__init__()
        self._r_agent = CodeAgent(
            model=self._model,
            tools=[describe_audio, describe_image],
            name="describe_agent",
            description="Transcribe extracts information from files and provides information about the contents of the file provided. As an example this might transcribe an audio file or describe an image. This agent should not be used to execute code or make requests to URLs.",
            verbosity_level=LogLevel.INFO,
            max_steps=1,
        )

    @property
    def agent(self):
        """
        Property that returns the CodeAgent instance
        """
        return self._r_agent


class CodingAgent(SubAgent):
    def __init__(self):
        super().__init__()
        self._c_agent = CodeAgent(
            model=self._model,
            tools=[PythonInterpreterTool()],
            name="coding_agent",
            description="Executes Python code to solve tasks, this agent should not be used to make external calls e.g. do not make any requests to URLs. This agent should just be used to process information retrieved from other tools and agents locally",
            verbosity_level=LogLevel.INFO,
            max_steps=5,
            # add_base_tools=True,
            additional_authorized_imports=[
                "pandas",
                "re",
                "requests",
                "json",
                "numpy",
                "bs4",
                "datetime",
                "os",
                "io",
                "csv",
            ],
        )

    @property
    def agent(self):
        """
        Property that returns the CodeAgent instance
        """
        return self._c_agent


class WebAgent(SubAgent):
    def __init__(self):
        super().__init__()
        self._w_agent = CodeAgent(
            model=self._model,
            tools=[WebSearchTool(), VisitWebpageTool(), WikipediaSearchTool()],
            name="web_agent",
            description="Browses the web to find information, to find information about a subject Wikipedia is the best source for factual information. For following direct links, use the VisitWebpageTool. For general searches, use the WebSearchTool. This agent should not be used to execute code or make requests to URLs.",
            verbosity_level=LogLevel.INFO,
            max_steps=3,
        )

    @property
    def agent(self):
        """
        Property that returns the CodeAgent instance
        """
        return self._w_agent


class BudleeAgent(SubAgent):

    # Constructor for SMOL Agent

    _PROMPT = """
    "You are a general AI assistant. "
    "I will ask you a question. Report your thoughts step by step. "
    "Finish your answer only with the final answer. In the final answer don't write explanations. "
    "The answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. "
    "Avoid units, abbreviations, or articles unless specified. "
    "Pay attention to each sentence in the question and verify the answer against every part. "
    "If the question requires a specific format (e.g., a number, a name, 'yes/no'), provide the answer in that format. "
    "If the answer is a numerical value, return only the number itself. "
    "If the detailed response does not clearly answer the question, output 'Unable to determine a concise answer from the provided information.' "
    "Focus solely on the factual answer."
    ""
    "QUESTION: "
    ""
    "{question}"
    """

    def __init__(self):
        """
        Initialize the BudleeAgent.
        This is where you can set up any necessary configurations or parameters.
        """
        # self._api_url = api_base_url
        # model = InferenceClientModel(
        #     max_tokens=4096,
        #     temperature=0.2,
        #     model_id="Qwen/Qwen2.5-Coder-32B-Instruct",  # it is possible that this model may be overloaded
        #     custom_role_conversions=None,
        #     provider="hf-inference",
        # )
        super().__init__()
        agent_web = WebAgent().agent
        agent_coding = CodingAgent().agent
        agent_reading = ReadingAgent().agent
        agent_describing = DescribeAgent().agent
        self._manager_agent = CodeAgent(
            model=self._model,
            managed_agents=[agent_web, agent_coding, agent_reading, agent_describing],
            tools=[
                FinalAnswerTool(),
            ],
            planning_interval=5,
            verbosity_level=LogLevel.DEBUG,
            # final_answer_checks=[self._check_response],
            max_steps=5,
        )

        self._manager_agent.visualize()

        # def _check_response(self, final_answer, agent_memory):
        #     """
        #     Check the final answer response for validity.
        #     This is a placeholder; implement your own checks as needed.
        #     """
        #     # Example: always return True (no check)
        #     return True

    def answer(self, question: str) -> str:
        """
        Try to solve the task as an Agent using the information
        """
        print(f"Agent received question (first 50 chars): {question[:50]}...")

        prompt = self._PROMPT.format(
            question=question,
            #   additional_context=additional_context
        )
        print(f"Agent prompt {prompt}")

        answer = self._manager_agent.run(prompt)
      
        print(f"Agent returning fixed answer: {answer}")

        return answer  # type: ignore

@tool
def read_excel_file(file_path: str) -> str:
    """Reads an Excel file and returns its content as a string.
    Args:
        file_path: The path to the Excel file.
    Returns:
        The content of the Excel file as a string, or an error message if the file cannot be read.
    """
    try:
        import pandas as pd

        # Read the Excel file
        df = pd.read_excel(file_path)

        # Convert the DataFrame to a string representation
        return df.to_string()

    except Exception as e:
        return f"Error reading the Excel file: {str(e)}"


@tool
def read_python_file(file_path: str) -> str:
    """Reads a Python file and returns its content as a string.
    Args:
        file_path: The path to the Python file.
    Returns:
        The content of the Python file as a string, or an error message if the file cannot be read.
    """
    try:
        with open(file_path, "r") as file:
            content = file.read()
        return content
    except Exception as e:
        return f"Error reading the Python file: {str(e)}"


import base64


def convert_image_to_base64(image_path: str) -> str:
    """
    Converts an image to a Base64 string.
    Args:
        image_path: The path to the image file.
    Returns:
        A Base64 encoded string of the image.
    """
    try:
        with open(image_path, "rb") as image_file:
            # Read the image file as binary data
            encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
        return encoded_string
    except Exception as e:
        return f"Error converting image to Base64: {str(e)}"


@tool
def describe_image(image_path: str) -> str:
    """
    Return the description of a given image.
    Args:
        image_path: the input path of the image to describe.
    """

    encoded_data = convert_image_to_base64(image_path)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Describe the image in detail."},
                {
                    "type": "file",
                    "file": {
                        "file_data": f"data:image/png;base64,{encoded_data}",  # 👈 SET MIME_TYPE + DATA
                    },
                },
            ],
        }
    ]

    # Make the API call to Gemini model
    response = litellm.completion(
        model="gemini/gemini-2.0-flash-lite",
        messages=messages,
    )

    # Extract the response content
    content = response.get("choices", [{}])[0].get("message", {}).get("content")

    return content


@tool
def describe_audio(audio_path: str) -> str:
    """
    Return the transcription of a given audio file.
    Args:
        audio_path: the input path of the audio to transcribe.
    """

    audio_bytes = Path(audio_path).read_bytes()
    encoded_data = base64.b64encode(audio_bytes).decode("utf-8")

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "Please transcribe the content of this audio.",
                },
                {
                    "type": "file",
                    "file": {
                        "file_data": "data:audio/mp3;base64,{}".format(
                            encoded_data
                        ),  # 👈 SET MIME_TYPE + DATA
                    },
                },
            ],
        }
    ]

    # Make the API call to Gemini model
    response = litellm.completion(
        model="gemini/gemini-2.0-flash-lite",
        messages=messages,
    )

    # Extract the response content
    content = response.get("choices", [{}])[0].get("message", {}).get("content")

    return content