provekit 0.1.0__py3-none-any.whl

provekit/routers/run.py

@@ -0,0 +1,203 @@
+"""Run a request (prompt | tool | agent), streaming unified events, evaluating any
+assertions, and persisting to history. Also batch datasets (run over N input rows)."""
+import json
+
+import anyio
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from ..database import SessionLocal, get_db
+from ..models import Environment, Run, Workspace, iso_utc
+from ..services import assertions as assertion_engine
+from ..services import dispatch, otel
+from ..services.limits import check_rate, clamp_max_tokens, enforce_dataset_size, prune_runs
+from ..services.masking import mask_body, mask_headers
+from ..services.workspace import current_workspace
+
+router = APIRouter(prefix="/api", tags=["run"])
+
+
+class RunPayload(BaseModel):
+    request: dict
+    variables: dict = {}
+    save: bool = True
+
+
+class DatasetPayload(BaseModel):
+    request: dict
+    rows: list[dict] = []
+
+
+def _active_vars(db: Session, ws_id: int) -> dict:
+    e = db.query(Environment).filter(Environment.workspace_id == ws_id, Environment.is_active.is_(True)).first()
+    return dict(e.variables or {}) if e else {}
+
+
+def _label(req: dict) -> str:
+    t = req.get("type")
+    if t == "prompt":
+        return f"{req.get('model', 'prompt')}: {(req.get('user') or '')[:60]}"
+    if t == "tool":
+        return f"tool: {req.get('tool', '?')}"
+    if t == "agent":
+        return f"{req.get('method', 'POST')} {req.get('path', '')}"
+    return t or "run"
+
+
+def _sanitize(req: dict) -> dict:
+    """Strip/mask secrets before persisting to run history (returned by GET /runs/{id})."""
+    out = {k: v for k, v in req.items() if k != "api_key"}
+    if isinstance(out.get("headers"), dict):
+        out["headers"] = mask_headers(out["headers"])
+    if out.get("body") is not None:  # agent bodies can carry tokens/passwords
+        out["body"] = mask_body(out["body"])
+    return out
+
+
+def _new_acc():
+    # status starts as "interrupted" — only a `done` event sets the real outcome, so a
+    # client disconnect mid-stream persists honestly instead of as "completed".
+    return {"parts": [], "output": None, "meta": {}, "events": [], "status": "interrupted", "dur": 0, "err": ""}
+
+
+def _apply(acc, ev):
+    et = ev["type"]
+    if et == "delta":
+        acc["parts"].append(ev.get("text", ""))
+    elif et == "result":
+        acc["output"], acc["meta"] = ev.get("data"), ev.get("meta", {})
+    elif et == "node":
+        acc["events"].append(ev.get("data"))
+    elif et == "error":
+        acc["err"] = ev.get("error", "")
+    elif et == "done":
+        acc["status"], acc["dur"] = ev.get("status", "completed"), ev.get("duration_ms", 0)
+
+
+def _run_dict(acc):
+    result = {"text": "".join(acc["parts"]) or None, "output": acc["output"], "meta": acc["meta"]}
+    if acc["events"]:
+        result["events"] = acc["events"]
+    return {"result": result, "status": acc["status"], "duration_ms": acc["dur"], "events": acc["events"], "error": acc["err"]}
+
+
+async def _collect(db, req, variables, ws_id):
+    acc = _new_acc()
+    async for ev in dispatch.run(db, req, variables, ws_id):
+        _apply(acc, ev)
+    return _run_dict(acc)
+
+
+async def _evaluate(db, assertions, rd, ws_id):
+    """Offload assertion evaluation to a thread — the llm_judge assertion runs its own
+    event loop (asyncio.run), which can't happen inside this request's loop."""
+    if not assertions:
+        return []
+    return await anyio.to_thread.run_sync(assertion_engine.evaluate, db, assertions, rd, ws_id)
+
+
+async def _persist_async(db, req, rd, asserts, ws_id):
+    """Persist off the event loop — a blocking DB commit (esp. a contended SQLite write)
+    must not stall every other in-flight stream sharing this loop."""
+    await anyio.to_thread.run_sync(_persist, db, req, rd, asserts, ws_id)
+
+
+def _persist(db, req, rd, asserts, ws_id):
+    result = dict(rd["result"])
+    if asserts:
+        result["assertions"] = asserts
+    try:
+        row = Run(workspace_id=ws_id, type=req.get("type", "?"), label=_label(req), request=_sanitize(req),
+                  result=result, status=rd["status"], duration_ms=rd["duration_ms"], error=rd["error"])
+        db.add(row); db.commit()
+        otel.emit_run(row)  # best-effort mirror to an OTLP collector if configured
+        prune_runs(db, ws_id)
+    except Exception:
+        db.rollback()
+
+
+@router.post("/run/stream")
+def run_stream(payload: RunPayload, db: Session = Depends(get_db), ws: Workspace = Depends(check_rate)):
+    variables = {**_active_vars(db, ws.id), **(payload.variables or {})}
+    req = payload.request
+    clamp_max_tokens(req)
+    assertions = req.get("assertions") or []
+    ws_id = ws.id
+
+    async def live():
+        # Fresh session: FastAPI tears down the request's Depends(get_db) before this
+        # generator runs (dependencies-with-yield exit when the endpoint returns).
+        session = SessionLocal()
+        acc = _new_acc()
+        asserts: list = []
+        try:
+            async for ev in dispatch.run(session, req, variables, ws_id):
+                _apply(acc, ev)
+                yield f"data: {json.dumps(ev)}\n\n"
+            rd = _run_dict(acc)
+            asserts = await _evaluate(session, assertions, rd, ws_id)
+            if asserts:
+                yield f"data: {json.dumps({'type': 'assert', 'results': asserts})}\n\n"
+            yield "data: [DONE]\n\n"
+        finally:
+            # Runs also on client disconnect (GeneratorExit) — history keeps the partial run.
+            try:
+                if payload.save:
+                    await _persist_async(session, req, _run_dict(acc), asserts, ws_id)
+            finally:
+                session.close()
+
+    return StreamingResponse(live(), media_type="text/event-stream",
+                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+
+
+@router.post("/run")
+async def run_once(payload: RunPayload, db: Session = Depends(get_db), ws: Workspace = Depends(check_rate)):
+    clamp_max_tokens(payload.request)
+    variables = {**_active_vars(db, ws.id), **(payload.variables or {})}
+    rd = await _collect(db, payload.request, variables, ws.id)
+    asserts = await _evaluate(db, payload.request.get("assertions") or [], rd, ws.id)
+    if payload.save:
+        await _persist_async(db, payload.request, rd, asserts, ws.id)
+    return {"result": rd["result"], "status": rd["status"], "duration_ms": rd["duration_ms"], "assertions": asserts}
+
+
+@router.post("/dataset/run")
+async def dataset_run(payload: DatasetPayload, db: Session = Depends(get_db), ws: Workspace = Depends(check_rate)):
+    enforce_dataset_size(len(payload.rows))
+    clamp_max_tokens(payload.request)
+    assertions = payload.request.get("assertions") or []
+    base = _active_vars(db, ws.id)
+    rows = []
+    for i, row in enumerate(payload.rows):
+        variables = {**base, **(row.get("variables") or {})}
+        rd = await _collect(db, payload.request, variables, ws.id)
+        asserts = await _evaluate(db, assertions, rd, ws.id)
+        passed = all(a["ok"] for a in asserts) if asserts else rd["status"] == "completed"
+        rows.append({
+            "name": row.get("name") or f"row {i + 1}", "status": rd["status"],
+            "text": rd["result"].get("text"), "output": rd["result"].get("output"),
+            "assertions": asserts, "pass": passed, "duration_ms": rd["duration_ms"],
+        })
+    return {"rows": rows, "summary": {"passed": sum(1 for r in rows if r["pass"]), "total": len(rows)}}
+
+
+@router.get("/runs")
+def list_runs(limit: int = 30, db: Session = Depends(get_db), ws: Workspace = Depends(current_workspace)):
+    limit = max(1, min(limit, 200))  # bound the page (SQLite treats LIMIT -1 as unlimited)
+    rows = db.query(Run).filter(Run.workspace_id == ws.id).order_by(Run.id.desc()).limit(limit).all()
+    return [{"id": r.id, "type": r.type, "label": r.label, "status": r.status,
+             "duration_ms": r.duration_ms, "created_at": iso_utc(r.created_at)}
+            for r in rows]
+
+
+@router.get("/runs/{rid}")
+def get_run(rid: int, db: Session = Depends(get_db), ws: Workspace = Depends(current_workspace)):
+    r = db.get(Run, rid)
+    if not r or r.workspace_id != ws.id:
+        raise HTTPException(404, "Run not found")
+    return {"id": r.id, "type": r.type, "label": r.label, "request": r.request,
+            "result": r.result, "status": r.status, "duration_ms": r.duration_ms,
+            "error": r.error, "created_at": iso_utc(r.created_at)}

provekit/routers/runtime.py

@@ -0,0 +1,103 @@
+"""Public runtime — invoke a deployed flow by slug with an API key. No session auth.
+
+POST /v1/d/{slug}          → JSON {output, status, run_id, duration_ms}
+POST /v1/d/{slug}?stream=1 → SSE passthrough of the flow's node/delta events
+
+This executes user-defined outbound calls, so it depends on the P0 safeguards
+(credential-override block, SSRF guard, secret masking) already in place.
+"""
+import json
+import time
+
+import anyio
+from fastapi import APIRouter, Header, HTTPException, Request
+from fastapi.responses import StreamingResponse
+
+from ..database import SessionLocal
+from ..models import Deployment, Run
+from ..services import deploy
+from ..services.limits import _window
+from ..config import get_settings
+
+router = APIRouter(prefix="/v1/d", tags=["runtime"])
+
+
+def _resolve(session, slug: str, api_key: str | None) -> Deployment:
+    d = (session.query(Deployment)
+         .filter(Deployment.slug == slug, Deployment.active.is_(True))
+         .order_by(Deployment.version.desc()).first())
+    if not d:
+        # Distinguish "never existed / deactivated" without leaking which.
+        any_row = session.query(Deployment).filter(Deployment.slug == slug).first()
+        raise HTTPException(410 if any_row else 404, "Deployment not available")
+    if not api_key or not deploy.verify_key(api_key, d.api_key_hash):
+        raise HTTPException(403, "Invalid API key")
+    return d
+
+
+def _rate_ok(deployment_id: int) -> bool:
+    limit = get_settings().rate_limit_per_min
+    if limit <= 0:
+        return True
+    key = f"rl:dep:{deployment_id}:{int(time.time()) // 60}"
+    return _window().hit(key, 60) <= limit
+
+
+@router.post("/{slug}")
+async def invoke(slug: str, request: Request, stream: bool = False,
+                 x_api_key: str | None = Header(default=None)):
+    body = {}
+    try:
+        body = await request.json()
+    except Exception:
+        pass
+
+    session = SessionLocal()
+    try:
+        d = _resolve(session, slug, x_api_key)
+        if not _rate_ok(d.id):
+            raise HTTPException(429, "Rate limit exceeded")
+        snapshot, ws_id, dep_id = d.snapshot, d.workspace_id, d.id
+    except HTTPException:
+        session.close()
+        raise
+
+    if stream:
+        async def gen():
+            try:
+                async for ev in deploy.run_snapshot(session, snapshot, body, ws_id):
+                    yield f"data: {json.dumps(ev)}\n\n"
+                yield "data: [DONE]\n\n"
+            finally:
+                session.close()
+        return StreamingResponse(gen(), media_type="text/event-stream",
+                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+
+    try:
+        t0 = time.monotonic()
+        timeout = get_settings().deployment_timeout_s
+        timed_out = False
+        events = []
+        with anyio.move_on_after(timeout if timeout > 0 else None) as scope:
+            async for ev in deploy.run_snapshot(session, snapshot, body, ws_id):
+                events.append(ev)
+        timed_out = scope.cancelled_caught
+        result = deploy.collect_output(events)
+        if timed_out:
+            result["status"] = "failed"
+            result["error"] = f"deployment invocation exceeded {timeout}s timeout"
+        dur = round((time.monotonic() - t0) * 1000)
+
+        def _save():
+            session.add(Run(workspace_id=ws_id, deployment_id=dep_id, type="deployment",
+                            label=f"deploy {slug}", request={"input": body},
+                            result={"output": result["output"]}, status=result["status"],
+                            duration_ms=dur, error=result["error"]))
+            session.commit()
+        await anyio.to_thread.run_sync(_save)  # keep the blocking write off the loop
+        if timed_out:
+            raise HTTPException(504, result["error"])
+        return {"output": result["output"], "status": result["status"],
+                "run_id": result["run_id"], "duration_ms": dur}
+    finally:
+        session.close()

provekit/routers/traces.py

@@ -0,0 +1,65 @@
+"""OTLP/HTTP trace ingest — point any OpenTelemetry GenAI exporter at /v1/traces and
+its LLM/agent spans show up in ProveKit's history, no SDK swap required.
+
+Auth: a workspace ingest key via `Authorization: Bearer <key>` (what real exporters can
+send), or a session cookie for interactive/local use. Mint the key from
+POST /api/workspace/ingest-key."""
+from fastapi import APIRouter, Depends, Header, HTTPException, Request
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from ..config import get_settings
+from ..database import get_db
+from ..models import Run, Workspace
+from ..services import deploy, otel
+from ..services.workspace import current_workspace
+
+router = APIRouter(prefix="/v1", tags=["traces"])
+ws_router = APIRouter(prefix="/api/workspace", tags=["workspace"])
+
+
+def _resolve_ingest_ws(db: Session, request: Request, authorization: str | None) -> Workspace:
+    """Bearer ingest key first (exporters), else fall back to the session cookie."""
+    if authorization and authorization.lower().startswith("bearer "):
+        key = authorization[7:].strip()
+        h = deploy.hash_key(key)
+        ws = db.query(Workspace).filter(Workspace.ingest_key_hash == h).first()
+        if ws:
+            return ws
+        raise HTTPException(403, "Invalid ingest key")
+    # No bearer key: only allowed via a logged-in session (or local mode's default user).
+    from ..services.auth import get_current_user
+    from ..services.workspace import get_or_create_default_workspace
+    user = get_current_user(request, db)
+    return get_or_create_default_workspace(db, user)
+
+
+@router.post("/traces")
+async def ingest_traces(request: Request, db: Session = Depends(get_db),
+                        authorization: str | None = Header(default=None)):
+    """Accept an OTLP ExportTraceServiceRequest (JSON) and persist gen_ai spans as runs.
+    Returns the OTLP success shape so standard exporters are satisfied."""
+    ws = _resolve_ingest_ws(db, request, authorization)
+    try:
+        payload = await request.json()
+    except Exception:
+        return {"partialSuccess": {"rejectedSpans": 0, "errorMessage": "invalid JSON"}}
+    rows = otel.ingest(payload)
+    for kw in rows:
+        db.add(Run(workspace_id=ws.id, **kw))
+    if rows:
+        db.commit()
+    return {"partialSuccess": {}}
+
+
+class _KeyOut(BaseModel):
+    ingest_key: str
+
+
+@ws_router.post("/ingest-key", response_model=_KeyOut)
+def rotate_ingest_key(db: Session = Depends(get_db), ws: Workspace = Depends(current_workspace)):
+    """Mint/rotate this workspace's OTLP ingest key (shown once; stored hashed)."""
+    plaintext, key_hash = deploy.new_api_key()
+    ws.ingest_key_hash = key_hash
+    db.commit()
+    return {"ingest_key": plaintext}

provekit/routers/usage.py

@@ -0,0 +1,42 @@
+"""Per-workspace token usage aggregation — the foundation for quotas/billing later.
+Normalizes the different providers' usage shapes (OpenAI prompt/completion, Anthropic
+and Responses input/output) into one in/out count."""
+from datetime import datetime, timedelta, timezone
+
+from fastapi import APIRouter, Depends
+from sqlalchemy.orm import Session
+
+from ..database import get_db
+from ..models import Run, Workspace
+from ..services.workspace import current_workspace
+
+router = APIRouter(prefix="/api/usage", tags=["usage"])
+
+
+def _tokens(usage: dict) -> tuple[int, int]:
+    if not isinstance(usage, dict):
+        return 0, 0
+    tin = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
+    tout = usage.get("output_tokens") or usage.get("completion_tokens") or 0
+    return int(tin or 0), int(tout or 0)
+
+
+@router.get("")
+def usage(days: int = 30, db: Session = Depends(get_db), ws: Workspace = Depends(current_workspace)):
+    """Totals + per-model breakdown of runs and tokens over the last `days`."""
+    since = datetime.now(timezone.utc) - timedelta(days=max(1, days))
+    rows = (db.query(Run).filter(Run.workspace_id == ws.id, Run.created_at >= since.replace(tzinfo=None)).all())
+    total_in = total_out = 0
+    by_model: dict[str, dict] = {}
+    for r in rows:
+        meta = (r.result or {}).get("meta") or {}
+        model = meta.get("model") or "—"
+        tin, tout = _tokens(meta.get("usage"))
+        total_in += tin; total_out += tout
+        m = by_model.setdefault(model, {"runs": 0, "tokens_in": 0, "tokens_out": 0})
+        m["runs"] += 1; m["tokens_in"] += tin; m["tokens_out"] += tout
+    return {
+        "window_days": days, "runs": len(rows),
+        "tokens_in": total_in, "tokens_out": total_out, "tokens_total": total_in + total_out,
+        "by_model": [{"model": k, **v} for k, v in sorted(by_model.items(), key=lambda kv: -kv[1]["tokens_in"] - kv[1]["tokens_out"])],
+    }

provekit/services/assertions.py

@@ -0,0 +1,135 @@
+"""Assertions / evals — check a run's result against a list of assertions. Types:
+contains · equals · regex · json_path · json_schema · tool_called · latency_lt · llm_judge.
+Mirrors the eval assertions in tools like Reticle/LangSmith."""
+from __future__ import annotations
+
+import json
+import re
+
+from jsonschema import ValidationError
+from jsonschema import validate as js_validate
+
+from ..models import Connection
+from .providers import llm
+
+# re.search holds the GIL, so a catastrophic-backtracking pattern (ReDoS) would freeze the
+# whole process — and a worker thread can't preempt it. Since stdlib re has no timeout, we
+# reject the common nested-quantifier signature and cap the searched text instead. This is a
+# heuristic guard; an interruptible engine (the `regex` module's timeout, or RE2) is the
+# complete fix if fully-untrusted patterns ever need to run.
+_MAX_REGEX_INPUT = 100_000
+_NESTED_QUANT = re.compile(r"\([^()]*[*+][^()]*\)\s*[*+{]")
+
+
+def _regex_search(pattern: str, text: str):
+    if len(pattern) > 2000:
+        raise ValueError("regex pattern too long")
+    if _NESTED_QUANT.search(pattern):
+        raise ValueError("regex rejected: nested quantifier (possible ReDoS)")
+    return re.search(pattern, text[:_MAX_REGEX_INPUT])
+
+
+def _text_of(result: dict) -> str:
+    if result.get("text"):
+        return result["text"]
+    out = result.get("output")
+    if isinstance(out, str):
+        return out
+    return json.dumps(out) if out is not None else ""
+
+
+def get_path(obj, path: str):
+    cur = obj
+    for part in str(path).split("."):
+        if part == "":
+            continue
+        if isinstance(cur, list):
+            try:
+                cur = cur[int(part)]
+            except (ValueError, IndexError):
+                return None
+        elif isinstance(cur, dict):
+            cur = cur.get(part)
+        else:
+            return None
+    return cur
+
+
+def _llm_judge(db, a: dict, text: str, workspace_id=None):
+    conn = db.get(Connection, a.get("connection_id")) if a.get("connection_id") else None
+    if conn and workspace_id is not None and conn.workspace_id != workspace_id:
+        conn = None  # tenancy: don't judge with another workspace's connection
+    if not conn:
+        return False, "no judge connection configured"
+    cfg = conn.config or {}
+    model = a.get("model") or (cfg.get("models") or ["gpt-4o-mini"])[0]
+    system = ("You are a strict evaluator. Given CRITERIA and OUTPUT, decide if the output "
+              "satisfies the criteria. Reply with exactly PASS or FAIL on the first line, "
+              "then a one-sentence reason.")
+    user = f"CRITERIA:\n{a.get('criteria', '')}\n\nOUTPUT:\n{text[:4000]}"
+    # Sync bridge (asyncio.run): evaluate() runs in a worker thread from async endpoints
+    # and directly in the CLI — both have no running loop, so this is safe.
+    reply = llm.collect_text_sync(provider=cfg.get("provider", "openai"), base_url=cfg.get("base_url", ""),
+                                  api_key=cfg.get("api_key", ""), model=model, system=system,
+                                  messages=[{"role": "user", "content": user}], temperature=0, max_tokens=200).strip()
+    return reply.upper().startswith("PASS"), reply[:160]
+
+
+def evaluate(db, assertions: list, run: dict, workspace_id=None) -> list[dict]:
+    """run = {result:{text,output,meta}, status, duration_ms, events}."""
+    result = run.get("result", {}) or {}
+    text = _text_of(result)
+    output = result.get("output")
+    meta = result.get("meta", {}) or {}
+    dur = run.get("duration_ms", 0)
+    events = run.get("events") or result.get("events") or []
+    out = []
+    for a in assertions or []:
+        t = a.get("type")
+        ok, detail = False, ""
+        try:
+            if t == "contains":
+                ok = str(a.get("value", "")) in text
+                detail = f"output {'contains' if ok else 'missing'} “{a.get('value')}”"
+            elif t == "equals":
+                target = get_path(output, a["path"]) if a.get("path") else text
+                ok = str(target) == str(a.get("value"))
+                detail = f"{target!r} {'==' if ok else '!='} {a.get('value')!r}"
+            elif t == "regex":
+                ok = _regex_search(a.get("value", ""), text) is not None
+                detail = f"/{a.get('value')}/ {'matched' if ok else 'no match'}"
+            elif t == "json_path":
+                val = get_path(output, a.get("path", ""))
+                if a.get("value", "") != "":
+                    ok = str(val) == str(a["value"]); detail = f"{a.get('path')} = {val!r}"
+                else:
+                    ok = val is not None; detail = f"{a.get('path')} {'exists' if ok else 'missing'} (={val!r})"
+            elif t == "json_schema":
+                schema = a.get("schema") if isinstance(a.get("schema"), dict) else json.loads(a.get("schema") or "{}")
+                try:
+                    js_validate(output, schema); ok = True; detail = "matches schema"
+                except ValidationError as ve:
+                    ok = False; detail = str(ve.message)[:140]
+            elif t == "tool_called":
+                # Match structured fields only — substring-matching the raw event JSON
+                # passed whenever the name appeared anywhere (e.g. inside prompt text).
+                name = a.get("value", "")
+                called = {meta.get("tool")}
+                for e in events:
+                    if isinstance(e, dict):
+                        called.update(str(e[k]) for k in ("tool", "tool_name", "name") if e.get(k))
+                        for tc in e.get("tool_calls") or []:
+                            if isinstance(tc, dict):
+                                called.add(str(tc.get("name") or (tc.get("function") or {}).get("name")))
+                ok = name in {c for c in called if c}
+                detail = f"tool '{name}' {'called' if ok else 'not called'}"
+            elif t == "latency_lt":
+                ok = dur < float(a.get("value", 0)); detail = f"{dur}ms {'<' if ok else '≥'} {a.get('value')}ms"
+            elif t == "llm_judge":
+                ok, detail = _llm_judge(db, a, text, workspace_id)
+            else:
+                detail = f"unknown assertion type: {t}"
+        except Exception as exc:
+            ok, detail = False, f"error: {exc}"[:140]
+        out.append({"type": t, "name": a.get("name") or t, "ok": ok, "detail": detail})
+    return out