143 lines
4.9 KiB
Python
143 lines
4.9 KiB
Python
"""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")
|
|
OPENAI_MAX_OUTPUT_TOKENS = int(os.getenv("OPENAI_MAX_OUTPUT_TOKENS") or 32768)
|
|
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": 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)
|
|
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
|
|
|
|
|
|
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())
|