File size: 3,397 Bytes
3581a58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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"
# Override with config if provided
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
# Initialize the modern google-genai client
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()
# Generate content using the google-genai client
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)}]"
|