app/gemini_client.py

from __future__ import annotations
from google.genai import types 

import logging
import os
import threading
from typing import Any

logger = logging.getLogger(__name__)

_client: Any = None
_client_lock = threading.Lock()


def _api_key() -> str:
    return (
        os.environ.get("GEMINI_API_KEY", "").strip()
        or os.environ.get("GOOGLE_API_KEY", "").strip()
    )


def generation_backend() -> str:
    """Return ``gemini`` or ``local``. ``auto`` picks Gemini when an API key is set."""
    raw = os.environ.get("GENERATION_BACKEND", "auto").strip().lower()
    if raw in ("gemini", "google"):
        if not _api_key():
            raise RuntimeError(
                "GENERATION_BACKEND=gemini but GEMINI_API_KEY / GOOGLE_API_KEY is not set."
            )
        return "gemini"
    if raw in ("local", "hf", "huggingface"):
        return "local"
    return "gemini" if _api_key() else "local"


def use_gemini() -> bool:
    return generation_backend() == "gemini"


def gemini_model() -> str:
    return os.environ.get("GEMINI_MODEL", "gemini-2.0-flash").strip() or "gemini-2.0-flash"


def skip_local_llm_hub_download() -> bool:
    if os.environ.get("SKIP_LOCAL_LLM_HUB_DOWNLOAD", "").strip().lower() in (
        "1",
        "true",
        "yes",
    ):
        return True
    if os.environ.get("GENERATION_BACKEND", "auto").strip().lower() in ("gemini", "google"):
        return True
    return use_gemini()


def _get_client() -> Any:
    global _client
    if _client is not None:
        return _client
    with _client_lock:
        if _client is not None:
            return _client
        key = _api_key()
        if not key:
            raise RuntimeError(
                "Gemini API key missing. Set GEMINI_API_KEY or GOOGLE_API_KEY "
                "(or GENERATION_BACKEND=local)."
            )
        from google import genai  # type: ignore[import-untyped]

        _client = genai.Client(api_key=key)
        logger.info("Gemini client ready (model=%s)", gemini_model())
        return _client


def gemini_generate_text(
    *,
    system_instruction: str,
    user_text: str,
    temperature: float = 0.0,
    max_output_tokens: int = 4096,  
    response_schema: Any | None = None,
) -> str:
    from google.genai import types  # type: ignore[import-untyped]

    client = _get_client()
    resp = client.models.generate_content(
        model=gemini_model(),
        contents=user_text,
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
            temperature=temperature,
            max_output_tokens=max_output_tokens,
            response_mime_type="application/json",
            response_schema=response_schema,
            
            thinking_config=types.ThinkingConfig(
                thinking_level="minimal"  # Options: 'minimal' or 'low' to free up response tokens
            )
        ),
    )
    text = (resp.text or "").strip()
    if not text:
        raise RuntimeError("Gemini returned empty text")
    return text


def gemini_generate_chat(
    messages: list[dict[str, str]],
    *,
    temperature: float = 0.2,
    max_output_tokens: int = 1024,
) -> str:
    from google.genai import types  # type: ignore[import-untyped]

    system_parts: list[str] = []
    contents: list[Any] = []
    for m in messages:
        role = m.get("role", "user")
        text = m.get("content", "")
        if role == "system":
            system_parts.append(text)
            continue
        gem_role = "model" if role == "assistant" else "user"
        contents.append(
            types.Content(role=gem_role, parts=[types.Part.from_text(text=text)])
        )

    if not contents:
        raise ValueError("No user/model messages for Gemini")

    config_kw: dict[str, Any] = {
        "temperature": temperature,
        "max_output_tokens": max_output_tokens,
    }
    if system_parts:
        config_kw["system_instruction"] = "\n\n".join(system_parts)

    client = _get_client()
    resp = client.models.generate_content(
        model=gemini_model(),
        contents=contents,
        config=types.GenerateContentConfig(**config_kw),
    )
    text = (resp.text or "").strip()
    if not text:
        raise RuntimeError("Gemini returned empty text")
    return text