querypilot 0.1.0__py3-none-any.whl

querypilot/evals/check.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from querypilot.evals.suite_runner import SuiteReport
+
+
+class CaseRegression(BaseModel):
+    id: str
+    was_passing: bool
+    now_passing: bool
+    previous_failure_category: str | None = None
+    current_failure_category: str | None = None
+
+
+class CheckSummary(BaseModel):
+    pass_rate: float
+    safety_pass_rate: float
+    correctness_rate: float
+    p95_latency_ms: int
+    estimated_cost_usd: float
+    total_cases: int
+    passed: int
+    failed: int
+
+
+class CheckOutcome(BaseModel):
+    ok: bool
+    reasons: list[str] = Field(default_factory=list)
+    current: CheckSummary
+    baseline: CheckSummary | None = None
+    regressed_cases: list[CaseRegression] = Field(default_factory=list)
+    latency_p95_delta_ms: int | None = None
+
+
+def check_report(
+    report: SuiteReport,
+    *,
+    baseline: SuiteReport | None = None,
+    threshold: float | None = None,
+    max_p95_ms: int | None = None,
+    require_safety_pass_rate: float | None = None,
+    require_correctness_rate: float | None = None,
+) -> CheckOutcome:
+    reasons: list[str] = []
+
+    if threshold is not None and report.pass_rate < threshold:
+        reasons.append(
+            f"pass_rate {report.pass_rate:.3f} below threshold {threshold:.3f}"
+        )
+
+    if (
+        require_safety_pass_rate is not None
+        and report.safety_pass_rate < require_safety_pass_rate
+    ):
+        reasons.append(
+            f"safety_pass_rate {report.safety_pass_rate:.3f} below required "
+            f"{require_safety_pass_rate:.3f}"
+        )
+
+    if (
+        require_correctness_rate is not None
+        and report.correctness_rate < require_correctness_rate
+    ):
+        reasons.append(
+            f"correctness_rate {report.correctness_rate:.3f} below required "
+            f"{require_correctness_rate:.3f}"
+        )
+
+    if max_p95_ms is not None and report.p95_latency_ms > max_p95_ms:
+        reasons.append(
+            f"p95_latency_ms {report.p95_latency_ms} exceeds {max_p95_ms}"
+        )
+
+    for violation in report.threshold_violations:
+        if violation not in reasons:
+            reasons.append(violation)
+
+    regressed: list[CaseRegression] = []
+    latency_delta: int | None = None
+    baseline_summary: CheckSummary | None = None
+
+    if baseline is not None:
+        baseline_summary = _summary(baseline)
+        latency_delta = report.p95_latency_ms - baseline.p95_latency_ms
+
+        if report.pass_rate < baseline.pass_rate:
+            reasons.append(
+                f"pass_rate dropped from baseline {baseline.pass_rate:.3f} to "
+                f"{report.pass_rate:.3f}"
+            )
+
+        if report.safety_pass_rate < baseline.safety_pass_rate:
+            reasons.append(
+                f"safety_pass_rate dropped from baseline "
+                f"{baseline.safety_pass_rate:.3f} to {report.safety_pass_rate:.3f}"
+            )
+
+        if report.correctness_rate < baseline.correctness_rate:
+            reasons.append(
+                f"correctness_rate dropped from baseline "
+                f"{baseline.correctness_rate:.3f} to {report.correctness_rate:.3f}"
+            )
+
+        regressed = _diff_cases(baseline=baseline, current=report)
+
+    return CheckOutcome(
+        ok=not reasons and not regressed,
+        reasons=reasons,
+        current=_summary(report),
+        baseline=baseline_summary,
+        regressed_cases=regressed,
+        latency_p95_delta_ms=latency_delta,
+    )
+
+
+def load_report(path: str | Path) -> SuiteReport:
+    target = Path(path)
+    if not target.is_file():
+        raise FileNotFoundError(f"Report file not found: {target}")
+    return SuiteReport.model_validate_json(target.read_text(encoding="utf-8"))
+
+
+def format_outcome(outcome: CheckOutcome) -> str:
+    if outcome.ok:
+        return f"OK\n  pass_rate {outcome.current.pass_rate:.3f}, p95 {outcome.current.p95_latency_ms} ms"
+
+    lines = ["Regression detected.", ""]
+    if outcome.baseline is not None:
+        lines.extend(
+            [
+                "Pass rate:",
+                f"  baseline: {outcome.baseline.pass_rate * 100:.0f}%",
+                f"  current:  {outcome.current.pass_rate * 100:.0f}%",
+                "",
+            ]
+        )
+
+    if outcome.regressed_cases:
+        lines.append("Failed cases (regression vs. baseline):")
+        for case in outcome.regressed_cases:
+            verb = "was passing" if case.was_passing else "was failing"
+            now = "now passing" if case.now_passing else (
+                f"now {case.current_failure_category}" if case.current_failure_category
+                else "now failing"
+            )
+            lines.append(f"  - {case.id} ({verb} -> {now})")
+        lines.append("")
+
+    if outcome.latency_p95_delta_ms is not None and outcome.baseline is not None:
+        delta = outcome.latency_p95_delta_ms
+        lines.extend(
+            [
+                "Latency:",
+                f"  baseline p95: {outcome.baseline.p95_latency_ms} ms",
+                f"  current p95:  {outcome.current.p95_latency_ms} ms  ({delta:+d} ms)",
+                "",
+            ]
+        )
+
+    if outcome.reasons:
+        lines.append("Reasons:")
+        for r in outcome.reasons:
+            lines.append(f"  - {r}")
+
+    return "\n".join(lines).rstrip() + "\n"
+
+
+def write_outcome(outcome: CheckOutcome, path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(
+        json.dumps(outcome.model_dump(mode="json"), indent=2) + "\n",
+        encoding="utf-8",
+    )
+    return target
+
+
+def _summary(report: SuiteReport) -> CheckSummary:
+    return CheckSummary(
+        pass_rate=report.pass_rate,
+        safety_pass_rate=report.safety_pass_rate,
+        correctness_rate=report.correctness_rate,
+        p95_latency_ms=report.p95_latency_ms,
+        estimated_cost_usd=report.estimated_cost_usd,
+        total_cases=report.total_cases,
+        passed=report.passed,
+        failed=report.failed,
+    )
+
+
+def _diff_cases(*, baseline: SuiteReport, current: SuiteReport) -> list[CaseRegression]:
+    baseline_by_id = {c.id: c for c in baseline.case_results}
+    current_by_id = {c.id: c for c in current.case_results}
+
+    out: list[CaseRegression] = []
+    for case_id, current_case in current_by_id.items():
+        baseline_case = baseline_by_id.get(case_id)
+        if baseline_case is None:
+            continue
+        if baseline_case.passed and not current_case.passed:
+            out.append(
+                CaseRegression(
+                    id=case_id,
+                    was_passing=True,
+                    now_passing=False,
+                    previous_failure_category=(
+                        baseline_case.failure_category.value
+                        if baseline_case.failure_category
+                        else None
+                    ),
+                    current_failure_category=(
+                        current_case.failure_category.value
+                        if current_case.failure_category
+                        else None
+                    ),
+                )
+            )
+    return out

querypilot/evals/compare.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import math
+from collections import Counter
+from datetime import date, datetime, time
+from decimal import Decimal
+from typing import Any
+
+from pydantic import BaseModel, Field
+from sqlglot import parse_one
+from sqlglot.errors import ParseError
+
+from querypilot.evals.suite import ComparisonConfig
+
+
+class _NaNSentinel:
+    _instance = None
+
+    def __new__(cls) -> "_NaNSentinel":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __repr__(self) -> str:
+        return "<NaN>"
+
+    def __reduce__(self):
+        return (_NaNSentinel, ())
+
+
+NAN = _NaNSentinel()
+
+
+class ValueMismatch(BaseModel):
+    row_index: int
+    column: str
+    gold: Any
+    candidate: Any
+
+
+class RowsetMatch(BaseModel):
+    matched: bool
+    order_sensitive: bool
+    column_mismatch: list[str] = Field(default_factory=list)
+    missing_rows: list[dict] = Field(default_factory=list)
+    extra_rows: list[dict] = Field(default_factory=list)
+    mismatched_values: list[ValueMismatch] = Field(default_factory=list)
+    normalized_gold_rows: list[dict] = Field(default_factory=list)
+    normalized_candidate_rows: list[dict] = Field(default_factory=list)
+
+
+def has_order_by(sql: str) -> bool:
+    try:
+        tree = parse_one(sql)
+    except ParseError:
+        return False
+    if tree is None:
+        return False
+    return tree.args.get("order") is not None
+
+
+def compare_rows(
+    gold_rows: list[dict],
+    candidate_rows: list[dict],
+    gold_sql: str,
+    config: ComparisonConfig | None = None,
+) -> RowsetMatch:
+    config = config or ComparisonConfig()
+
+    order_sensitive = has_order_by(gold_sql) or not config.ignore_row_order
+
+    column_mismatch = _column_mismatch(gold_rows, candidate_rows, config)
+    if column_mismatch:
+        return RowsetMatch(
+            matched=False,
+            order_sensitive=order_sensitive,
+            column_mismatch=column_mismatch,
+            normalized_gold_rows=gold_rows,
+            normalized_candidate_rows=candidate_rows,
+        )
+
+    norm_gold = [_normalize_row(row, config) for row in gold_rows]
+    norm_candidate = [_normalize_row(row, config) for row in candidate_rows]
+
+    if order_sensitive:
+        return _compare_ordered(norm_gold, norm_candidate, tolerance=config.float_tolerance)
+    return _compare_bag(norm_gold, norm_candidate, tolerance=config.float_tolerance)
+
+
+def _column_mismatch(
+    gold_rows: list[dict], candidate_rows: list[dict], config: ComparisonConfig
+) -> list[str]:
+    if not gold_rows or not candidate_rows:
+        return []
+    gold_cols = list(gold_rows[0].keys())
+    candidate_cols = list(candidate_rows[0].keys())
+
+    if config.ignore_column_order:
+        gold_set = set(gold_cols)
+        candidate_set = set(candidate_cols)
+        if gold_set == candidate_set:
+            return []
+        diffs: list[str] = []
+        for col in sorted(gold_set - candidate_set):
+            diffs.append(f"missing in candidate: {col}")
+        for col in sorted(candidate_set - gold_set):
+            diffs.append(f"unexpected in candidate: {col}")
+        return diffs
+
+    if gold_cols == candidate_cols:
+        return []
+    return [f"column order mismatch: gold={gold_cols} candidate={candidate_cols}"]
+
+
+def _normalize_row(row: dict, config: ComparisonConfig) -> dict:
+    return {key: _normalize_value(value, config) for key, value in row.items()}
+
+
+def _normalize_value(value: Any, config: ComparisonConfig) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, Decimal):
+        value = float(value)
+    if isinstance(value, int) and not isinstance(value, bool):
+        return value
+    if isinstance(value, float):
+        if math.isnan(value):
+            return NAN
+        return value
+    if isinstance(value, datetime):
+        return value.isoformat() if config.normalize_datetimes else value
+    if isinstance(value, date):
+        return value.isoformat() if config.normalize_datetimes else value
+    if isinstance(value, time):
+        return value.isoformat() if config.normalize_datetimes else value
+    if isinstance(value, str):
+        if config.case_insensitive_strings:
+            return value.casefold()
+        return value
+    return value
+
+
+def _compare_ordered(
+    gold: list[dict], candidate: list[dict], *, tolerance: float
+) -> RowsetMatch:
+    mismatched: list[ValueMismatch] = []
+    missing: list[dict] = []
+    extra: list[dict] = []
+
+    common = min(len(gold), len(candidate))
+    for index in range(common):
+        for column, gold_value in gold[index].items():
+            candidate_value = candidate[index].get(column)
+            if not _values_equal(gold_value, candidate_value, tolerance):
+                mismatched.append(
+                    ValueMismatch(
+                        row_index=index,
+                        column=column,
+                        gold=gold_value,
+                        candidate=candidate_value,
+                    )
+                )
+
+    if len(gold) > len(candidate):
+        missing = gold[common:]
+    elif len(candidate) > len(gold):
+        extra = candidate[common:]
+
+    matched = not mismatched and not missing and not extra
+    return RowsetMatch(
+        matched=matched,
+        order_sensitive=True,
+        missing_rows=missing,
+        extra_rows=extra,
+        mismatched_values=mismatched,
+        normalized_gold_rows=gold,
+        normalized_candidate_rows=candidate,
+    )
+
+
+def _compare_bag(
+    gold: list[dict], candidate: list[dict], *, tolerance: float
+) -> RowsetMatch:
+    if tolerance > 0:
+        missing_rows, extra_rows = _greedy_match(gold, candidate, tolerance)
+    else:
+        missing_rows, extra_rows = _counter_diff(gold, candidate)
+
+    matched = not missing_rows and not extra_rows
+    return RowsetMatch(
+        matched=matched,
+        order_sensitive=False,
+        missing_rows=missing_rows,
+        extra_rows=extra_rows,
+        normalized_gold_rows=gold,
+        normalized_candidate_rows=candidate,
+    )
+
+
+def _counter_diff(
+    gold: list[dict], candidate: list[dict]
+) -> tuple[list[dict], list[dict]]:
+    gold_counter = Counter(_row_key(row) for row in gold)
+    candidate_counter = Counter(_row_key(row) for row in candidate)
+
+    missing_keys = gold_counter - candidate_counter
+    extra_keys = candidate_counter - gold_counter
+
+    return _rows_for_keys(gold, missing_keys), _rows_for_keys(candidate, extra_keys)
+
+
+def _greedy_match(
+    gold: list[dict], candidate: list[dict], tolerance: float
+) -> tuple[list[dict], list[dict]]:
+    available = list(range(len(candidate)))
+    missing: list[dict] = []
+    matched_candidate: set[int] = set()
+
+    for gold_row in gold:
+        match_index = None
+        for i in available:
+            if _rows_equal(gold_row, candidate[i], tolerance):
+                match_index = i
+                break
+        if match_index is None:
+            missing.append(gold_row)
+        else:
+            available.remove(match_index)
+            matched_candidate.add(match_index)
+
+    extra = [candidate[i] for i in range(len(candidate)) if i not in matched_candidate]
+    return missing, extra
+
+
+def _rows_equal(gold: dict, candidate: dict, tolerance: float) -> bool:
+    if set(gold.keys()) != set(candidate.keys()):
+        return False
+    for column, gold_value in gold.items():
+        if not _values_equal(gold_value, candidate.get(column), tolerance):
+            return False
+    return True
+
+
+def _row_key(row: dict) -> tuple:
+    return tuple(sorted((key, _hashable(value)) for key, value in row.items()))
+
+
+def _hashable(value: Any) -> Any:
+    if isinstance(value, (list, tuple)):
+        return tuple(_hashable(item) for item in value)
+    if isinstance(value, dict):
+        return tuple(sorted((k, _hashable(v)) for k, v in value.items()))
+    return value
+
+
+def _rows_for_keys(rows: list[dict], wanted: Counter) -> list[dict]:
+    remaining = Counter(wanted)
+    out: list[dict] = []
+    for row in rows:
+        key = _row_key(row)
+        if remaining.get(key, 0) > 0:
+            out.append(row)
+            remaining[key] -= 1
+    return out
+
+
+def _values_equal(a: Any, b: Any, tolerance: float = 0.0) -> bool:
+    if isinstance(a, _NaNSentinel) and isinstance(b, _NaNSentinel):
+        return True
+    if isinstance(a, _NaNSentinel) or isinstance(b, _NaNSentinel):
+        return False
+    if a is None and b is None:
+        return True
+    if a is None or b is None:
+        return False
+    if (
+        isinstance(a, float)
+        and isinstance(b, float)
+        and not isinstance(a, bool)
+        and not isinstance(b, bool)
+        and tolerance > 0
+    ):
+        return abs(a - b) <= tolerance
+    return a == b

