Spaces:
Sleeping
Sleeping
| from smolagents import ( | |
| CodeAgent, | |
| LiteLLMModel, | |
| Tool, | |
| ) | |
| from typing import Callable | |
| class MyAgent: | |
| def __init__( | |
| self, | |
| provider: str = "litellm", | |
| model_id: str = "gemini/gemini-2.0-flash-lite", | |
| api_base: str | None = None, | |
| api_key: str | None = None, | |
| planning_interval: int = 3, | |
| num_ctx: int = 8192, | |
| tools: list[Tool] = [], | |
| add_base_tools: bool = True, | |
| temperature: float = 0.2, | |
| additional_authorized_imports: list[str] = [], | |
| step_callbacks: list[Callable] = [], | |
| max_steps: int = 20, | |
| verbosity_level: int = 2, | |
| ): | |
| """ | |
| Initializes the agent depending on the provider and model ID. | |
| Args: | |
| provider (str): The provider of the model (e.g., "litellm", "huggingface"). | |
| model_id (str): The ID of the model to be used. | |
| tools (list[Tool]): The tools to be used by the agent. | |
| api_base (str | None): The base URL of the API. | |
| api_key (str | None): The API key. | |
| planning_interval (int): The interval for planning. | |
| num_ctx (int): The number of context tokens. | |
| add_base_tools (bool): Whether to add base tools. | |
| temperature (float): The temperature for the model. | |
| additional_authorized_imports (list[str]): The additional authorized imports. | |
| step_callbacks (list[Callable]): The step callbacks. | |
| max_steps (int): The maximum steps. | |
| verbosity_level (int): The verbosity level. | |
| Returns: | |
| None: None | |
| """ | |
| self.provider = provider | |
| self.model_id = model_id | |
| self.api_base = api_base | |
| self.api_key = api_key | |
| self.planning_interval = planning_interval | |
| self.num_ctx = num_ctx | |
| self.temperature = temperature | |
| model = LiteLLMModel( | |
| model_id=self.model_id, | |
| api_base=self.api_base, | |
| api_key=self.api_key, | |
| num_ctx=self.num_ctx, | |
| add_base_tools=add_base_tools, | |
| temperature=self.temperature, | |
| ) | |
| # Initialize the agent with the specified provider and model ID | |
| if provider == "litellm": | |
| self.agent = CodeAgent( | |
| model=model, | |
| tools=tools, | |
| planning_interval=self.planning_interval, | |
| additional_authorized_imports=additional_authorized_imports, | |
| step_callbacks=step_callbacks, | |
| max_steps=max_steps, | |
| verbosity_level=verbosity_level, | |
| ) | |
| else: | |
| raise ValueError(f"Unsupported provider: {provider}") | |
| print(f"Agent initialized with provider: {provider}, model ID: {model_id}") | |
| def __call__(self, question: str) -> str: | |
| """ | |
| Given a question, run the agent and return the answer. | |
| Args: | |
| question (str): The question to be answered. | |
| Returns: | |
| str: The answer to the question. | |
| """ | |
| final_answer = self.agent.run(question) | |
| print(f"Agent received question (last 50 chars): {question[-50:]}...") | |
| print(f"Agent returning fixed answer: {final_answer}") | |
| return final_answer | |