HR-ATS-Portal/backend/llm_setup.py

"""OpenAI async client and a single llm_call helper.

Pure module: no FastAPI imports and no HTTPException.

Config is module-level `os.getenv` (house style for non-DB secrets); the client is
lazy, created on first use like `db_setup.get_engine()`.

    text = await llm_call(system, user)
    data = await llm_call(system, user, json_mode=True)

`init_llm()` confirms the key on startup and `close_llm()` disposes of the connection
pool, so both can hang off the FastAPI lifespan beside `init_db()` / `close_db()`.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os

from dotenv import load_dotenv
from openai import APIError, APIStatusError, AsyncOpenAI

load_dotenv()

logger = logging.getLogger("llm")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL") or None
OPENAI_ORGANIZATION = os.getenv("OPENAI_ORGANIZATION") or None
OPENAI_PROJECT = os.getenv("OPENAI_PROJECT") or None
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-5.4-mini")
# Default kept under gpt-4o-mini's 16384 completion ceiling: a larger value is a 400
# on every call, not a bigger budget.
OPENAI_MAX_OUTPUT_TOKENS = int(os.getenv("OPENAI_MAX_OUTPUT_TOKENS") or 4096)
OPENAI_TIMEOUT = float(os.getenv("OPENAI_TIMEOUT") or 60)
OPENAI_MAX_RETRIES = int(os.getenv("OPENAI_MAX_RETRIES") or 3)
OPENAI_CONNECT_RETRIES = int(os.getenv("OPENAI_CONNECT_RETRIES") or 3)

# Blank OPENAI_TEMPERATURE means omit the param (some models reject it).
_raw_temp = (os.getenv("OPENAI_TEMPERATURE") or "").strip()
OPENAI_TEMPERATURE = float(_raw_temp) if _raw_temp else None

_client: AsyncOpenAI | None = None


def get_client() -> AsyncOpenAI:
    """The process-wide AsyncOpenAI client, created on first use."""
    global _client
    if _client is None:
        if not OPENAI_API_KEY:
            raise RuntimeError("OPENAI_API_KEY is not configured")
        _client = AsyncOpenAI(
            api_key=OPENAI_API_KEY,
            base_url=OPENAI_BASE_URL,
            organization=OPENAI_ORGANIZATION,
            project=OPENAI_PROJECT,
            timeout=OPENAI_TIMEOUT,
            max_retries=OPENAI_MAX_RETRIES,
        )
    return _client


async def llm_call(system, user, *, model=None, temperature=None, json_mode=False):
    """One system+user turn. Returns text, or a parsed dict when json_mode=True.

    With json_mode the prompt must mention JSON somewhere or the API rejects the call.
    """
    kwargs = {
        "model": model or OPENAI_MODEL,
        "max_completion_tokens": OPENAI_MAX_OUTPUT_TOKENS,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
    }
    resolved = OPENAI_TEMPERATURE if temperature is None else temperature
    if resolved is not None:
        kwargs["temperature"] = resolved
    if json_mode:
        kwargs["response_format"] = {"type": "json_object"}

    response = await get_client().chat.completions.create(**kwargs)
    _log_usage(response, kwargs["model"])
    content = (response.choices[0].message.content or "").strip()
    if not json_mode:
        return content
    try:
        return json.loads(content)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"model did not return valid JSON: {content[:200]}") from exc


def _log_usage(response, model) -> None:
    """Per-call token and cache visibility. Never logs prompt or reply text."""
    usage = getattr(response, "usage", None)
    if usage is None:
        return
    prompt_details = getattr(usage, "prompt_tokens_details", None)
    logger.info(
        "llm usage: model=%s request_id=%s prompt_tokens=%s completion_tokens=%s cached_tokens=%s",
        model,
        getattr(response, "_request_id", None),
        getattr(usage, "prompt_tokens", None),
        getattr(usage, "completion_tokens", None),
        getattr(prompt_details, "cached_tokens", None),
    )


async def check_connection(retries=None, delay=1.0):
    """Confirm the key works, retrying with a capped backoff."""
    attempts = OPENAI_CONNECT_RETRIES if retries is None else retries
    for attempt in range(1, max(attempts, 1) + 1):
        try:
            await get_client().models.list()
            logger.info("openai reachable, default model %s", OPENAI_MODEL)
            return
        except APIStatusError as exc:
            if exc.status_code in (401, 403):
                raise RuntimeError(f"OPENAI_API_KEY rejected ({exc.status_code})") from exc
            if attempt >= attempts:
                raise
            logger.warning("openai not ready (%s/%s): %s", attempt, attempts, exc)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 10.0)
        except APIError as exc:
            if attempt >= attempts:
                raise
            logger.warning("openai not ready (%s/%s): %s", attempt, attempts, exc)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 10.0)


async def init_llm(*, verify=True):
    """Build the client and, unless told otherwise, confirm the key is live."""
    get_client()
    if verify:
        await check_connection()


async def close_llm():
    """Close the underlying httpx pool and reset the cached client."""
    global _client
    if _client is not None:
        await _client.close()
        logger.info("openai client closed")
    _client = None


# if __name__ == "__main__":
#     logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")

#     async def _main():
#         try:
#             await init_llm()
#             print(await llm_call("You are terse.", "Reply with the single word: ready"))
#         finally:
#             await close_llm()

#     asyncio.run(_main())