| import os |
| import logging |
| import time |
| from typing import Optional |
| from google import genai |
|
|
| logger = logging.getLogger(__name__) |
|
|
| class LargeLLMGenerator: |
| """ |
| Wrapper for Google Gemini API to serve as an upper-bound baseline for fast, reliable generation. |
| """ |
| def __init__(self, config=None, backend: str = "api", api_model: str = "gemini-3.6-flash", local_model_name: str = None): |
| self.backend = backend |
| self.api_model = api_model |
| self.local_model_name = local_model_name |
| self.client = None |
| self.config = config |
| self.api_key_env_name = "GEMINI_API_KEY" |
|
|
| |
| if config and hasattr(config, "llm_baseline"): |
| llm_cfg = config.llm_baseline |
| self.backend = getattr(llm_cfg, "backend", self.backend) |
| self.api_model = getattr(llm_cfg, "api_model", self.api_model) |
| self.local_model_name = getattr(llm_cfg, "local_model", self.local_model_name) |
| self.api_key_env_name = getattr(llm_cfg, "api_key_env", None) or self.api_key_env_name |
|
|
| self._initialize_client() |
|
|
| def _initialize_client(self): |
| if self.backend == "api": |
| self._init_gemini() |
| else: |
| logger.info(f"Using local/off-the-shelf LLM backend: {self.local_model_name}") |
|
|
| def _init_gemini(self): |
| try: |
| api_key = os.getenv(self.api_key_env_name) |
|
|
| if not api_key: |
| logger.error( |
| f"API key not found. Set the '{self.api_key_env_name}' environment " |
| "variable before initializing LargeLLMGenerator." |
| ) |
| return |
|
|
| |
| self.client = genai.Client(api_key=api_key) |
| logger.info(f"✅ Initialized Gemini client for model: {self.api_model}") |
| except ImportError: |
| logger.error("google-genai package not installed. Run: pip install google-genai") |
| except Exception as e: |
| logger.error(f"Failed to initialize Gemini client: {e}") |
|
|
| def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.2) -> str: |
| if self.backend == "api": |
| return self._generate_api(prompt, max_tokens, temperature) |
| else: |
| logger.warning("Local LLM generation not fully implemented in this baseline wrapper. Returning empty string.") |
| return "" |
|
|
| def _generate_api(self, prompt: str, max_tokens: int, temperature: float) -> str: |
| if not self.client: |
| return "[API Client Not Initialized]" |
| try: |
| start_time = time.time() |
|
|
| |
| response = self.client.models.generate_content( |
| model=self.api_model, |
| contents=prompt, |
| config={ |
| "max_output_tokens": max_tokens, |
| "temperature": temperature |
| } |
| ) |
| text = response.text.strip() |
|
|
| latency = round((time.time() - start_time) * 1000, 2) |
| logger.info(f"Gemini API ({self.api_model}) generated {len(text)} chars in {latency}ms") |
| return text |
| except Exception as e: |
| logger.error(f"API Generation failed: {e}") |
| return f"[API Error: {str(e)}]" |
|
|
