Ruurd's picture
Deploy OpenAI classification and protected terminology streaming
c2b7494 verified
Raw History Blame Contribute Delete
13.7 kB
"""Shared model clients for generative JSON calls and Jev typed decisions.
Automatic provider selection preserves the original Anthropic-then-OpenAI
precedence; Google is selected explicitly by experiments that require it.
"""
from __future__ import annotations
import json
import os
import re
import time
import uuid
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from src.utils.logging import get_logger
logger = get_logger(__name__)
_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
DEFAULT_ANTHROPIC_MODEL = "claude-haiku-4-8"
DEFAULT_OPENAI_MODEL = "gpt-6-luna"
DEFAULT_GOOGLE_MODEL = "gemini-3-flash-preview"
DEFAULT_JEV_MODEL = "jev-latest"
DEFAULT_JEV_API_URL = "https://api.typesafe.ai/v1/systemone"
_PLACEHOLDER_KEY_VALUES = {"", "your-api-key-here"}
class LLMCallError(RuntimeError):
"""Raised when the LLM call fails or returns unparsable output after retries."""
def _has_real_key(env_var: str) -> bool:
"""Return True if *env_var* is set to a non-empty, non-placeholder value.
Guards against `.env` files copied from `.env.example` where a key is
present but still holds the literal placeholder text.
"""
value = os.environ.get(env_var, "").strip()
return value.lower() not in _PLACEHOLDER_KEY_VALUES
def resolve_provider(provider: str = "auto") -> str:
"""Resolve the LLM provider to use, based on explicit choice or available API keys.
Parameters
----------
provider : str
One of ``"auto"``, ``"anthropic"``, ``"openai"``, ``"google"``, or ``"jev"``.
Returns
-------
str
The resolved provider name.
"""
provider = (provider or "auto").lower()
if provider in {"anthropic", "openai", "google", "jev", "openai_decisions", "laya", "jpt", "mlx"}:
return provider
if provider != "auto":
raise LLMCallError(
f"Unknown LLM provider: {provider!r}. "
"Use 'auto', 'anthropic', 'openai', 'google', or 'jev'."
)
if _has_real_key("ANTHROPIC_API_KEY"):
return "anthropic"
if _has_real_key("OPENAI_API_KEY"):
return "openai"
if _has_real_key("JEV_API_KEY"):
return "jev"
raise LLMCallError(
"No LLM API key found. Set ANTHROPIC_API_KEY and/or OPENAI_API_KEY in your .env file."
)
def default_model_for(provider: str) -> str:
"""Return the default model identifier for the given provider."""
if provider == "anthropic":
return DEFAULT_ANTHROPIC_MODEL
if provider == "google":
return DEFAULT_GOOGLE_MODEL
if provider == "jev":
return DEFAULT_JEV_MODEL
if provider in {"laya", "jpt", "mlx"}:
return {"laya": "multilingual", "jpt": "kirp/jpt-4b",
"mlx": "mlx-community/Qwen3.5-9B-8bit"}[provider]
return DEFAULT_OPENAI_MODEL
def _get_anthropic_client():
"""Lazily construct the Anthropic client so importing this module does not
require an API key to be set (e.g. for unit tests that mock calls)."""
import anthropic
api_key = os.environ.get("ANTHROPIC_API_KEY")
if not api_key:
raise LLMCallError("ANTHROPIC_API_KEY environment variable is not set.")
return anthropic.Anthropic(api_key=api_key)
def _get_openai_client():
"""Lazily construct the OpenAI client so importing this module does not
require an API key to be set (e.g. for unit tests that mock calls)."""
import openai
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
raise LLMCallError("OPENAI_API_KEY environment variable is not set.")
return openai.OpenAI(api_key=api_key)
def _get_google_client():
"""Lazily construct the Google Gen AI client."""
from google import genai
api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
if not api_key:
raise LLMCallError("GOOGLE_API_KEY or GEMINI_API_KEY environment variable is not set.")
return genai.Client(api_key=api_key)
def extract_json(text: str) -> dict[str, Any]:
"""Extract and parse the first JSON object found in *text*.
Handles cases where the model wraps JSON in markdown code fences or adds
surrounding commentary despite instructions not to.
"""
text = text.strip()
try:
return json.loads(text)
except json.JSONDecodeError:
pass
match = _JSON_BLOCK_RE.search(text)
if match:
return json.loads(match.group(0))
raise ValueError(f"Could not extract JSON from LLM response: {text[:200]!r}")
def _call_anthropic(system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
client = _get_anthropic_client()
response = client.messages.create(
model=model,
max_tokens=max_tokens,
system=system_prompt,
messages=[{"role": "user", "content": user_prompt}],
)
return "".join(
block.text for block in response.content if getattr(block, "type", None) == "text"
)
_REASONING_MODEL_PREFIXES = ("gpt-5", "gpt-6", "o1", "o3", "o4")
def _is_reasoning_model(model: str) -> bool:
"""Whether *model* is an OpenAI reasoning model (GPT-5/6 and o-series).
Reasoning models spend part of their completion token budget on hidden
reasoning tokens before producing visible output, and support a
``reasoning_effort`` parameter to control that trade-off. Non-reasoning
models (e.g. gpt-4o) don't accept that parameter, so it must only be sent
for models in this family.
"""
return model.lower().startswith(_REASONING_MODEL_PREFIXES)
def _call_openai(
system_prompt: str,
user_prompt: str,
model: str,
max_tokens: int,
reasoning_effort: str | None = None,
) -> str:
client = _get_openai_client()
kwargs: dict[str, Any] = {}
if reasoning_effort and _is_reasoning_model(model):
kwargs["reasoning_effort"] = reasoning_effort
response = client.chat.completions.create(
model=model,
max_completion_tokens=max_tokens,
response_format={"type": "json_object"},
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
**kwargs,
)
choice = response.choices[0]
content = choice.message.content or ""
if not content.strip():
# Most commonly caused by a reasoning model spending its entire
# max_completion_tokens budget on hidden reasoning, leaving nothing
# for the visible answer (finish_reason == "length" in that case).
raise ValueError(
f"Empty response content from model '{model}' "
f"(finish_reason={choice.finish_reason!r}, usage={response.usage!r}). "
"Consider increasing max_tokens or lowering reasoning_effort."
)
return content
def _call_google(system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
from google.genai import types
client = _get_google_client()
response = client.models.generate_content(
model=model,
contents=user_prompt,
config=types.GenerateContentConfig(
system_instruction=system_prompt,
response_mime_type="application/json",
max_output_tokens=max_tokens,
),
)
text = response.text or ""
if not text.strip():
raise ValueError(f"Empty response content from Google model '{model}'.")
return text
def call_jev_decisions(
state: str,
questions: dict[str, dict[str, Any]],
model: str = DEFAULT_JEV_MODEL,
max_retries: int = 3,
) -> dict[str, Any]:
"""Submit typed questions to Jev and return its structured response.
Jev is not a generative chat model, so it does not use
:func:`call_llm_json`. A request shares one phenotype state across up to 20
candidate-classification questions and receives declared labels plus
probability distributions.
"""
api_key = (
os.environ.get("TYPESAFE_API_KEY", "").strip()
or os.environ.get("JEV_API_KEY", "").strip()
)
if not api_key:
raise LLMCallError(
"TYPESAFE_API_KEY or JEV_API_KEY environment variable is not set."
)
if not questions or len(questions) > 20:
raise ValueError("Jev requests must contain between 1 and 20 questions.")
url = os.environ.get("JEV_API_URL", DEFAULT_JEV_API_URL).strip() or DEFAULT_JEV_API_URL
payload = json.dumps(
{"model": model, "state": state, "questions": questions},
separators=(",", ":"),
).encode("utf-8")
idempotency_key = str(uuid.uuid5(uuid.NAMESPACE_URL, payload.decode("utf-8")))
retryable_statuses = {429, 502, 503, 504, 529}
last_error: Exception | None = None
for attempt in range(1, max_retries + 1):
request = Request(
url,
data=payload,
method="POST",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"Accept": "application/json",
"Idempotency-Key": idempotency_key,
},
)
try:
with urlopen(request, timeout=60) as response:
result = json.loads(response.read().decode("utf-8"))
if not isinstance(result.get("answers"), dict):
raise ValueError("Jev response is missing an 'answers' object.")
return result
except HTTPError as exc:
try:
error_body = exc.read().decode("utf-8", errors="replace")[:500]
except Exception: # noqa: BLE001 - preserve the original HTTP status
error_body = ""
last_error = RuntimeError(
f"Jev HTTP {exc.code} from {url}"
+ (f": {error_body}" if error_body else "")
)
if exc.code not in retryable_statuses or attempt >= max_retries:
break
retry_after = exc.headers.get("Retry-After") if exc.headers else None
delay = min(float(retry_after), 30.0) if retry_after else min(2 ** (attempt - 1), 8)
time.sleep(delay)
except (URLError, TimeoutError, ValueError, json.JSONDecodeError) as exc:
last_error = exc
if attempt < max_retries:
time.sleep(min(2 ** (attempt - 1), 8))
raise LLMCallError(
f"Jev decision call failed after {max_retries} attempts: {last_error}"
) from last_error
def call_llm_json(
system_prompt: str,
user_prompt: str,
provider: str = "auto",
model: str | None = None,
max_retries: int = 3,
max_tokens: int = 8192,
reasoning_effort: str | None = None,
thinking_budget: int | None = None,
) -> dict[str, Any]:
"""Call the configured LLM provider and parse a JSON object from the response,
retrying on transient failures or malformed JSON output.
Parameters
----------
system_prompt : str
System prompt establishing the LLM's role and output format.
user_prompt : str
The task-specific user prompt.
provider : str
``"auto"``, ``"anthropic"``, ``"openai"``, or ``"google"``. ``"auto"`` selects
Anthropic if ``ANTHROPIC_API_KEY`` is set, otherwise OpenAI.
model : str, optional
Model identifier. Defaults to a sensible default for the resolved provider.
max_retries : int
Number of attempts before raising ``LLMCallError``.
max_tokens : int
Maximum completion tokens to generate. For OpenAI reasoning models
(gpt-5/o1/o3/o4), this budget is shared with hidden reasoning tokens,
so it should be generous enough to leave room for the visible answer.
reasoning_effort : str, optional
Passed as ``reasoning_effort`` to OpenAI reasoning models only (e.g.
``"low"``/``"medium"``/``"high"``); ignored for Anthropic and for
non-reasoning OpenAI models.
Returns
-------
dict
Parsed JSON response body.
"""
resolved_provider = resolve_provider(provider)
if resolved_provider in {"jev", "openai_decisions", "laya", "jpt"}:
raise LLMCallError(
"Jev uses typed decisions rather than generative JSON. "
"Call call_jev_decisions(), or use llm_classify(provider='jev')."
)
resolved_model = model or default_model_for(resolved_provider)
last_error: Exception | None = None
for attempt in range(1, max_retries + 1):
try:
if resolved_provider == "anthropic":
text = _call_anthropic(system_prompt, user_prompt, resolved_model, max_tokens)
elif resolved_provider == "openai":
text = _call_openai(
system_prompt, user_prompt, resolved_model, max_tokens, reasoning_effort
)
elif resolved_provider == "mlx":
from src.llm.local import call_mlx_json
return call_mlx_json(system_prompt, user_prompt, resolved_model, max_tokens, reasoning_effort,
thinking_budget if thinking_budget is not None else 1024)
else:
text = _call_google(system_prompt, user_prompt, resolved_model, max_tokens)
return extract_json(text)
except Exception as exc: # noqa: BLE001 - broad by design, retried below
last_error = exc
logger.warning(
"LLM call attempt %d/%d (provider=%s) failed: %s",
attempt, max_retries, resolved_provider, exc,
)
raise LLMCallError(
f"LLM call failed after {max_retries} attempts (provider={resolved_provider}): {last_error}"
) from last_error