downshift 0.1.0.dev0__py3-none-any.whl

downshift/runner.py

@@ -0,0 +1,312 @@
+"""Run eval sets against models and store scored results.
+
+For each call site and model, the runner renders the site's prompt with every
+eval case, calls the model with the site's own settings, scores the output and
+appends one JSON line per case to results/<slug>/<model>.jsonl.
+
+Runs are resumable: cases that already have a scored row are skipped, and rows
+that failed (model or judge error) are retried. The last row for a case wins.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import asdict, dataclass, replace
+from dataclasses import fields as dataclass_fields
+from pathlib import Path
+from typing import Any
+
+from downshift.evals import EvalCase, EvalSet, graded_fields, render_messages, slug_for
+from downshift.llm import LLMClient, LLMError
+from downshift.schema import CallSite
+from downshift.scorer import Judge, score_case
+
+RESULTS_SUFFIX = ".jsonl"
+WARMUP_MESSAGES: list[dict[str, str]] = [{"role": "user", "content": "Reply with the word OK."}]
+
+
+def model_filename(model: str) -> str:
+    """Filesystem-safe model name: `qwen2.5:7b` -> `qwen2.5-7b`."""
+    return re.sub(r"[^A-Za-z0-9._-]+", "-", model).strip("-") or "model"
+
+
+def results_path(results_dir: Path, site_id: str, model: str) -> Path:
+    return results_dir / slug_for(site_id) / f"{model_filename(model)}{RESULTS_SUFFIX}"
+
+
+@dataclass(frozen=True)
+class ResultRow:
+    """One eval case run on one model."""
+
+    case_id: str
+    model: str
+    output: str = ""
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    latency_s: float = 0.0
+    score: float | None = None
+    passed: bool | None = None
+    detail: str = ""
+    judge_score: int | None = None
+    error: str | None = None
+    judge_model: str | None = None
+
+    @property
+    def ok(self) -> bool:
+        """Scored without errors."""
+        return self.error is None and self.score is not None
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["latency_s"] = round(self.latency_s, 3)
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> ResultRow:
+        known = {f.name for f in dataclass_fields(cls)}
+        return cls(**{k: v for k, v in data.items() if k in known})
+
+
+def load_results(path: Path) -> dict[str, ResultRow]:
+    """Rows by case id. Unreadable lines are skipped; the last row for a case wins."""
+    rows: dict[str, ResultRow] = {}
+    if not path.is_file():
+        return rows
+    for raw in path.read_text(encoding="utf-8").splitlines():
+        if not raw.strip():
+            continue
+        try:
+            row = ResultRow.from_dict(json.loads(raw))
+        except (json.JSONDecodeError, TypeError, AttributeError):
+            continue
+        rows[row.case_id] = row
+    return rows
+
+
+def append_row(path: Path, row: ResultRow) -> None:
+    """Append one row, starting a new line if a previous write was cut off."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    prefix = ""
+    if path.is_file() and path.stat().st_size:
+        with path.open("rb") as fh:
+            fh.seek(-1, os.SEEK_END)
+            if fh.read(1) != b"\n":
+                prefix = "\n"
+    with path.open("a", encoding="utf-8") as fh:
+        fh.write(prefix + json.dumps(row.to_dict(), ensure_ascii=False) + "\n")
+
+
+@dataclass(frozen=True)
+class RunSummary:
+    """Totals for one call site on one model, over the cases in scope."""
+
+    site_id: str
+    model: str
+    cases: int
+    scored: int
+    passed: int
+    errors: int
+    new: int
+    mean_score: float | None
+    avg_latency_s: float | None
+    avg_prompt_tokens: float | None
+    avg_completion_tokens: float | None
+
+    @property
+    def pass_rate(self) -> float | None:
+        return self.passed / self.scored if self.scored else None
+
+
+def _mean(values: Sequence[float]) -> float | None:
+    return sum(values) / len(values) if values else None
+
+
+def summarize_rows(
+    site_id: str,
+    model: str,
+    case_ids: Sequence[str],
+    rows: Mapping[str, ResultRow],
+    new: int = 0,
+) -> RunSummary:
+    """Summarize the rows for these case ids (rows for other cases are ignored)."""
+    selected = [rows[c] for c in case_ids if c in rows]
+    ok = [r for r in selected if r.ok]
+    return RunSummary(
+        site_id=site_id,
+        model=model,
+        cases=len(case_ids),
+        scored=len(ok),
+        passed=sum(1 for r in ok if r.passed),
+        errors=sum(1 for r in selected if r.error is not None),
+        new=new,
+        mean_score=_mean([r.score or 0.0 for r in ok]),
+        avg_latency_s=_mean([r.latency_s for r in ok]),
+        avg_prompt_tokens=_mean([float(r.prompt_tokens) for r in ok]),
+        avg_completion_tokens=_mean([float(r.completion_tokens) for r in ok]),
+    )
+
+
+class Runner:
+    """Runs eval sets on models and appends scored rows to results files."""
+
+    def __init__(
+        self,
+        client: LLMClient,
+        *,
+        results_dir: Path,
+        judge: Judge | None = None,
+        limit: int | None = None,
+        warmup: bool = True,
+        on_case: Callable[[ResultRow], None] | None = None,
+    ) -> None:
+        self.client = client
+        self.results_dir = results_dir
+        self.judge = judge
+        self.limit = limit
+        self.warmup = warmup
+        self.on_case = on_case
+        self._warmed: set[str] = set()
+
+    def warm_up(self, model: str, client: LLMClient | None = None) -> None:
+        """One untimed call so model load time does not count as latency.
+
+        LLMError propagates: a model that cannot answer this cannot run evals.
+        """
+        if not self.warmup or model in self._warmed:
+            return
+        (client or self.client).complete(model, WARMUP_MESSAGES, temperature=0.0, max_tokens=5)
+        self._warmed.add(model)
+
+    def cases_for(self, eval_set: EvalSet) -> list[EvalCase]:
+        return eval_set.cases[: self.limit] if self.limit else list(eval_set.cases)
+
+    def run_site(self, site: CallSite, eval_set: EvalSet, model: str) -> RunSummary:
+        cases = self.cases_for(eval_set)
+        path = results_path(self.results_dir, site.id, model)
+        rows = load_results(path)
+        todo = [c for c in cases if not (c.id in rows and rows[c.id].ok)]
+        if todo:
+            self.warm_up(model)
+            if self.judge is not None and any(c.grading == "judge" for c in todo):
+                self.warm_up(self.judge.model, self.judge.client)
+        fields = graded_fields(site)
+        for case in todo:
+            row = self._run_case(site, eval_set, case, model, fields)
+            append_row(path, row)
+            rows[case.id] = row
+            if self.on_case is not None:
+                self.on_case(row)
+        return summarize_rows(site.id, model, [c.id for c in cases], rows, new=len(todo))
+
+    def _run_case(
+        self,
+        site: CallSite,
+        eval_set: EvalSet,
+        case: EvalCase,
+        model: str,
+        fields: list[str] | None,
+    ) -> ResultRow:
+        messages = render_messages(site, eval_set.inputs_for(case))
+        temperature = float(site.temperature) if site.temperature is not None else 0.0
+        try:
+            completion = self.client.complete(
+                model,
+                messages,
+                temperature=temperature,
+                max_tokens=site.max_tokens,
+                json_mode=site.output_format == "json",
+            )
+        except LLMError as exc:
+            return ResultRow(case.id, model, error=str(exc))
+        row = ResultRow(
+            case.id,
+            model,
+            output=completion.text,
+            prompt_tokens=completion.prompt_tokens,
+            completion_tokens=completion.completion_tokens,
+            latency_s=completion.latency_s,
+        )
+        try:
+            score = score_case(
+                case.grading,
+                case.expected,
+                completion.text,
+                fields=fields,
+                prompt_messages=messages,
+                judge=self.judge,
+            )
+        except (LLMError, ValueError) as exc:
+            return replace(row, error=f"scoring failed: {exc}")
+        return replace(
+            row,
+            score=score.value,
+            passed=score.passed,
+            detail=score.detail,
+            judge_score=score.judge_score,
+            judge_model=self.judge.model if case.grading == "judge" and self.judge else None,
+        )
+
+
+def rescore_site(
+    site: CallSite,
+    eval_set: EvalSet,
+    model: str,
+    *,
+    results_dir: Path,
+    judge: Judge,
+    on_case: Callable[[ResultRow], None] | None = None,
+) -> tuple[RunSummary, RunSummary]:
+    """Re-grade saved outputs of a judge-graded site with `judge`. Returns (before, after).
+
+    Model outputs are never regenerated. Rows already graded by this judge are
+    skipped, rows whose model call failed (no output) are left for `run`, and
+    judge failures become error rows that are retried on the next rescore.
+    """
+    path = results_path(results_dir, site.id, model)
+    rows = load_results(path)
+    case_ids = [c.id for c in eval_set.cases]
+    before = summarize_rows(site.id, model, case_ids, rows)
+    new = 0
+    for case in eval_set.cases:
+        row = rows.get(case.id)
+        if row is None or case.grading != "judge":
+            continue
+        if row.error is not None and not row.error.startswith("scoring failed"):
+            continue
+        if row.ok and row.judge_model == judge.model:
+            continue
+        messages = render_messages(site, eval_set.inputs_for(case))
+        try:
+            score = score_case(
+                case.grading, case.expected, row.output, prompt_messages=messages, judge=judge
+            )
+        except (LLMError, ValueError) as exc:
+            updated = replace(
+                row,
+                score=None,
+                passed=None,
+                detail="",
+                judge_score=None,
+                judge_model=judge.model,
+                error=f"scoring failed: {exc}",
+            )
+        else:
+            updated = replace(
+                row,
+                score=score.value,
+                passed=score.passed,
+                detail=score.detail,
+                judge_score=score.judge_score,
+                judge_model=judge.model,
+                error=None,
+            )
+        append_row(path, updated)
+        rows[case.id] = updated
+        new += 1
+        if on_case is not None:
+            on_case(updated)
+    after = summarize_rows(site.id, model, case_ids, rows, new=new)
+    return before, after

