src/agent/google_vertex_ai_agent.py

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