agent.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel # HfApiModel, OpenAIServerModel
import asyncio

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",
            api_key=os.getenv("GEMINI_KEY"),
            max_tokens=8192
#        model = HfApiModel(
#            max_tokens=2096,
#            temperature=0.5,
#            model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
#            custom_role_conversions=None,
        )
        self.agent = CodeAgent(
            model= model,  # OpenAIServerModel(model_id="gpt-4o"), 
            tools=[
                DuckDuckGoSearchTool(),
                WikipediaSearchTool()
            ]
        )
        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:
            # Run the agent asynchronously
            response = await asyncio.to_thread(
                self.agent.run,
                task=f"Answer the following question accurately and concisely: {question}"
            )
            print(f"MagAgent response: {response[:50]}...")
            return response
        except Exception as e:
            error_msg = f"Error processing question: {e}"
            print(error_msg)
            return error_msg