downshift/scanner.py

@@ -0,0 +1,310 @@
+"""Find LLM call sites in Python source code using the ast module."""
+
+from __future__ import annotations
+
+import ast
+import fnmatch
+import os
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+from downshift.config import ScanConfig
+from downshift.resolve import Ctx, FunctionNode, Module, ModuleIndex, Resolver, keyword_arg
+from downshift.schema import CallSite, ModelRef, PromptMessage, ScanResult
+
+SKIP_DIRS = frozenset(
+    {
+        ".venv",
+        "venv",
+        "env",
+        ".git",
+        "node_modules",
+        "__pycache__",
+        "site-packages",
+        ".tox",
+        ".nox",
+        ".mypy_cache",
+        ".pytest_cache",
+        ".ruff_cache",
+        "build",
+        "dist",
+        ".downshift",
+    }
+)
+
+# (attribute suffix, api name, only counts if a model= keyword is present)
+API_PATTERNS: tuple[tuple[tuple[str, ...], str, bool], ...] = (
+    (("chat", "completions", "create"), "openai.chat.completions", False),
+    (("chat", "completions", "parse"), "openai.chat.completions.parse", False),
+    (("responses", "create"), "openai.responses", True),
+    (("messages", "create"), "anthropic.messages", True),
+)
+
+MAX_TOKEN_KWARGS = ("max_tokens", "max_completion_tokens", "max_output_tokens")
+JSON_FORMAT_TYPES = frozenset({"json_object", "json_schema"})
+
+
+# --- public API ---------------------------------------------------------------
+
+
+def scan_path(root: Path, scan_config: ScanConfig | None = None) -> ScanResult:
+    """Scan a directory (or a single .py file) for LLM call sites."""
+    scan_config = scan_config or ScanConfig()
+    if not root.exists():
+        raise FileNotFoundError(f"path does not exist: {root}")
+    root = root.resolve()
+
+    modules: list[Module] = []
+    warnings: list[str] = []
+    files_scanned = 0
+    for path, rel in iter_python_files(root, scan_config):
+        files_scanned += 1
+        try:
+            source = path.read_text(encoding="utf-8")
+        except (OSError, UnicodeDecodeError) as exc:
+            warnings.append(f"{rel}: could not read file, skipped ({exc})")
+            continue
+        try:
+            tree = ast.parse(source, filename=rel)
+        except SyntaxError as exc:
+            warnings.append(f"{rel}:{exc.lineno}: syntax error, skipped ({exc.msg})")
+            continue
+        modules.append(Module.build(tree, rel))
+
+    return ScanResult(
+        root=root.name,
+        files_scanned=files_scanned,
+        call_sites=_scan_modules(modules),
+        warnings=warnings,
+    )
+
+
+def scan_source(source: str, rel: str = "<string>") -> list[CallSite]:
+    """Scan one module's source text. Handy for tests and editor integrations."""
+    return scan_module(ast.parse(source, filename=rel), rel)
+
+
+def scan_module(tree: ast.Module, rel: str) -> list[CallSite]:
+    """Return the call sites in one parsed module, resolving names within it only."""
+    return _scan_modules([Module.build(tree, rel)])
+
+
+def iter_python_files(root: Path, scan_config: ScanConfig) -> Iterator[tuple[Path, str]]:
+    """Yield (absolute path, posix path relative to root) for files to scan."""
+    if root.is_file():
+        yield root, root.name
+        return
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames[:] = sorted(d for d in dirnames if d not in SKIP_DIRS and not d.startswith("."))
+        for name in sorted(filenames):
+            if not name.endswith(".py"):
+                continue
+            path = Path(dirpath) / name
+            rel = path.relative_to(root).as_posix()
+            if not any(fnmatch.fnmatch(rel, pattern) for pattern in scan_config.include):
+                continue
+            if any(fnmatch.fnmatch(rel, pattern) for pattern in scan_config.exclude):
+                continue
+            yield path, rel
+
+
+# --- scanning -----------------------------------------------------------------
+
+
+def _scan_modules(modules: list[Module]) -> list[CallSite]:
+    resolver = Resolver(ModuleIndex(modules))
+    sites: list[CallSite] = []
+    calls: list[tuple[str, str]] = []
+    for module in modules:
+        finder = _CallFinder()
+        finder.visit(module.tree)
+        calls.extend((f"{module.rel}::{caller}", callee) for caller, callee in finder.calls)
+        sites.extend(_build_sites(module, finder.found, resolver))
+    _attach_callers(sites, calls)
+    sites.sort(key=lambda site: (site.file, site.line))
+    return sites
+
+
+def _build_sites(module: Module, found: list[_FoundCall], resolver: Resolver) -> list[CallSite]:
+    sites: list[CallSite] = []
+    seen: dict[str, int] = {}
+    for item in found:
+        base_id = f"{module.rel}::{item.qualname}"
+        seen[base_id] = seen.get(base_id, 0) + 1
+        site_id = base_id if seen[base_id] == 1 else f"{base_id}#{seen[base_id]}"
+
+        ctx = Ctx(module, item.func, item.node.lineno)
+        model = resolver.resolve_model(item.node, ctx)
+        messages = resolver.resolve_messages(item.node, item.api, ctx)
+        sites.append(
+            CallSite(
+                id=site_id,
+                file=module.rel,
+                line=item.node.lineno,
+                end_line=item.node.end_lineno,
+                function=item.qualname,
+                api=item.api,
+                model=model,
+                is_async=item.is_async,
+                messages=messages,
+                output_format=_output_format(item.node, item.api),
+                temperature=_temperature(item.node),
+                max_tokens=_max_tokens(item.node),
+                notes=_notes(model, messages),
+            )
+        )
+    return sites
+
+
+def _attach_callers(sites: list[CallSite], calls: list[tuple[str, str]]) -> None:
+    """Name-based: record which functions call the function containing each call site."""
+    for site in sites:
+        if site.function == "<module>":
+            continue
+        short = site.function.split(".")[-1]
+        own = f"{site.file}::{site.function}"
+        site.callers = sorted(
+            {caller for caller, callee in calls if callee == short and caller != own}
+        )
+        if len(site.callers) > 1 and site.messages is None:
+            site.notes.append(
+                f"shared helper called from {len(site.callers)} places; each caller may be "
+                "a separate feature with its own prompt"
+            )
+
+
+@dataclass
+class _FoundCall:
+    node: ast.Call
+    api: str
+    qualname: str
+    func: FunctionNode | None
+    is_async: bool
+
+
+class _CallFinder(ast.NodeVisitor):
+    """Walks a module, recording LLM calls and every function call's caller/callee."""
+
+    def __init__(self) -> None:
+        self.scope: list[str] = []
+        self.functions: list[FunctionNode] = []
+        self.found: list[_FoundCall] = []
+        self.calls: list[tuple[str, str]] = []
+        self._awaited: set[int] = set()
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        self.scope.append(node.name)
+        self.generic_visit(node)
+        self.scope.pop()
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self._visit_function(node)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        self._visit_function(node)
+
+    def _visit_function(self, node: FunctionNode) -> None:
+        self.scope.append(node.name)
+        self.functions.append(node)
+        self.generic_visit(node)
+        self.functions.pop()
+        self.scope.pop()
+
+    def visit_Await(self, node: ast.Await) -> None:
+        if isinstance(node.value, ast.Call):
+            self._awaited.add(id(node.value))
+        self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call) -> None:
+        qualname = ".".join(self.scope) or "<module>"
+        if isinstance(node.func, ast.Name):
+            self.calls.append((qualname, node.func.id))
+        elif isinstance(node.func, ast.Attribute):
+            self.calls.append((qualname, node.func.attr))
+
+        api = _match_api(node)
+        if api is not None:
+            func = self.functions[-1] if self.functions else None
+            self.found.append(_FoundCall(node, api, qualname, func, id(node) in self._awaited))
+        self.generic_visit(node)
+
+
+def _attribute_chain(node: ast.expr) -> list[str]:
+    """client.chat.completions.create -> ["chat", "completions", "create"]."""
+    parts: list[str] = []
+    while isinstance(node, ast.Attribute):
+        parts.append(node.attr)
+        node = node.value
+    parts.reverse()
+    return parts
+
+
+def _match_api(call: ast.Call) -> str | None:
+    chain = _attribute_chain(call.func)
+    for suffix, api, needs_model in API_PATTERNS:
+        if len(chain) >= len(suffix) and tuple(chain[-len(suffix) :]) == suffix:
+            if needs_model and not any(kw.arg == "model" for kw in call.keywords):
+                return None
+            return api
+    return None
+
+
+# --- reading call arguments ---------------------------------------------------
+
+
+def _output_format(call: ast.Call, api: str) -> str:
+    if api.endswith(".parse"):
+        return "json"
+    fmt = keyword_arg(call, "response_format")
+    if isinstance(fmt, ast.Dict):
+        for key, val in zip(fmt.keys, fmt.values, strict=True):
+            if (
+                isinstance(key, ast.Constant)
+                and key.value == "type"
+                and isinstance(val, ast.Constant)
+                and val.value in JSON_FORMAT_TYPES
+            ):
+                return "json"
+    return "text"
+
+
+def _temperature(call: ast.Call) -> float | None:
+    value = keyword_arg(call, "temperature")
+    if isinstance(value, ast.Constant) and isinstance(value.value, (int, float)):
+        if isinstance(value.value, bool):
+            return None
+        return float(value.value)
+    return None
+
+
+def _max_tokens(call: ast.Call) -> int | None:
+    value = keyword_arg(call, *MAX_TOKEN_KWARGS)
+    if (
+        isinstance(value, ast.Constant)
+        and isinstance(value.value, int)
+        and not isinstance(value.value, bool)
+        and value.value >= 1
+    ):
+        return value.value
+    return None
+
+
+def _notes(model: ModelRef, messages: list[PromptMessage] | None) -> list[str]:
+    notes: list[str] = []
+    if model.source == "kwargs":
+        notes.append(f"model is passed via {model.expression} and could not be resolved statically")
+    elif model.source == "missing":
+        notes.append("no model argument found at this call")
+    elif model.source == "env_default":
+        notes.append(f"model comes from env var {model.env_var}; assuming default {model.value!r}")
+    elif model.source == "env":
+        notes.append(f"model comes from env var {model.env_var} with no default")
+    elif model.source == "dynamic":
+        notes.append(f"model is computed at runtime ({model.expression})")
+
+    if messages is None:
+        notes.append("messages are built at runtime; the prompt template could not be recovered")
+    elif not all(m.resolved for m in messages):
+        notes.append("part of the prompt comes from a runtime value the scanner could not see")
+    return notes