system-prompts-and-models-of-ai-tools/task/llm.py

"""
Unified LLM interface — Gemini-only build for the Carta onboarding service.

This service runs exclusively on Google Gemini (google-genai SDK, API-key auth
via the GCP_API_KEY env var). The AWS Bedrock and GCP Vertex backends that the
shared lambdaLLM module supports have been removed here: there is a single
provider and a single code path.

All call sites import from this module so the provider stays centralized:

    from llm import get_chat_model      # chat model with .invoke()
    from llm import get_llm_client      # low-level client with .invoke_model()
    from llm import model_call, invoke_messages, get_converse_client

Retry strategy: on rate-limit / quota errors, wait a fixed interval until the
provider's limit window resets and retry until the call succeeds (no
exponential backoff). Non-retryable errors are raised immediately.
"""
import logging
import os
import time

logger = logging.getLogger(__name__)

# Backend is fixed to Gemini. Exposed as a constant (and written back to the
# environment) so any code or dependency that still reads LLM_BACKEND sees the
# only supported value.
LLM_BACKEND = "gemini"
os.environ["LLM_BACKEND"] = LLM_BACKEND

# Model id is config-driven so ops can flip Gemini models without code changes.
GEMINI_MODEL_ID = os.environ.get("GEMINI_MODEL_ID", "gemini-3.5-flash")
ACTIVE_MODEL_ID = GEMINI_MODEL_ID

logger.info("[llm] backend=gemini | active_model=%s", ACTIVE_MODEL_ID)
print(f"[llm] backend=gemini | active_model={ACTIVE_MODEL_ID}", flush=True)

# Retry config. LLM_RETRY_WAIT is the fixed wait (seconds) between retries;
# LLM_MAX_RETRIES is a safety cap (default 1000 ~= 16h at 60s wait).
LLM_RETRY_WAIT = float(os.environ.get("LLM_RETRY_WAIT", "60.0"))
LLM_MAX_RETRIES = int(os.environ.get("LLM_MAX_RETRIES", "1000"))


def _is_rate_limit(exc):
    """Return True for throttling / rate-limit / quota errors from google-genai."""
    msg = str(exc).lower()
    if "google.genai" in type(exc).__module__:
        code = getattr(exc, "code", None)
        if isinstance(code, int) and code in {408, 429, 500, 502, 503, 504}:
            return True
    for token in ("throttl", "rate limit", "too many requests", "quota",
                  "overloaded", "resource_exhausted", "429", "503",
                  "service unavailable"):
        if token in msg:
            return True
    return False


def _with_retry(fn, *args, **kwargs):
    """Invoke fn and wait out provider rate-limit windows until it succeeds."""
    for attempt in range(LLM_MAX_RETRIES + 1):
        try:
            return fn(*args, **kwargs)
        except Exception as e:  # noqa: BLE001
            if not _is_rate_limit(e) or attempt >= LLM_MAX_RETRIES:
                raise
            logger.warning(
                "LLM rate-limited (%s) on attempt %d; waiting %.0fs for limit window to reset: %s",
                type(e).__name__, attempt + 1, LLM_RETRY_WAIT, e,
            )
            time.sleep(LLM_RETRY_WAIT)
    raise RuntimeError("LLM retry loop exhausted without success")


def _gemini():
    """Import the gemini_client module (top-level or task-package layout)."""
    try:
        import gemini_client as gc
    except ImportError:
        from task import gemini_client as gc  # type: ignore
    return gc


def get_llm_client():
    """Return a low-level client exposing .invoke_model(), backed by
    google-genai (GeminiBedrockClient)."""
    return _gemini().GeminiBedrockClient()


def model_call(prompt, max_tokens=16384, use_thinking=False, system=None,
               json_mode=False, response_schema=None, use_search=False):
    """Single-shot Gemini call. Returns the text response. Retries on throttling."""
    return _with_retry(
        _gemini().gemini_model_call, prompt,
        max_tokens=max_tokens, use_thinking=use_thinking, system=system,
        json_mode=json_mode, response_schema=response_schema, use_search=use_search,
    )


def invoke_messages(messages, system=None, max_tokens=16384, temperature=0,
                    use_thinking=False, use_search=False):
    """Multi-turn Gemini call. Returns the text response. Retries on throttling."""
    return _with_retry(
        _gemini().gemini_invoke_model, messages=messages, system=system,
        max_tokens=max_tokens, temperature=temperature,
        use_thinking=use_thinking, use_search=use_search,
    )


def get_chat_model(model_id=None, temperature=0, max_tokens=16384,
                   model_kwargs=None, use_search=False):
    """Return a chat model with .invoke() (GeminiChatModel).

    model_id is accepted for call-site compatibility but ignored — the active
    model is always GEMINI_MODEL_ID. GeminiChatModel records token usage
    internally."""
    kwargs = dict(model_kwargs) if model_kwargs else {}
    kwargs.setdefault("temperature", temperature)
    kwargs.setdefault("max_tokens", max_tokens)
    if use_search:
        kwargs["use_search"] = True
    return _gemini().GeminiChatModel(model_id=model_id or GEMINI_MODEL_ID,
                                     model_kwargs=kwargs)


def get_converse_client(max_tokens=16384, thinking_budget=1024, use_search=False):
    """Return a converse-style Gemini client (completion / completion_with_image)."""
    return _gemini().GeminiConverseClient(
        max_tokens=max_tokens, thinking_budget=thinking_budget, use_search=use_search,
    )