agent.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel, tool, load_tool # HfApiModel, OpenAIServerModel
import asyncio
import os
import yaml
from PIL import Image
import requests
from io import BytesIO

# Simulated additional tools (implementation depends on external APIs or setup)
#@tool
#def GoogleSearchTool(query: str) -> str:
#    """Tool for performing Google searches using Custom Search JSON API
#    Args:
#        query (str): Search query string
#    Returns:
#        str: Formatted search results
#    """
#    cse_id = os.environ.get("GOOGLE_CSE_ID")
#    if not api_key or not cse_id:
#        raise ValueError("GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables.")
#    url = "https://www.googleapis.com/customsearch/v1"
#    params = {
#        "key": api_key,
#        "cx": cse_id,
#        "q": query,
#        "num": 5  # Number of results to return
#    }
#    try:
#        response = requests.get(url, params=params)
#        response.raise_for_status()
#        results = response.json().get("items", [])
#        return "\n".join([f"{item['title']}: {item['link']}" for item in results]) or "No results found."
#    except Exception as e:
#        return f"Error performing Google search: {str(e)}"
        
@tool
def ImageAnalysisTool(image_path: str) -> str:
    """Tool for analyzing images using computer vision
    Args:
        image_path (str): Path to image file
    Returns:
        str: Image description
    """

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36" 
    }
    response = requests.get(image_path,headers=headers)
    image = Image.open(BytesIO(response.content)).convert("RGB")

    model = LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key= os.environ.get("GEMINI_KEY"),
        max_tokens=8192
    )

    agent = CodeAgent(
        tools=[],
        model=model,
        max_steps=20,
        verbosity_level=2
    )

    response = agent.run(
        """
        Describe what you see in the image in details.
        """,
        images=image
    )
    
    return f"The image description: '{response}'"


#@tool
#class LocalFileAudioTool:
#    """Tool for transcribing audio files"""
#    
#    @tool
#    def transcribe(self, file_path: str) -> str:
#        """Transcribe audio from file
#        Args:
#            file_path (str): Path to audio file
#        Returns:
#            str: Transcription text
#        """
#        return f"Transcribed audio from '{file_path}' (simulated)."

class MagAgent:
    def __init__(self):
        """Initialize the MagAgent with search tools."""
        print("Initializing MagAgent with search tools...")
        model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash-lite",
            api_key= os.environ.get("GEMINI_KEY"),
            max_tokens=8192
        )

        # Load prompt templates
        with open("prompts.yaml", 'r') as stream:
            prompt_templates = yaml.safe_load(stream)
        
        self.agent = CodeAgent(
            model= model,
            tools=[
#                GoogleSearchTool,
                DuckDuckGoSearchTool(),
                WikipediaSearchTool(),
                ImageAnalysisTool,
#                LocalFileAudioTool()
            ]
        )
        print("MagAgent initialized.")

    async def __call__(self, question: str) -> str:
        """Process a question asynchronously using the MagAgent."""
        print(f"MagAgent received question (first 50 chars): {question[:50]}...")
        try:
            # Define a task with fallback search logic
            task = (
                f"Answer the following question accurately and concisely: {question}\n"
            )
            response = await asyncio.to_thread(
                self.agent.run,
                task=task
            )

            # Ensure response is a string, fixing the integer error
            response = str(response) if response is not None else "No answer found."
       
            if not response or "No Wikipedia page found" in response:
                # Fallback response if search fails
                response = "Unable to retrieve exact data. Please refine the question or check external sources."
            print(f"MagAgent response: {response[:50]}...")
            return response
        except Exception as e:
            error_msg = f"Error processing question: {str(e)}. Check API key or network connectivity."
            print(error_msg)
            return error_msg