from .base import BaseModel
from google import genai

class GoogleGeminiModel(BaseModel):
    def __init__(self, model_name: str, api_key: str):
        super().__init__()
        self.model_name = model_name
        self.client = genai.Client(api_key=api_key)

    def get_response(self, input_data: str) -> str:
        """
        Gets a response from the Google Gemini model.
        
        Args:
            input_data (str): The input data for the model.

        Returns:
            str: The response generated by the model.
        """
        try:
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=input_data,
            )
            return response.text if response and hasattr(response, "text") else "No response generated."
        except Exception as e:
            return f"Error: {str(e)}"
    