agent.py

import os

import PIL.Image
from dotenv import load_dotenv
from loguru import logger
from smolagents import AzureOpenAIServerModel, CodeAgent

from src.file_handler.parse import parse_file

load_dotenv()


class Agent:
    def __init__(self):
        model = AzureOpenAIServerModel(
            model_id=os.getenv("AZURE_OPENAI_MODEL_ID"),
            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
            api_version=os.getenv("OPENAI_API_VERSION"),
        )
        self.agent = CodeAgent(
            tools=[],
            model=model,
            add_base_tools=True,  # Add any additional base tools
            # planning_interval=3,
        )
        logger.info("BasicAgent initialized.")

    def __call__(self, question: str, file_name: str) -> str:
        logger.info(
            f"Agent received question (first 50 chars): {question[:50]}..."
        )
        images = None

        if file_name:
            content = parse_file(task_id, file_name, api_url)
            if content:
                if isinstance(
                    content, PIL.Image.Image
                ):  # Parse content as image
                    images = [content]
                else:  # Append content to question
                    question += f"\n\nAttached content:\n{content}"
                    logger.info(f"Question with content: {question}")

        answer = self.agent.run(question, images=images)
        logger.info(f"Agent returning answer: {answer}")
        return answer


if __name__ == "__main__":
    import requests

    api_url = "https://agents-course-unit4-scoring.hf.space"
    question_url = f"{api_url}/random-question"

    data = requests.get(question_url).json()
    agent = Agent()

    task_id = data["task_id"]
    question = data["question"]
    file_name = data["file_name"]
    logger.info(
        f"Task ID: {task_id}\nQuestion: {question}\nFile Name: {file_name}\n\n"
    )

    answer = agent(question, file_name)