querypilot 0.1.0__py3-none-any.whl

querypilot/evals/factory.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable
+
+from querypilot import QueryPilot
+from querypilot.evals.cost import (
+    AnthropicCostTracker,
+    CostTracker,
+    NullCostTracker,
+    OpenAICostTracker,
+)
+from querypilot.evals.loader import load_suite, load_suite_dir
+from querypilot.evals.pipeline import QueryPilotFactory
+from querypilot.evals.suite import BenchmarkCase, BenchmarkSuite
+from querypilot.generation.sql_generator import DemoSQLGenerator, SQLGenerator
+
+
+GENERATOR_NAMES = ("demo", "openai", "anthropic")
+
+DEFAULT_OPENAI_MODEL = "gpt-4o-mini"
+DEFAULT_ANTHROPIC_MODEL = "claude-sonnet-4-6"
+
+
+def build_generator(name: str, *, model: str | None = None) -> SQLGenerator:
+    normalized = name.lower()
+    if normalized == "demo":
+        return DemoSQLGenerator()
+    if normalized == "openai":
+        from querypilot.generation.llm import OpenAISQLGenerator
+
+        return OpenAISQLGenerator(model=model or DEFAULT_OPENAI_MODEL)
+    if normalized == "anthropic":
+        from querypilot.generation.llm import AnthropicSQLGenerator
+
+        return AnthropicSQLGenerator(model=model or DEFAULT_ANTHROPIC_MODEL)
+    raise ValueError(
+        f"Unknown generator {name!r}. Supported: {', '.join(GENERATOR_NAMES)}."
+    )
+
+
+def build_cost_tracker_factory(name: str) -> Callable[[], CostTracker]:
+    normalized = name.lower()
+    if normalized == "demo":
+        return NullCostTracker
+    if normalized == "openai":
+        return OpenAICostTracker
+    if normalized == "anthropic":
+        return AnthropicCostTracker
+    raise ValueError(
+        f"Unknown generator {name!r}. Supported: {', '.join(GENERATOR_NAMES)}."
+    )
+
+
+def build_qp_factory(
+    *,
+    database_url: str,
+    dialect: str = "sqlite",
+    generator: SQLGenerator,
+    max_rows: int = 100,
+    timeout_seconds: int = 10,
+    max_generation_attempts: int = 2,
+) -> QueryPilotFactory:
+    def _make(case: BenchmarkCase) -> QueryPilot:
+        url = case.fixture_db or database_url
+        case_dialect = case.fixture_dialect or dialect
+        return QueryPilot.connect(
+            database_url=url,
+            dialect=case_dialect,
+            max_rows=max_rows,
+            timeout_seconds=timeout_seconds,
+            max_generation_attempts=max_generation_attempts,
+            generator=generator,
+        )
+
+    return _make
+
+
+def load_suite_or_dir(path: str | Path) -> BenchmarkSuite:
+    target = Path(path)
+    if target.is_dir():
+        return load_suite_dir(target)
+    return load_suite(target)

querypilot/evals/init.py

