AI Engineer
Deploy Streamlit Content Agent
0b29030
Raw
History Blame Contribute Delete
3.78 kB
"""Provider-agnostic LLM client.
Supports any OpenAI-compatible chat-completions endpoint (Groq by default,
OpenRouter optional). Falls back to a deterministic MOCK backend so the whole
pipeline can run and be tested with no API key and no network.
The client always asks the model for JSON and parses it defensively.
"""
from __future__ import annotations
import json
import re
from typing import Any, Dict, Optional
from .config import Settings
from . import mock_backend
try:
import requests # type: ignore
except Exception: # pragma: no cover - requests may be absent in some envs
requests = None
class LLMError(RuntimeError):
pass
def _extract_json(text: str) -> Any:
"""Pull a JSON object/array out of a model response."""
text = text.strip()
# Strip markdown code fences if present.
fence = re.search(r"```(?:json)?\s*(.*?)```", text, re.DOTALL)
if fence:
text = fence.group(1).strip()
# Direct parse first.
try:
return json.loads(text)
except Exception:
pass
# Fallback: grab the outermost braces/brackets.
for open_c, close_c in (("{", "}"), ("[", "]")):
start = text.find(open_c)
end = text.rfind(close_c)
if start != -1 and end != -1 and end > start:
try:
return json.loads(text[start:end + 1])
except Exception:
continue
raise LLMError("Could not parse JSON from model response: " + text[:200])
class LLMClient:
def __init__(self, settings: Settings):
self.settings = settings
@property
def live(self) -> bool:
return self.settings.is_live()
def generate_json(self, task: str, system: str, user: str,
context: Optional[Dict[str, Any]] = None) -> Any:
"""Return parsed JSON for a given task.
`task` is used by the mock backend to return realistic canned data.
"""
if not self.live:
return mock_backend.respond(task, context or {})
return self._call_api(system, user)
def _call_api(self, system: str, user: str) -> Any:
if requests is None:
raise LLMError("The 'requests' package is required for live API calls.")
s = self.settings
headers = {
"Authorization": f"Bearer {s.api_key}",
"Content-Type": "application/json",
}
if s.provider == "openrouter":
# OpenRouter asks for these (optional but recommended).
headers["HTTP-Referer"] = "https://content-agent.local"
headers["X-Title"] = "Content Agent"
payload = {
"model": s.resolved_model(),
"messages": [
{"role": "system", "content": system},
{"role": "user", "content": user},
],
"temperature": s.temperature,
"max_tokens": s.max_tokens,
"response_format": {"type": "json_object"},
}
try:
resp = requests.post(s.endpoint(), headers=headers,
data=json.dumps(payload), timeout=60)
except Exception as e: # network error
raise LLMError(f"Request to {s.provider} failed: {e}")
if resp.status_code >= 400:
# Retry once without response_format (some models reject it).
payload.pop("response_format", None)
resp = requests.post(s.endpoint(), headers=headers,
data=json.dumps(payload), timeout=60)
if resp.status_code >= 400:
raise LLMError(f"{s.provider} returned {resp.status_code}: {resp.text[:300]}")
data = resp.json()
content = data["choices"][0]["message"]["content"]
return _extract_json(content)