woogi
7488cf68c5
- [LLM] line per call: stage, model, prompt size, output size, cost, parsed item counts - LLM_RAW_DUMP: full prompt/response JSON per call under outputs/<job>/llm_raw/ - Cost block at tail of job.log (per-stage, per-model, cached vs live) - Agent mode: reset llm cost counters per job; review finalization now teed into job.log + dumps
491 lines
18 KiB
Python
491 lines
18 KiB
Python
"""
|
|
llm.py - Thin OpenRouter client + a JSON-returning chat helper.
|
|
|
|
Mirrors the call pattern proven in the IronBid pipeline: OpenAI SDK pointed at
|
|
OpenRouter, base64 image data URLs, markdown-fence stripping, and a single
|
|
retry on JSON parse failure. One helper serves both the extraction (Stage 1,
|
|
with an image) and conflict-reasoning (Stage 3, with images) calls.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import hashlib
|
|
import threading
|
|
from typing import List, Optional, Dict, Any
|
|
|
|
from openai import OpenAI
|
|
|
|
from backend import config
|
|
|
|
_clients: Dict[str, OpenAI] = {}
|
|
|
|
# Per-run text backend choice (hybrid). Set by the runner at job start; read by
|
|
# call_json when routing a no-image (text) call. Module-global mirrors the
|
|
# set_stage/cost pattern (single-user tool).
|
|
_text_local = False
|
|
# Per-job model overrides (user picked models in the UI). Same module-global
|
|
# pattern: set by the job runner before the pipeline starts, cleared after.
|
|
# Vision applies to image calls, text to no-image calls on OpenRouter (and to
|
|
# the local->cloud fallback). The LOCAL endpoint's model name is never taken
|
|
# from these overrides - hybrid local keeps LOCAL_TEXT_MODEL.
|
|
_vision_model_override: Optional[str] = None
|
|
_text_model_override: Optional[str] = None
|
|
|
|
# Per-job raw request/response dumps (missed/hallucinated-item debugging).
|
|
# Set by the job runner to <out_dir>/llm_raw at job start, cleared after.
|
|
# Same module-global pattern as the model overrides (single-user tool).
|
|
_raw_dump_dir: Optional[str] = None
|
|
_seq_lock = threading.Lock()
|
|
_call_seq = 0
|
|
|
|
|
|
def set_raw_dump_dir(path: Optional[str]) -> None:
|
|
"""Point raw LLM request/response dumps at a directory. None disables."""
|
|
global _raw_dump_dir, _call_seq
|
|
with _seq_lock:
|
|
_raw_dump_dir = path
|
|
_call_seq = 0
|
|
|
|
|
|
def _next_seq() -> int:
|
|
global _call_seq
|
|
with _seq_lock:
|
|
_call_seq += 1
|
|
return _call_seq
|
|
|
|
|
|
def _summarize_parsed(parsed: Any) -> str:
|
|
"""
|
|
Compact digest of a parsed response for the job log. List values become
|
|
item counts (e.g. conflicts[3]) so a stage that returned nothing (miss)
|
|
or invented items (hallucination) is visible without opening the raw dump.
|
|
"""
|
|
if isinstance(parsed, list):
|
|
return f"list[{len(parsed)}]"
|
|
if not isinstance(parsed, dict):
|
|
return type(parsed).__name__
|
|
parts = []
|
|
for k, v in parsed.items():
|
|
if isinstance(v, list):
|
|
parts.append(f"{k}[{len(v)}]")
|
|
elif isinstance(v, dict):
|
|
parts.append(f"{k}{{{len(v)}}}")
|
|
else:
|
|
s = str(v)
|
|
parts.append(f"{k}={s[:40]!r}{'...' if len(s) > 40 else ''}")
|
|
out = ", ".join(parts)
|
|
return out[:300] + ("..." if len(out) > 300 else "")
|
|
|
|
|
|
def _dump_raw(seq: int, be: Dict[str, Any], usage_stage: str,
|
|
system_prompt: str, user_text: str,
|
|
images_b64: Optional[List[str]], max_tokens: int,
|
|
raw: str, parsed: Any) -> None:
|
|
"""Write the full request/response for one call to the job's llm_raw dir."""
|
|
if not _raw_dump_dir:
|
|
return
|
|
try:
|
|
os.makedirs(_raw_dump_dir, exist_ok=True)
|
|
safe_stage = re.sub(r"[^A-Za-z0-9_.-]+", "_", usage_stage)[:40]
|
|
safe_model = re.sub(r"[^A-Za-z0-9_.-]+", "_", be["model"])
|
|
payload = {
|
|
"seq": seq,
|
|
"stage": usage_stage,
|
|
"model": be["model"],
|
|
"backend": "local" if be.get("local") else "cloud",
|
|
"max_tokens": max_tokens,
|
|
# base64 image payloads deliberately excluded (multi-MB each);
|
|
# the count + source.pdf in the job dir identify what was sent.
|
|
"n_images": len(images_b64 or []),
|
|
"system_prompt": system_prompt,
|
|
"user_text": user_text,
|
|
"raw_response": raw,
|
|
"parsed": parsed,
|
|
}
|
|
path = os.path.join(_raw_dump_dir, f"{seq:04d}_{safe_stage}_{safe_model}.json")
|
|
tmp = f"{path}.{threading.get_ident()}.tmp"
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, indent=2)
|
|
os.replace(tmp, path) # atomic so thread-pooled stages can't tear it
|
|
except OSError as e:
|
|
print(f"[LLM] raw dump failed: {e}")
|
|
|
|
|
|
def _log_call(seq: int, be: Dict[str, Any], usage_stage: str,
|
|
user_text: str, images_b64: Optional[List[str]],
|
|
raw: str, parsed: Any, usd: Optional[float],
|
|
cached: bool = False) -> None:
|
|
"""One verbose per-call line for the job log (tee'd by job_log.py)."""
|
|
if not config.LLM_VERBOSE:
|
|
return
|
|
backend = "local" if be.get("local") else "cloud"
|
|
if cached:
|
|
cost_str = "cache hit"
|
|
elif usd is not None:
|
|
cost_str = f"${usd:.4f}"
|
|
else:
|
|
cost_str = "cost n/a"
|
|
print(f"[LLM] #{seq:04d} {usage_stage} | {be['model']} ({backend}) | "
|
|
f"in {len(user_text)}ch+{len(images_b64 or [])}img | "
|
|
f"out {len(raw)}ch | {cost_str} | {_summarize_parsed(parsed)}")
|
|
|
|
|
|
def set_text_backend(local: bool) -> None:
|
|
"""Choose whether text (no-image) calls go to the local endpoint this run."""
|
|
global _text_local
|
|
_text_local = bool(local)
|
|
|
|
|
|
def set_model_overrides(vision: Optional[str] = None, text: Optional[str] = None) -> None:
|
|
"""Per-run OpenRouter vision/text model picks. None/blank clears to defaults."""
|
|
global _vision_model_override, _text_model_override
|
|
_vision_model_override = (vision or "").strip() or None
|
|
_text_model_override = (text or "").strip() or None
|
|
|
|
# --- per-job cost accounting -------------------------------------------------
|
|
# OpenRouter returns the real USD cost of each call when we request usage
|
|
# accounting. We accumulate it in a module-level counter; the runner resets it
|
|
# at the start of a job and reads it at the end. (Single-user tool: overlapping
|
|
# jobs would share this counter, which is acceptable here.)
|
|
_cost_lock = threading.Lock()
|
|
_cost = {"usd": 0.0, "calls": 0, "cached": 0}
|
|
_by_stage: Dict[str, Dict[str, Any]] = {}
|
|
_current_stage = "?"
|
|
|
|
# Model usage tracking - records which models were actually called this run.
|
|
_models: Dict[str, Any] = {
|
|
"vision": set(), # models used for vision (image) calls
|
|
"text_local": set(), # local model names attempted
|
|
"text_cloud": set(), # cloud text model names used
|
|
"fallback_count": 0, # times local fell back to cloud
|
|
}
|
|
|
|
|
|
def set_stage(name: str) -> None:
|
|
"""Label subsequent calls so cost can be attributed per pipeline stage."""
|
|
global _current_stage
|
|
_current_stage = name
|
|
|
|
|
|
def reset_cost() -> None:
|
|
with _cost_lock:
|
|
_cost["usd"] = 0.0
|
|
_cost["calls"] = 0
|
|
_cost["cached"] = 0
|
|
_by_stage.clear()
|
|
_models["vision"] = set()
|
|
_models["text_local"] = set()
|
|
_models["text_cloud"] = set()
|
|
_models["fallback_count"] = 0
|
|
|
|
|
|
def get_cost() -> Dict[str, Any]:
|
|
with _cost_lock:
|
|
out = dict(_cost)
|
|
out["by_stage"] = {k: dict(v) for k, v in _by_stage.items()}
|
|
out["models"] = {
|
|
"vision": sorted(_models["vision"]),
|
|
"text_local": sorted(_models["text_local"]),
|
|
"text_cloud": sorted(_models["text_cloud"]),
|
|
"fallback_count": _models["fallback_count"],
|
|
}
|
|
return out
|
|
|
|
|
|
def _record_model(be: Dict[str, Any], has_images: bool, fell_back: bool) -> None:
|
|
"""Thread-safe recording of which model was actually used for a call."""
|
|
model = be["model"]
|
|
with _cost_lock:
|
|
if has_images:
|
|
_models["vision"].add(model)
|
|
elif fell_back or not be.get("local"):
|
|
_models["text_cloud"].add(model)
|
|
else:
|
|
_models["text_local"].add(model)
|
|
|
|
|
|
def _stage_bucket() -> Dict[str, Any]:
|
|
return _by_stage.setdefault(_current_stage, {"usd": 0.0, "calls": 0, "cached": 0})
|
|
|
|
|
|
def _add_cost(usd: float) -> None:
|
|
with _cost_lock:
|
|
_cost["usd"] += usd
|
|
_cost["calls"] += 1
|
|
b = _stage_bucket()
|
|
b["usd"] += usd
|
|
b["calls"] += 1
|
|
|
|
|
|
def _add_cached() -> None:
|
|
with _cost_lock:
|
|
_cost["cached"] += 1
|
|
_stage_bucket()["cached"] += 1
|
|
|
|
|
|
# --- disk-backed response cache (testing loop; opt-in via config.LLM_CACHE) ---
|
|
|
|
def _cache_key(model: str, system_prompt: str, user_text: str,
|
|
images_b64: Optional[List[str]], max_tokens: int,
|
|
json_mode: bool) -> str:
|
|
h = hashlib.sha256()
|
|
parts = [model, str(max_tokens), str(json_mode), system_prompt, user_text]
|
|
for b in (images_b64 or []):
|
|
parts.append(b)
|
|
for p in parts:
|
|
h.update(b"\x00")
|
|
h.update(p.encode("utf-8", "ignore"))
|
|
return h.hexdigest()
|
|
|
|
|
|
def _cache_get(key: str) -> Optional[Dict[str, Any]]:
|
|
path = os.path.join(config.LLM_CACHE_DIR, key + ".json")
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
|
|
|
|
def _cache_set(key: str, value: Dict[str, Any]) -> None:
|
|
try:
|
|
os.makedirs(config.LLM_CACHE_DIR, exist_ok=True)
|
|
path = os.path.join(config.LLM_CACHE_DIR, key + ".json")
|
|
tmp = f"{path}.{os.getpid()}.{threading.get_ident()}.tmp"
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(value, f)
|
|
os.replace(tmp, path) # atomic so thread-pooled stages can't tear it
|
|
except OSError as e:
|
|
print(f"[LLM] cache write failed: {e}")
|
|
|
|
|
|
def get_client(base_url: str, api_key: str) -> OpenAI:
|
|
"""Lazily build + cache one OpenAI-compatible client per endpoint."""
|
|
client = _clients.get(base_url)
|
|
if client is None:
|
|
client = OpenAI(base_url=base_url, api_key=api_key)
|
|
_clients[base_url] = client
|
|
return client
|
|
|
|
|
|
def _resolve_backend(has_images: bool, model_override: Optional[str]) -> Dict[str, Any]:
|
|
"""
|
|
Pick (base_url, api_key, model, usage) for this call.
|
|
|
|
Vision (has_images) always uses OpenRouter. Text (no images) uses the local
|
|
endpoint when the run chose hybrid AND it's configured; otherwise OpenRouter.
|
|
"""
|
|
if not has_images and _text_local and config.LOCAL_BASE_URL:
|
|
return {
|
|
"base_url": config.LOCAL_BASE_URL,
|
|
"api_key": config.LOCAL_API_KEY,
|
|
# Local model name comes from per-call args or LOCAL_TEXT_MODEL —
|
|
# never the UI's OpenRouter picks, which a local server won't serve.
|
|
"model": model_override or config.LOCAL_TEXT_MODEL or config.TEXT_MODEL,
|
|
"usage": False, # local has no OpenRouter usage accounting
|
|
"local": True,
|
|
}
|
|
# Vision, or text-on-OpenRouter (default / fallback). A per-job override
|
|
# (user's UI model pick) wins over per-call and env defaults.
|
|
if has_images:
|
|
model = _vision_model_override or model_override or config.MODEL
|
|
else:
|
|
model = _text_model_override or model_override or config.TEXT_MODEL
|
|
return {
|
|
"base_url": config.AI_BASE_URL,
|
|
"api_key": config.AI_API_KEY,
|
|
"model": model,
|
|
"usage": True,
|
|
"local": False,
|
|
}
|
|
|
|
|
|
def _strip_fences(raw: str) -> str:
|
|
raw = raw.strip()
|
|
if raw.startswith("```"):
|
|
# ```json ... ``` or ``` ... ```
|
|
raw = raw.split("```")[1]
|
|
if raw.startswith("json"):
|
|
raw = raw[4:]
|
|
raw = raw.strip()
|
|
return raw
|
|
|
|
|
|
def _repair_truncated(raw: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Best-effort recovery of a truncated JSON object/array.
|
|
|
|
Walks the text tracking bracket depth and string state, finds the last
|
|
point where we can safely cut (after a complete element), drops the partial
|
|
trailing element, and closes the still-open brackets. Lets a response that
|
|
got cut off by max_tokens still yield its complete leading items instead of
|
|
nothing. Returns None if nothing salvageable.
|
|
"""
|
|
stack: List[str] = []
|
|
in_str = esc = False
|
|
cut = None # (index_exclusive, stack_snapshot) of a safe truncation point
|
|
for i, ch in enumerate(raw):
|
|
if in_str:
|
|
if esc:
|
|
esc = False
|
|
elif ch == "\\":
|
|
esc = True
|
|
elif ch == '"':
|
|
in_str = False
|
|
continue
|
|
if ch == '"':
|
|
in_str = True
|
|
elif ch in "{[":
|
|
stack.append("}" if ch == "{" else "]")
|
|
elif ch in "}]":
|
|
if stack:
|
|
stack.pop()
|
|
cut = (i + 1, list(stack)) # after a complete sub-structure
|
|
elif ch == ",":
|
|
cut = (i, list(stack)) # before an (incomplete) next element
|
|
|
|
if not cut:
|
|
return None
|
|
idx, open_brackets = cut
|
|
candidate = raw[:idx].rstrip().rstrip(",")
|
|
candidate += "".join(reversed(open_brackets))
|
|
try:
|
|
return json.loads(candidate)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
|
|
def _response_cost(response) -> Optional[float]:
|
|
"""Pull OpenRouter's per-call USD cost out of the usage object, if present."""
|
|
try:
|
|
dump = response.model_dump()
|
|
except Exception:
|
|
return None
|
|
usage = dump.get("usage") or {}
|
|
cost = usage.get("cost")
|
|
if cost is None:
|
|
cost = (usage.get("cost_details") or {}).get("upstream_inference_cost")
|
|
return float(cost) if isinstance(cost, (int, float)) else None
|
|
|
|
|
|
def _parse(raw: str) -> Optional[Dict[str, Any]]:
|
|
"""Parse model JSON, falling back to truncation repair."""
|
|
try:
|
|
return json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
repaired = _repair_truncated(raw)
|
|
if repaired is not None:
|
|
print("[LLM] recovered truncated JSON (dropped partial trailing item)")
|
|
return repaired
|
|
|
|
|
|
def call_json(
|
|
system_prompt: str,
|
|
user_text: str,
|
|
images_b64: Optional[List[str]] = None,
|
|
max_tokens: int = 4096,
|
|
model: Optional[str] = None,
|
|
usage_tracker: Optional[Any] = None,
|
|
usage_stage: str = "?",
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Send one chat completion expecting a JSON object back.
|
|
|
|
Uses the provider's JSON mode (response_format) so the model returns a bare
|
|
JSON object instead of prose/empty text, and recovers from max_tokens
|
|
truncation. images_b64: optional base64 JPEGs attached as high-detail image
|
|
parts. Returns the parsed dict, or None on a hard failure (caller degrades).
|
|
"""
|
|
has_images = bool(images_b64)
|
|
be = _resolve_backend(has_images, model)
|
|
seq = _next_seq()
|
|
|
|
cache_key = None
|
|
if config.LLM_CACHE:
|
|
cache_key = _cache_key(be["model"], system_prompt, user_text,
|
|
images_b64, max_tokens, json_mode=True)
|
|
hit = _cache_get(cache_key)
|
|
if hit is not None:
|
|
_add_cached()
|
|
if usage_tracker:
|
|
usage_tracker.record(
|
|
usage_stage, be["model"], cached=True, has_images=has_images
|
|
)
|
|
_log_call(seq, be, usage_stage, user_text, images_b64,
|
|
"", hit, None, cached=True)
|
|
return hit
|
|
|
|
content: List[Dict[str, Any]] = []
|
|
for b64 in images_b64 or []:
|
|
content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/jpeg;base64,{b64}", "detail": "high"},
|
|
})
|
|
content.append({"type": "text", "text": user_text})
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": content},
|
|
]
|
|
|
|
use_json_mode = True
|
|
fell_back = False
|
|
for attempt in range(2):
|
|
try:
|
|
client = get_client(be["base_url"], be["api_key"])
|
|
kwargs = dict(model=be["model"], messages=messages,
|
|
max_tokens=max_tokens, timeout=config.LLM_TIMEOUT)
|
|
if be["usage"]:
|
|
kwargs["extra_body"] = {"usage": {"include": True}}
|
|
if use_json_mode:
|
|
kwargs["response_format"] = {"type": "json_object"}
|
|
response = client.chat.completions.create(**kwargs)
|
|
usd = _response_cost(response)
|
|
if usd is not None:
|
|
_add_cost(usd)
|
|
if usage_tracker:
|
|
usage_tracker.record(
|
|
usage_stage, be["model"], usd=usd or 0.0, has_images=has_images
|
|
)
|
|
raw = _strip_fences(response.choices[0].message.content or "")
|
|
parsed = _parse(raw)
|
|
if parsed is not None:
|
|
_record_model(be, has_images, fell_back)
|
|
if cache_key:
|
|
_cache_set(cache_key, parsed)
|
|
_log_call(seq, be, usage_stage, user_text, images_b64,
|
|
raw, parsed, usd)
|
|
if config.LLM_RAW_DUMP:
|
|
_dump_raw(seq, be, usage_stage, system_prompt, user_text,
|
|
images_b64, max_tokens, raw, parsed)
|
|
return parsed
|
|
if attempt == 0:
|
|
print("[LLM] JSON parse error (retrying)")
|
|
continue
|
|
print(f"[LLM] JSON parse error (giving up); raw head: {raw[:300]}")
|
|
return None
|
|
except Exception as e:
|
|
# Some models reject response_format -> drop it and retry once.
|
|
if use_json_mode and "response_format" in str(e).lower():
|
|
print("[LLM] model rejected JSON mode; retrying without it")
|
|
use_json_mode = False
|
|
continue
|
|
# Local endpoint unreachable -> fall back to OpenRouter text backend once.
|
|
if be["local"] and not fell_back:
|
|
print(f"[LLM] local endpoint failed ({e}); falling back to OpenRouter")
|
|
with _cost_lock:
|
|
_models["text_local"].add(be["model"])
|
|
_models["fallback_count"] += 1
|
|
be = {"base_url": config.AI_BASE_URL, "api_key": config.AI_API_KEY,
|
|
"model": _text_model_override or config.TEXT_MODEL,
|
|
"usage": True, "local": False}
|
|
cache_key = None # don't cache fallback under the local-model key
|
|
fell_back = True
|
|
continue
|
|
print(f"[LLM] call error: {e}")
|
|
if attempt == 0:
|
|
continue
|
|
return None
|
|
return None
|