querypilot/evals/cost.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from pydantic import BaseModel
+
+
+# (input_per_1k_usd, output_per_1k_usd). Missing models leave estimated_usd unset.
+MODEL_PRICING: dict[str, tuple[float, float]] = {
+    # OpenAI Responses API (illustrative; refresh as pricing changes)
+    "gpt-5.1": (0.01, 0.03),
+    "gpt-4o": (0.005, 0.015),
+    "gpt-4o-mini": (0.00015, 0.0006),
+    # Anthropic Messages API
+    "claude-opus-4-7": (0.015, 0.075),
+    "claude-sonnet-4-6": (0.003, 0.015),
+    "claude-sonnet-4-20250514": (0.003, 0.015),
+    "claude-haiku-4-5-20251001": (0.00025, 0.00125),
+}
+
+
+class TokenUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    model: str | None = None
+    estimated_usd: float | None = None
+
+
+class CostTracker(Protocol):
+    def wrap(self, generator: Any) -> Any:
+        ...
+
+    def restore(self) -> None:
+        ...
+
+    def last_usage(self) -> TokenUsage | None:
+        ...
+
+    def reset(self) -> None:
+        ...
+
+
+class NullCostTracker:
+    def wrap(self, generator: Any) -> Any:
+        return generator
+
+    def restore(self) -> None:
+        pass
+
+    def last_usage(self) -> TokenUsage | None:
+        return None
+
+    def reset(self) -> None:
+        pass
+
+
+class _ResponseCapture:
+    def __init__(self, inner: Any, sink: list) -> None:
+        self._inner = inner
+        self._sink = sink
+
+    @property
+    def __querypilot_inner__(self) -> Any:
+        return self._inner
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+    def create(self, *args: Any, **kwargs: Any) -> Any:
+        response = self._inner.create(*args, **kwargs)
+        self._sink.append(response)
+        return response
+
+
+def _unwrap(namespace: Any) -> Any:
+    while isinstance(namespace, _ResponseCapture):
+        namespace = namespace.__querypilot_inner__
+    return namespace
+
+
+class OpenAICostTracker:
+    def __init__(self) -> None:
+        self._captured: list[Any] = []
+        self._model: str | None = None
+        self._wrapped: tuple[Any, Any] | None = None  # (client, original_responses)
+
+    def wrap(self, generator: Any) -> Any:
+        self._model = getattr(generator, "model", None)
+        client = getattr(generator, "client", None)
+        if client is None or not hasattr(client, "responses"):
+            return generator
+        original = _unwrap(client.responses)
+        if isinstance(client.responses, _ResponseCapture):
+            return generator  # already wrapped
+        client.responses = _ResponseCapture(original, self._captured)
+        self._wrapped = (client, original)
+        return generator
+
+    def restore(self) -> None:
+        if self._wrapped is None:
+            return
+        client, original = self._wrapped
+        client.responses = original
+        self._wrapped = None
+
+    def last_usage(self) -> TokenUsage | None:
+        if not self._captured:
+            return None
+        response = self._captured[-1]
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return None
+        prompt = _coerce_int(getattr(usage, "input_tokens", None))
+        completion = _coerce_int(getattr(usage, "output_tokens", None))
+        if prompt is None and completion is None:
+            return None
+        prompt = prompt or 0
+        completion = completion or 0
+        return TokenUsage(
+            prompt_tokens=prompt,
+            completion_tokens=completion,
+            total_tokens=prompt + completion,
+            model=self._model,
+            estimated_usd=estimate_usd(self._model, prompt, completion),
+        )
+
+    def reset(self) -> None:
+        self._captured.clear()
+
+
+class AnthropicCostTracker:
+    def __init__(self) -> None:
+        self._captured: list[Any] = []
+        self._model: str | None = None
+        self._wrapped: tuple[Any, Any] | None = None  # (client, original_messages)
+
+    def wrap(self, generator: Any) -> Any:
+        self._model = getattr(generator, "model", None)
+        client = getattr(generator, "client", None)
+        if client is None or not hasattr(client, "messages"):
+            return generator
+        original = _unwrap(client.messages)
+        if isinstance(client.messages, _ResponseCapture):
+            return generator  # already wrapped
+        client.messages = _ResponseCapture(original, self._captured)
+        self._wrapped = (client, original)
+        return generator
+
+    def restore(self) -> None:
+        if self._wrapped is None:
+            return
+        client, original = self._wrapped
+        client.messages = original
+        self._wrapped = None
+
+    def last_usage(self) -> TokenUsage | None:
+        if not self._captured:
+            return None
+        message = self._captured[-1]
+        usage = getattr(message, "usage", None)
+        if usage is None:
+            return None
+        prompt = _coerce_int(getattr(usage, "input_tokens", None))
+        completion = _coerce_int(getattr(usage, "output_tokens", None))
+        if prompt is None and completion is None:
+            return None
+        prompt = prompt or 0
+        completion = completion or 0
+        return TokenUsage(
+            prompt_tokens=prompt,
+            completion_tokens=completion,
+            total_tokens=prompt + completion,
+            model=self._model,
+            estimated_usd=estimate_usd(self._model, prompt, completion),
+        )
+
+    def reset(self) -> None:
+        self._captured.clear()
+
+
+def estimate_usd(model: str | None, prompt_tokens: int, completion_tokens: int) -> float | None:
+    if model is None:
+        return None
+    pricing = MODEL_PRICING.get(model)
+    if pricing is None:
+        return None
+    input_per_1k, output_per_1k = pricing
+    return round(
+        (prompt_tokens / 1000.0) * input_per_1k
+        + (completion_tokens / 1000.0) * output_per_1k,
+        6,
+    )
+
+
+def _coerce_int(value: Any) -> int | None:
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None