File size: 1,620 Bytes
d1568ce
9e0ec52
 
 
 
3cf8730
 
d1568ce
 
 
 
 
027bbe7
 
 
 
89d512b
9e0ec52
7f24a38
9e0ec52
 
 
 
 
89d512b
9e0ec52
 
3cf8730
 
9e0ec52
 
 
 
975bd5d
9e0ec52
3cf8730
9e0ec52
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
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