Spaces:
Paused
Paused
| """ | |
| Google Gemini AI endpoint implementation. | |
| This module provides integration with Google's Gemini API for LLM inference. | |
| """ | |
| from google import genai | |
| from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError | |
| DEFAULT_MODEL = "gemini-2.0-flash-exp" | |
| class GeminiEndpoint(BaseAIEndpoint): | |
| """Google Gemini endpoint for cloud-based LLM inference.""" | |
| def _initialize_client(self) -> None: | |
| """Initialize the Gemini client.""" | |
| api_key = self.ai_config.get("api_key", "") | |
| if not api_key: | |
| raise AIEndpointRequestError("Gemini API key is required") | |
| # Default timeout of 30 seconds, configurable via ai_config | |
| timeout = self.ai_config.get("timeout", 30) | |
| self.client = genai.Client( | |
| api_key=api_key, | |
| http_options={'timeout': timeout} | |
| ) | |
| def _get_default_model(self) -> str: | |
| """Get the default Gemini model.""" | |
| return DEFAULT_MODEL | |
| def query(self, prompt: str, prompt_format: dict) -> str: | |
| """ | |
| Send a query to Gemini and return the response. | |
| Args: | |
| prompt: The prompt to send to the model | |
| Returns: | |
| The model's response as a string | |
| Raises: | |
| AIEndpointRequestError: If the request fails | |
| """ | |
| try: | |
| response = self.client.models.generate_content( | |
| model=self.model, | |
| contents=prompt, | |
| generation_config={ | |
| 'max_output_tokens': self.max_tokens, | |
| 'temperature': self.temperature, | |
| 'response_schema': prompt_format.model_json_schema(), | |
| } | |
| ) | |
| return response.text | |
| except Exception as e: | |
| raise AIEndpointRequestError(f"Gemini request failed: {e}") | |