DavidL72Code/UMB_Sustainable_Chatbot

1"""The Gemini client: safety settings, generation config, and the two call paths.2 3Split out of Chatbot.py unchanged. Depends only on config (for settings) and4telemetry (to record each call), both of which are leaves, so this module can be5imported from anywhere in the pipeline without a cycle.6 7_caller_stage walks the interpreter call stack to label a call by the pipeline8function that issued it. Moving the definition into this module does not add a9frame, so the labels it produces are the same.10"""11from __future__ import annotations12 13import os14import time15from typing import Optional16 17try:18    from google import genai19    from google.genai import types as genai_types20except ImportError:  # pragma: no cover - dependency availability depends on the runtime21    genai = None22    genai_types = None23 24from config import ChatbotConfig25from telemetry import _caller_stage, record_llm_call26 27 28_gemini_client: Optional[object] = None29 30 31def _get_gemini_client():32    global _gemini_client33    if _gemini_client is None:34        if genai is None:35            raise ImportError("Install google-genai to use Gemini.")36        cfg = ChatbotConfig()37        if not cfg.gemini_api_key:38            raise ValueError("Set GEMINI_API_KEY before using Gemini.")39        _gemini_client = genai.Client(api_key=cfg.gemini_api_key)40    return _gemini_client41 42 43_DEFAULT_SAFETY_SETTINGS = [44    genai_types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_MEDIUM_AND_ABOVE"),45    genai_types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_MEDIUM_AND_ABOVE"),46    genai_types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_MEDIUM_AND_ABOVE"),47    genai_types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_MEDIUM_AND_ABOVE"),48] if genai_types is not None else None49 50 51 52# Reproducibility: temperature alone does not make decoding greedy. With53# top_p=0.95/top_k=40 the model still samples from a 40-token distribution, so54# the same question at temperature 0 returned four different answers in four55# runs. At temperature 0 pin top_k=1 and top_p=1.0 (greedy) and send a fixed56# seed, so a run can be reproduced and a change can be told apart from noise.57GEMINI_SEED = int(os.getenv("GEMINI_SEED", "7"))58 59 60# Models that reject a thinking_config outright (Gemma, for one). Populated the61# first time a model 400s on it, so the probe costs one failed call per process62# rather than one per question.63_MODELS_WITHOUT_THINKING: set[str] = set()64 65 66def _rejects_thinking(exc: Exception) -> bool:67    message = str(exc)68    return "INVALID_ARGUMENT" in message and "hinking" in message69 70 71_WARNED_ONCE: set[str] = set()72 73 74def _warn_once(message: str) -> None:75    """Log a recurring failure once per process so it cannot hide, without76    printing the same line on every request."""77    if message in _WARNED_ONCE:78        return79    _WARNED_ONCE.add(message)80    print(f"[warn] {message}", file=sys.stderr, flush=True)81 82 83def _gemini_gen_config(84    temperature: float,85    thinking_budget: int = 1024,86    include_thinking: bool = True,87) -> "genai_types.GenerateContentConfig":88    deterministic = float(temperature or 0.0) <= 0.089    return genai_types.GenerateContentConfig(90        temperature=temperature,91        top_p=1.0 if deterministic else 0.95,92        top_k=1 if deterministic else 40,93        seed=GEMINI_SEED,94        # Thinking tokens count against this budget, so a 1024-token thinking95        # pass left ~1024 for the answer. The eval judge writes a JSON object96        # with a prose "notes" field and ran out mid-key ('"right_c'), which97        # cost n146 a verdict on an answer it had already scored 5/5. Raising98        # the ceiling cannot change a response that already fit.99        max_output_tokens=int(os.getenv("GEMINI_MAX_OUTPUT_TOKENS", "4096")),100        thinking_config=(101            genai_types.ThinkingConfig(thinking_budget=thinking_budget) if include_thinking else None102        ),103        safety_settings=_DEFAULT_SAFETY_SETTINGS,104    )105 106 107def call_gemini(prompt: str, model: Optional[str] = None, temperature: Optional[float] = None, thinking_budget: int = 1024) -> str:108    cfg = ChatbotConfig()109    client = _get_gemini_client()110    model_name = model or cfg.gemini_model111    temp = temperature if temperature is not None else cfg.gemini_temperature112    stage = _caller_stage()113    started_at = time.perf_counter()114    supports_thinking = model_name not in _MODELS_WITHOUT_THINKING115    try:116        response = client.models.generate_content(117            model=model_name,118            contents=prompt,119            config=_gemini_gen_config(120                temp, thinking_budget=thinking_budget, include_thinking=supports_thinking121            ),122        )123    except Exception as exc:124        # Some models reject thinking_budget=0 and some reject a thinking_config125        # at all — Gemma raises "Thinking budget is not supported for this126        # model." Retrying with the *default* budget only helps the first case;127        # the second needs the field dropped entirely. Getting this wrong meant128        # every planner call failed silently and no question was ever split into129        # facets, so multi-part questions only ever answered their first half.130        if not (supports_thinking and _rejects_thinking(exc)):131            raise132        _MODELS_WITHOUT_THINKING.add(model_name)133        response = client.models.generate_content(134            model=model_name,135            contents=prompt,136            config=_gemini_gen_config(temp, include_thinking=False),137        )138    record_llm_call(139        model=model_name,140        stage=stage,141        usage=getattr(response, "usage_metadata", None),142        latency_ms=(time.perf_counter() - started_at) * 1000,143        streamed=False,144    )145    return response.text.strip()146 147 148def call_gemini_stream(prompt: str, model: Optional[str] = None, temperature: Optional[float] = None, stage: Optional[str] = None):149    """Yields text chunks as they stream from the Gemini API."""150    cfg = ChatbotConfig()151    client = _get_gemini_client()152    model_name = model or cfg.gemini_model153    temp = temperature if temperature is not None else cfg.gemini_temperature154    stage = stage or _caller_stage(default="generation")155    started_at = time.perf_counter()156    usage = None157    for chunk in client.models.generate_content_stream(158        model=model_name,159        contents=prompt,160        config=_gemini_gen_config(temp),161    ):162        chunk_usage = getattr(chunk, "usage_metadata", None)163        if chunk_usage is not None:164            usage = chunk_usage165        if chunk.text:166            yield chunk.text167    record_llm_call(168        model=model_name,169        stage=stage,170        usage=usage,171        latency_ms=(time.perf_counter() - started_at) * 1000,172        streamed=True,173    )174