@@ -0,0 +1,104 @@
+"""Scaffold a starter eval layout (suites/ and .eval/) for a project."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+_SMOKE_YAML = """name: smoke
+
+# Update fixture_db to point at your own SQLite or Postgres test database.
+# Relative sqlite:/// paths are resolved against this file's directory.
+fixture_db: sqlite:///REPLACE_ME_WITH_YOUR_DB.db
+fixture_dialect: sqlite
+
+thresholds:
+  pass_rate: 0.95
+  safety_pass_rate: 1.0
+  correctness_rate: 0.9
+  max_p95_latency_ms: 5000
+
+comparison:
+  ignore_row_order: true
+  ignore_column_order: true
+  float_tolerance: 0.001
+  normalize_datetimes: true
+  case_insensitive_strings: false
+
+# Replace these with cases that exercise your own schema. gold_sql is
+# the SQL we expect to be functionally equivalent to what the generator
+# produces; results are compared row-by-row after normalization.
+cases:
+  - id: example_count
+    question: "Count of customers"
+    gold_sql: "SELECT COUNT(*) AS count FROM customers"
+    expected_tables: [customers]
+    must_not_contain: [DELETE, UPDATE, DROP, INSERT]
+    tags: [smoke, aggregation]
+"""
+
+_SAFETY_YAML = """name: safety
+
+# Safety cases run raw SQL strings through the validator and assert it
+# blocks them. Update fixture_db to match your smoke suite.
+fixture_db: sqlite:///REPLACE_ME_WITH_YOUR_DB.db
+fixture_dialect: sqlite
+
+thresholds:
+  safety_pass_rate: 1.0
+
+cases:
+  - id: blocks_drop_table
+    sql: "DROP TABLE customers"
+    should_pass: false
+    expected_failure_kind: validation
+    expected_error_contains: ["Only SELECT queries are allowed"]
+    tags: [safety, ddl]
+
+  - id: blocks_update
+    sql: "UPDATE customers SET revenue = 0"
+    should_pass: false
+    expected_failure_kind: validation
+    tags: [safety, mutation]
+
+  - id: blocks_multi_statement
+    sql: "SELECT * FROM customers; DROP TABLE customers"
+    should_pass: false
+    expected_failure_kind: validation
+    tags: [safety, multi_statement]
+"""
+
+_BASELINE_README = """# .eval/
+
+This directory holds the committed baseline `SuiteReport` JSON used by
+`querypilot eval check` to detect regressions.
+
+To refresh on `main` after a deliberate change:
+
+    querypilot eval run --suite suites/smoke.yaml --report .eval/baseline.json
+    git commit -am "Refresh eval baseline"
+"""
+
+
+def scaffold(target: Path, *, force: bool = False) -> list[Path]:
+    written: list[Path] = []
+    suites_dir = target / "suites"
+    eval_dir = target / ".eval"
+    suites_dir.mkdir(parents=True, exist_ok=True)
+    eval_dir.mkdir(parents=True, exist_ok=True)
+
+    smoke_path = suites_dir / "smoke.yaml"
+    safety_path = suites_dir / "safety.yaml"
+    readme_path = eval_dir / "README.md"
+
+    for path, contents in (
+        (smoke_path, _SMOKE_YAML),
+        (safety_path, _SAFETY_YAML),
+        (readme_path, _BASELINE_README),
+    ):
+        if path.exists() and not force:
+            continue
+        path.write_text(contents, encoding="utf-8")
+        written.append(path)
+
+    return written

querypilot/evals/loader.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from querypilot.evals.suite import BenchmarkCase, BenchmarkSuite
+
+
+class SuiteLoadError(ValueError):
+    """Raised when a benchmark suite cannot be loaded or parsed."""
+
+
+_SQLITE_PREFIX = "sqlite:///"
+
+
+def write_suite(suite: BenchmarkSuite, path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    payload = suite.model_dump(mode="json", exclude_none=True)
+    text = _dump(payload, target.suffix.lower())
+    target.write_text(text, encoding="utf-8")
+    return target
+
+
+def _dump(payload: dict[str, Any], suffix: str) -> str:
+    if suffix == ".json":
+        return json.dumps(payload, indent=2, sort_keys=False) + "\n"
+    if suffix in {".yaml", ".yml"}:
+        try:
+            import yaml
+        except ImportError as exc:
+            raise SuiteLoadError(
+                "PyYAML is required to write YAML suites. Install with `pip install querypilot[eval]`."
+            ) from exc
+        return yaml.safe_dump(payload, sort_keys=False)
+    raise SuiteLoadError(f"Unsupported suite file extension for write: {suffix}")
+
+
+def load_suite(path: str | Path) -> BenchmarkSuite:
+    suite_path = Path(path).resolve()
+    if not suite_path.is_file():
+        raise SuiteLoadError(f"Suite file not found: {suite_path}")
+
+    payload = _parse_file(suite_path)
+    if not isinstance(payload, dict):
+        raise SuiteLoadError(f"Suite file must contain a mapping at the top level: {suite_path}")
+
+    suite = _build_suite(payload, suite_path.parent)
+    return suite
+
+
+def load_suite_dir(path: str | Path) -> BenchmarkSuite:
+    suite_dir = Path(path).resolve()
+    if not suite_dir.is_dir():
+        raise SuiteLoadError(f"Suite directory not found: {suite_dir}")
+
+    suite_files = sorted(
+        p for p in suite_dir.iterdir() if p.suffix.lower() in {".yaml", ".yml", ".json"}
+    )
+    if not suite_files:
+        raise SuiteLoadError(f"No .yaml/.yml/.json suite files found in: {suite_dir}")
+
+    sub_suites = [load_suite(suite_file) for suite_file in suite_files]
+    head, *rest = sub_suites
+    head_path = suite_files[0]
+
+    for sub, sub_path in zip(rest, suite_files[1:]):
+        _require_dir_match(head, head_path, sub, sub_path, "fixture_db")
+        _require_dir_match(head, head_path, sub, sub_path, "fixture_dialect")
+        _require_dir_match(head, head_path, sub, sub_path, "thresholds")
+        _require_dir_match(head, head_path, sub, sub_path, "comparison")
+
+    merged_cases: list[BenchmarkCase] = []
+    seen_ids: set[str] = set()
+    for sub in sub_suites:
+        for case in sub.cases:
+            if case.id in seen_ids:
+                raise SuiteLoadError(
+                    f"Duplicate case id across suite directory {suite_dir}: {case.id!r}"
+                )
+            seen_ids.add(case.id)
+            merged_cases.append(case)
+
+    return BenchmarkSuite(
+        name=head.name,
+        fixture_db=head.fixture_db,
+        fixture_dialect=head.fixture_dialect,
+        thresholds=head.thresholds,
+        comparison=head.comparison,
+        cases=merged_cases,
+    )
+
+
+def _require_dir_match(
+    head: BenchmarkSuite,
+    head_path: Path,
+    sub: BenchmarkSuite,
+    sub_path: Path,
+    field: str,
+) -> None:
+    head_value = getattr(head, field)
+    sub_value = getattr(sub, field)
+    if head_value != sub_value:
+        raise SuiteLoadError(
+            f"{field!r} differs across suite directory: "
+            f"{head_path.name} has {head_value!r} but {sub_path.name} has {sub_value!r}. "
+            "All suite files in a directory must share fixture_db, fixture_dialect, "
+            "thresholds, and comparison settings."
+        )
+
+
+def _parse_file(path: Path) -> Any:
+    text = path.read_text(encoding="utf-8")
+    suffix = path.suffix.lower()
+
+    if suffix == ".json":
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise SuiteLoadError(f"Invalid JSON in {path}: {exc}") from exc
+
+    if suffix in {".yaml", ".yml"}:
+        try:
+            import yaml
+        except ImportError as exc:
+            raise SuiteLoadError(
+                "PyYAML is required to load YAML suites. Install with `pip install querypilot[eval]`."
+            ) from exc
+        try:
+            return yaml.safe_load(text)
+        except yaml.YAMLError as exc:
+            raise SuiteLoadError(f"Invalid YAML in {path}: {exc}") from exc
+
+    raise SuiteLoadError(f"Unsupported suite file extension: {path.suffix} ({path})")
+
+
+def _build_suite(payload: dict[str, Any], suite_dir: Path) -> BenchmarkSuite:
+    payload = dict(payload)
+
+    suite_fixture_db = payload.get("fixture_db")
+    if suite_fixture_db is not None:
+        payload["fixture_db"] = _resolve_fixture_db(str(suite_fixture_db), suite_dir)
+
+    raw_cases = payload.get("cases") or []
+    if not isinstance(raw_cases, list):
+        raise SuiteLoadError("Suite 'cases' must be a list.")
+
+    resolved_cases: list[dict[str, Any]] = []
+    for raw in raw_cases:
+        if not isinstance(raw, dict):
+            raise SuiteLoadError(f"Each case must be a mapping; got {type(raw).__name__}.")
+        case = dict(raw)
+        case_fixture_db = case.get("fixture_db")
+        if case_fixture_db is not None:
+            case["fixture_db"] = _resolve_fixture_db(str(case_fixture_db), suite_dir)
+        resolved_cases.append(case)
+    payload["cases"] = resolved_cases
+
+    try:
+        return BenchmarkSuite.model_validate(payload)
+    except ValueError as exc:
+        raise SuiteLoadError(f"Invalid suite definition: {exc}") from exc
+
+
+def _resolve_fixture_db(fixture_db: str, suite_dir: Path) -> str:
+    if not fixture_db.startswith(_SQLITE_PREFIX):
+        return fixture_db
+
+    raw_path = fixture_db[len(_SQLITE_PREFIX):]
+    if not raw_path or raw_path.startswith(":memory:"):
+        return fixture_db
+
+    candidate = Path(raw_path)
+    if candidate.is_absolute():
+        return fixture_db
+
+    resolved = (suite_dir / candidate).resolve()
+    return f"{_SQLITE_PREFIX}{resolved}"