File size: 669 Bytes
37c9a6b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
import os
from smolagents import Tool, LiteLLMModel
from agent.base_agent import BaseAgent
class OpenAiAgent(BaseAgent):
def __init__(self, model_name: str = "gemini/gemini-2.0-flash-lite", tools: list[Tool] | None = None,
use_all_custom_tools: bool = True):
super().__init__(model_name=model_name, tools=tools, use_all_custom_tools=use_all_custom_tools)
self.model_name: str = model_name
self.agent = self.init_agent()
def get_model(self):
model = LiteLLMModel(
model_id=self.model_name,
temperature=0.2,
api_key=os.getenv("GEMINI_API_KEY")
)
return model
|