255 lines
10 KiB
Python
255 lines
10 KiB
Python
"""The single Ollama wrapper. Every LLM call in Art goes through generate().
|
|
|
|
Each call sets the model, options, and num_ctx explicitly (Ollama's default
|
|
context silently truncates long prompts, the first failure mode in Stan's
|
|
taxonomy) and writes one row to llm_calls: host, model tag and digest,
|
|
Modelfile hash, prompt name and version, num_ctx, token counts, options,
|
|
latency, the raw response, and the outcome. Raw responses are kept even when
|
|
they fail validation.
|
|
|
|
The model is built from modelfiles/researcher.Modelfile and tagged with the
|
|
file's hash (art-researcher:<hash>), so editing the Modelfile creates a new
|
|
model and the tag alone identifies the Modelfile behind any text.
|
|
|
|
Prompts live in prompts/<task>_v<N>.txt; the version recorded is
|
|
"v<N>+<content hash>".
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import time
|
|
from dataclasses import dataclass
|
|
from functools import cache
|
|
from typing import Callable
|
|
|
|
import ollama
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from art import env, store
|
|
from art._paths import ROOT
|
|
|
|
PROMPTS = ROOT / "prompts"
|
|
|
|
|
|
def cfg() -> dict:
|
|
return env.config()["llm"]
|
|
|
|
|
|
def _sha(text: str, n: int = 12) -> str:
|
|
return hashlib.sha256(text.encode()).hexdigest()[:n]
|
|
|
|
|
|
# -- prompts -----------------------------------------------------------------
|
|
|
|
@dataclass(frozen=True)
|
|
class Prompt:
|
|
name: str
|
|
version: str
|
|
template: str
|
|
|
|
def render(self, **values: str) -> str:
|
|
text = self.template
|
|
for key, value in values.items():
|
|
text = text.replace("{" + key + "}", value)
|
|
missing = re.findall(r"\{[a-z_]+\}", text)
|
|
if missing:
|
|
raise KeyError(f"prompt {self.name}: unfilled placeholders {missing}")
|
|
return text
|
|
|
|
|
|
@cache
|
|
def prompt(task: str) -> Prompt:
|
|
"""The highest-numbered prompts/<task>_v<N>.txt."""
|
|
files = sorted(PROMPTS.glob(f"{task}_v*.txt"),
|
|
key=lambda p: int(re.search(r"_v(\d+)\.txt$", p.name).group(1)))
|
|
if not files:
|
|
raise FileNotFoundError(f"no prompt file for task {task!r} in {PROMPTS}")
|
|
path = files[-1]
|
|
template = path.read_text()
|
|
number = re.search(r"_v(\d+)\.txt$", path.name).group(1)
|
|
return Prompt(task, f"v{number}+{_sha(template, 8)}", template)
|
|
|
|
|
|
# -- model -------------------------------------------------------------------
|
|
|
|
def modelfile_hash() -> str:
|
|
return _sha((ROOT / cfg()["modelfile"]).read_text())
|
|
|
|
|
|
def model_tag() -> str:
|
|
return f"{cfg()['model']}:{modelfile_hash()}"
|
|
|
|
|
|
def client() -> ollama.Client:
|
|
return ollama.Client(host=os.environ.get("OLLAMA_HOST"), timeout=cfg()["timeout_s"])
|
|
|
|
|
|
def available() -> bool:
|
|
"""True if the Ollama server answers."""
|
|
try:
|
|
ollama.Client(host=os.environ.get("OLLAMA_HOST"), timeout=5).list()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
@cache
|
|
def ensure_model() -> tuple[str, str]:
|
|
"""(tag, digest) of the researcher model, creating it from the Modelfile if missing."""
|
|
tag = model_tag()
|
|
|
|
def digest():
|
|
for m in client().list().models:
|
|
if m.model == tag:
|
|
return m.digest
|
|
return None
|
|
|
|
if digest() is None:
|
|
print(f" creating {tag} from {cfg()['modelfile']}")
|
|
subprocess.run(["ollama", "create", tag, "-f", str(ROOT / cfg()["modelfile"])],
|
|
check=True, capture_output=True, text=True)
|
|
return tag, digest() or ""
|
|
|
|
|
|
# -- calls -------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class Result:
|
|
text: str | None # raw response, or None on a connection error
|
|
parsed: object | None # validated value for structured or parsed calls
|
|
outcome: str # ok, invalid_json, schema_fail, parse_fail, error
|
|
call_id: int
|
|
prompt_version: str
|
|
|
|
|
|
def _validate(raw: str, schema: type[BaseModel]) -> tuple[str, dict | None, dict]:
|
|
"""Classify a structured response. Fenced JSON is accepted but recorded."""
|
|
details: dict = {}
|
|
text = raw.strip()
|
|
fence = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", text, re.S)
|
|
if fence:
|
|
text, details["fenced"] = fence.group(1), True
|
|
try:
|
|
obj = json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
return "invalid_json", None, {**details, "error": str(exc)}
|
|
if not isinstance(obj, dict):
|
|
return "schema_fail", None, {**details, "error": f"top level is {type(obj).__name__}"}
|
|
fields = set(schema.model_fields)
|
|
details.update(missing=sorted(fields - set(obj)), extra=sorted(set(obj) - fields))
|
|
try:
|
|
return "ok", schema.model_validate(obj).model_dump(), details
|
|
except ValidationError as exc:
|
|
details["errors"] = [{"loc": list(e["loc"]), "type": e["type"], "msg": e["msg"]}
|
|
for e in exc.errors()]
|
|
return "schema_fail", None, details
|
|
|
|
|
|
def generate(conn, task: str, values: dict[str, str], *,
|
|
schema: type[BaseModel] | None = None,
|
|
parse: Callable[[str], tuple[str, object, dict]] | None = None,
|
|
item_id: str | None = None, profile=None, run_id: int | None = None,
|
|
options: dict | None = None) -> Result:
|
|
"""Run prompt `task` with `values`, log every attempt, and return the last.
|
|
|
|
With `schema` (a pydantic model), the response must be a JSON object that
|
|
validates; failures are retried cfg()["retries"] times. With `parse`, a
|
|
function raw -> (outcome, value, details) classifies free-text responses.
|
|
"""
|
|
c = cfg()
|
|
tag, digest = ensure_model()
|
|
p = prompt(task)
|
|
text = p.render(**values)
|
|
opts = {"num_ctx": c["num_ctx"], **c["tasks"].get(task, {}), **(options or {})}
|
|
fmt = schema.model_json_schema() if schema and c["enforce_schema"] else None
|
|
attempts = 1 + (c["retries"] if schema else 0)
|
|
result = None
|
|
for attempt in range(1, attempts + 1):
|
|
t0 = time.perf_counter()
|
|
try:
|
|
r = client().generate(model=tag, prompt=text, options=opts, format=fmt,
|
|
think=c["think"])
|
|
except Exception as exc:
|
|
call_id = _log(conn, tag, digest, p, item_id, profile, run_id, opts, fmt, c,
|
|
time.perf_counter() - t0, None, None, None, None, "error",
|
|
{"error": str(exc)}, attempt, None)
|
|
return Result(None, None, "error", call_id, p.version)
|
|
latency = time.perf_counter() - t0
|
|
raw = r.response or ""
|
|
outcome, value, details = "ok", None, {}
|
|
if schema:
|
|
outcome, value, details = _validate(raw, schema)
|
|
elif parse:
|
|
outcome, value, details = parse(raw)
|
|
if r.prompt_eval_count and r.prompt_eval_count >= 0.9 * opts["num_ctx"]:
|
|
details["near_num_ctx"] = True
|
|
call_id = _log(conn, tag, digest, p, item_id, profile, run_id, opts, fmt, c, latency,
|
|
r.prompt_eval_count, r.eval_count, raw, r.done_reason, outcome,
|
|
details, attempt, getattr(r, "thinking", None))
|
|
result = Result(raw, value, outcome, call_id, p.version)
|
|
if outcome == "ok":
|
|
break
|
|
return result
|
|
|
|
|
|
def _log(conn, tag, digest, p, item_id, profile, run_id, opts, fmt, c, latency,
|
|
prompt_tokens, eval_tokens, raw, done_reason, outcome, details, attempt, thinking) -> int:
|
|
cur = conn.execute(
|
|
"INSERT INTO llm_calls (ts, run_id, host, model, modelfile_hash, prompt_name, prompt_version, "
|
|
"item_id, profile_id, profile_version, num_ctx, prompt_tokens, eval_tokens, options, format, "
|
|
"think, latency_s, raw_response, outcome, validation, attempt, model_digest, done_reason, thinking) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
|
(store.now(), run_id, env.host(), tag, modelfile_hash(), p.name, p.version, item_id,
|
|
profile.id if profile else None, profile.version if profile else None, opts["num_ctx"],
|
|
prompt_tokens, eval_tokens, json.dumps(opts), json.dumps(fmt) if fmt else None,
|
|
int(bool(c["think"])), round(latency, 3), raw, outcome,
|
|
json.dumps(details) if details else None, attempt, digest, done_reason, thinking or None))
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
|
|
|
|
# -- stored generations --------------------------------------------------------
|
|
|
|
def _key(task: str, item_id: str, profile) -> tuple:
|
|
return (item_id, task, profile.id if profile else "", profile.version if profile else "",
|
|
model_tag(), modelfile_hash(), prompt(task).version)
|
|
|
|
|
|
def lookup(conn, task: str, item_id: str, profile=None) -> dict | None:
|
|
"""The stored generation for the current model, Modelfile, prompt, and profile versions.
|
|
|
|
Needs no Ollama: the key is computed from the files on disk.
|
|
"""
|
|
row = conn.execute(
|
|
"SELECT text, parsed FROM generations WHERE item_id = ? AND task = ? AND profile_id = ? "
|
|
"AND profile_version = ? AND model = ? AND modelfile_hash = ? AND prompt_version = ?",
|
|
_key(task, item_id, profile)).fetchone()
|
|
if row is None:
|
|
return None
|
|
return {"text": row["text"], "parsed": json.loads(row["parsed"]) if row["parsed"] else None}
|
|
|
|
|
|
def run_task(conn, task: str, item_id: str, values: dict[str, str], *, profile=None,
|
|
run_id: int | None = None, schema: type[BaseModel] | None = None,
|
|
clean: Callable[[str], str] = str.strip) -> tuple[dict | None, bool]:
|
|
"""Stored generation if current, else generate and store. Returns (generation, called).
|
|
|
|
Failed generations are not stored, so they are retried on the next run;
|
|
their raw responses are in llm_calls.
|
|
"""
|
|
have = lookup(conn, task, item_id, profile)
|
|
if have is not None:
|
|
return have, False
|
|
r = generate(conn, task, values, schema=schema, item_id=item_id, profile=profile, run_id=run_id)
|
|
if r.outcome != "ok":
|
|
return None, True
|
|
gen = {"text": clean(r.text) if not schema else None, "parsed": r.parsed}
|
|
conn.execute("INSERT OR REPLACE INTO generations VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
|
(*_key(task, item_id, profile)[:7], cfg()["num_ctx"], gen["text"],
|
|
json.dumps(gen["parsed"]) if gen["parsed"] is not None else None,
|
|
r.call_id, store.now()))
|
|
conn.commit()
|
|
return gen, True
|