Spaces:
Sleeping
Sleeping
File size: 1,736 Bytes
cfd75dc 0845b6c cfd75dc 0845b6c cfd75dc |
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 |
"""Utility functions for working with the language model."""
import logging
from langchain_google_genai import ChatGoogleGenerativeAI
from config import settings
logger = logging.getLogger(__name__)
_API_KEYS: list[str] = []
_current_key_idx = 0
MODEL_NAME = "gemini-2.5-flash-preview-05-20"
def _get_api_key() -> str:
"""Return an API key using round-robin selection."""
global _API_KEYS, _current_key_idx
if not _API_KEYS:
keys_str = settings.gemini_api_key.get_secret_value()
if keys_str:
_API_KEYS = [k.strip() for k in keys_str.split(",") if k.strip()]
if not _API_KEYS:
msg = "Google API keys are not configured or invalid"
logger.error(msg)
raise ValueError(msg)
key = _API_KEYS[_current_key_idx]
_current_key_idx = (_current_key_idx + 1) % len(_API_KEYS)
logger.debug("Using Google API key index %s", _current_key_idx)
return key
def create_llm(
temperature: float = settings.temperature,
top_p: float = settings.top_p,
) -> ChatGoogleGenerativeAI:
"""Create a standard LLM instance."""
return ChatGoogleGenerativeAI(
model=MODEL_NAME,
google_api_key=_get_api_key(),
temperature=temperature,
top_p=top_p,
thinking_budget=1024,
)
def create_light_llm(temperature: float = settings.temperature, top_p: float = settings.top_p):
return ChatGoogleGenerativeAI(
model="gemini-2.0-flash",
google_api_key=_get_api_key(),
temperature=temperature,
top_p=top_p
)
def create_precise_llm() -> ChatGoogleGenerativeAI:
"""Return an LLM tuned for deterministic output."""
return create_llm(temperature=0, top_p=1)
|