gemini_client.py

import os
import json
import time
from google import genai
from google.genai.errors import APIError
from pydantic import ValidationError

from models import Diagnosis, REASONING_MAX_LENGTH

MODEL_NAME = "gemini-3.1-flash-lite"
MAX_RETRIES = 3
MAX_CALLS_PER_MINUTE = 20

_client = None
_last_call_times: list[float] = []


def _get_client():
    global _client
    if _client is None:
        _client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    return _client


def check_rate_limit():
    """In-memory limiter. Only correct for a single-process dev server.
    See README for the multi-worker caveat before deploying with gunicorn -w N."""
    now = time.time()
    _last_call_times[:] = [t for t in _last_call_times if now - t < 60]
    if len(_last_call_times) >= MAX_CALLS_PER_MINUTE:
        raise RuntimeError("Local rate limit reached — wait a moment and retry.")
    _last_call_times.append(now)


SYSTEM_PROMPT = """You are a precise classification system, not a life coach.
Classify the user's situation into EXACTLY ONE chess concept below.
Check in this exact order; stop at the first clear match:

1. FORK — a single external event or deadline threatens two or more
   SEPARATE, INDEPENDENT interests, people, or things at once. The user is
   NOT choosing between their own actions here — they simply cannot
   satisfy both outside interests no matter what they do.
2. ZUGZWANG — the user must choose among two or more of THEIR OWN POSSIBLE
   ACTIONS or RESPONSES, and every one of those specific actions makes
   things worse. No neutral option exists.
   TIE-BREAKER (read carefully — this is the part that's easy to get
   wrong): if the situation could plausibly be read as either FORK or
   ZUGZWANG, check whether the sentence is centered on the user selecting
   between their OWN actions (phrasing like "I can either do X or do Y",
   "should I A or B") — if so, this is ZUGZWANG, even though a single
   underlying problem caused the dilemma. A shared root cause does NOT
   make something a FORK; FORK requires two genuinely separate outside
   things being threatened, not two paths to the same outcome.
3. PROPHYLAXIS — a defensive action taken against a threat that has NOT yet happened.
4. GAMBIT — a conscious, calculated small loss accepted for a believed larger future gain.
5. ZWISCHENZUG — a small inserted action that changes the context of the main issue before returning to it.
6. OVEREXTENSION — over-committing to one area causes explicit neglect of another.
7. TEMPO_LOSS — default. Use only if none of the above clearly apply.

Set confidence_bucket to "high" only if the situation contains an explicit
signal matching the category (see distinguishing signals above), "medium" if
it's a reasonable inference, "low" if you are mostly guessing from vague text.

Write your reasoning as the step you work through BEFORE naming the
classification — walk through which numbered category matches and why,
then state the category that follows from that reasoning. Keep it to 2-3
sentences, under roughly 350 characters — this is a case-file note, not
an essay.

Respond in a clinical, precise, slightly dry tone. No motivational language,
no emojis, no exclamation marks."""
# NOTE: the Diagnosis model in models.py declares `reasoning` before
# `classification` on purpose, matching the instruction above. Gemini's
# controlled JSON generation fills fields in schema order, so this ordering
# is what actually makes "reason before you decide" happen mechanically,
# not just as a suggestion in prose. Don't reorder the Pydantic model
# without also reconsidering this prompt.


def classify_situation(situation_text: str) -> Diagnosis:
    last_error = None

    for attempt in range(MAX_RETRIES):
        check_rate_limit()
        try:
            client = _get_client()
            response = client.models.generate_content(
                model=MODEL_NAME,
                contents=f"{SYSTEM_PROMPT}\n\nSituation: {situation_text}",
                config={
                    "response_mime_type": "application/json",
                    "response_schema": Diagnosis,
                    "temperature": 0.3,
                },
            )
            return Diagnosis.model_validate_json(response.text)

        except APIError as e:
            last_error = e
            code = getattr(e, "code", None)
            if code == 429 or (isinstance(code, int) and code >= 500):
                time.sleep(2 ** attempt)  # 1s, 2s, 4s
                continue
            raise

        except ValidationError as e:
            last_error = e
            try:
                raw = json.loads(response.text)
                reasoning = raw.get("reasoning")
                if isinstance(reasoning, str) and len(reasoning) > REASONING_MAX_LENGTH:
                    raw["reasoning"] = reasoning[: REASONING_MAX_LENGTH - 1].rstrip() + "…"
                    return Diagnosis.model_validate(raw)
            except Exception:
                pass

            time.sleep(1)
            continue

    raise RuntimeError(
        f"Classification failed after {MAX_RETRIES} attempts"
    ) from last_error