132 lines
4.6 KiB
Python
132 lines
4.6 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional, Type
|
|
|
|
from openai import OpenAI
|
|
from pydantic import BaseModel
|
|
|
|
class NoSchema(BaseModel): ...
|
|
|
|
|
|
logger = logging.getLogger("llm client")
|
|
|
|
class ResponseCache:
|
|
"""File-backed cache: prompt-hash -> validated-model JSON. Makes runs idempotent."""
|
|
|
|
def __init__(self, path: Optional[Path] = None) -> None:
|
|
self.path = path
|
|
self._mem: dict[str, str] = {}
|
|
if path and path.exists():
|
|
try:
|
|
self._mem = json.loads(path.read_text("utf-8"))
|
|
except Exception: # noqa: BLE001
|
|
logger.warning("Corrupt cache at %s; starting fresh", path)
|
|
|
|
@staticmethod
|
|
def key(system: str, user: str, schema: Type[BaseModel]) -> str:
|
|
h = hashlib.sha256()
|
|
h.update(schema.__name__.encode())
|
|
h.update(b"\x00")
|
|
h.update(system.encode())
|
|
h.update(b"\x00")
|
|
h.update(user.encode())
|
|
return h.hexdigest()
|
|
|
|
def get(self, k: str) -> Optional[str]:
|
|
return self._mem.get(k)
|
|
|
|
def put(self, k: str, v: str) -> None:
|
|
self._mem[k] = v
|
|
if self.path:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
|
|
tmp.write_text(json.dumps(self._mem, indent=2, sort_keys=True), "utf-8")
|
|
tmp.replace(self.path)
|
|
else:
|
|
logger.warning("No path specified for cache!")
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self, base_url, api_key: str | None = None, model: str = "gpt-4o-mini", cache: Optional[ResponseCache] = None):
|
|
self.model = model
|
|
self.client = OpenAI(base_url=base_url, api_key=api_key or "ollama")
|
|
self.cache = cache
|
|
|
|
def chat_with_schema(self, prompt: str, schema: dict, system: str | None = None) -> str | None:
|
|
# --- idempotency: serve from cache when prompt+system+schema match ---
|
|
|
|
cache_key = ResponseCache.key(system or "", prompt, NoSchema) if self.cache else None
|
|
if cache_key and self.cache and (hit := self.cache.get(cache_key)) is not None:
|
|
return json.loads(hit).get('response', "")
|
|
|
|
assert_strict_mode_clean(schema)
|
|
kwargs = {
|
|
"model": self.model,
|
|
"messages": [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt},
|
|
] if system else [
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"response_format": schema,
|
|
}
|
|
resp = self.client.chat.completions.create(**kwargs)
|
|
|
|
if cache_key and self.cache:
|
|
self.cache.put(cache_key, json.dumps({"response": resp.choices[0].message.content or ""}))
|
|
|
|
return resp.choices[0].message.content
|
|
|
|
def chat(self, system: str, user: str, json_mode: bool = False, images: list[str] | None = None) -> str | None:
|
|
images = images or []
|
|
if images:
|
|
content = [{"type": "text", "text": user}]
|
|
for b64 in images:
|
|
content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{b64}"},
|
|
})
|
|
else:
|
|
content = user
|
|
|
|
kwargs = {
|
|
"model": self.model,
|
|
"messages": [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": content},
|
|
],
|
|
}
|
|
if json_mode:
|
|
kwargs["response_format"] = {"type": "json_object"}
|
|
resp = self.client.chat.completions.create(**kwargs)
|
|
return resp.choices[0].message.content
|
|
|
|
|
|
def assert_strict_mode_clean(schema: dict, path: str = "$") -> None:
|
|
"""Recursively verify a JSON Schema obeys OpenAI strict mode."""
|
|
if schema.get("type") == "object":
|
|
props = schema.get("properties", {})
|
|
required = set(schema.get("required", []))
|
|
|
|
if schema.get("additionalProperties") is not False:
|
|
raise ValueError(f"{path}: additionalProperties must be false")
|
|
|
|
for name in props:
|
|
if name not in required:
|
|
raise ValueError(f"{path}.{name}: not in required")
|
|
|
|
for name, sub in props.items():
|
|
assert_strict_mode_clean(sub, f"{path}.{name}")
|
|
|
|
if schema.get("type") == "array":
|
|
assert_strict_mode_clean(schema["items"], f"{path}[]")
|
|
|
|
for name, sub in schema.get("$defs", {}).items():
|
|
assert_strict_mode_clean(sub, f"$defs.{name}")
|
|
|
|
# OpenAI rejects these keywords in strict mode
|
|
for banned in ("default", "oneOf"):
|
|
if banned in schema:
|
|
raise ValueError(f"{path}: '{banned}' not allowed in strict mode") |