querypilot 0.1.0__py3-none-any.whl

querypilot/__init__.py

@@ -0,0 +1,36 @@
+from querypilot.access import AccessPolicy, MaskingRule
+from querypilot.core.client import QueryPilot
+from querypilot.core.types import (
+    ColumnSchema,
+    DatabaseSchema,
+    GeneratedSQL,
+    PolicyCheck,
+    QueryPilotAnswer,
+    QueryResult,
+    SchemaMatch,
+    TableSchema,
+    ValidationResult,
+)
+from querypilot.generation import AnthropicSQLGenerator, OpenAISQLGenerator
+from querypilot.audit import AuditMetadata, InMemoryAuditSink, JSONLAuditSink, QueryAuditRecord
+
+__all__ = [
+    "ColumnSchema",
+    "DatabaseSchema",
+    "GeneratedSQL",
+    "PolicyCheck",
+    "QueryPilot",
+    "QueryPilotAnswer",
+    "QueryResult",
+    "AnthropicSQLGenerator",
+    "AccessPolicy",
+    "AuditMetadata",
+    "InMemoryAuditSink",
+    "JSONLAuditSink",
+    "MaskingRule",
+    "OpenAISQLGenerator",
+    "QueryAuditRecord",
+    "SchemaMatch",
+    "TableSchema",
+    "ValidationResult",
+]

querypilot/access/__init__.py

@@ -0,0 +1,3 @@
+from querypilot.access.policy import AccessPolicy, MaskingRule
+
+__all__ = ["AccessPolicy", "MaskingRule"]

querypilot/access/policy.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class MaskingRule(BaseModel):
+    mode: Literal["redact", "null", "hash"] = "redact"
+
+
+class AccessPolicy(BaseModel):
+    allowed_columns: dict[str, list[str]] = Field(default_factory=dict)
+    blocked_columns: dict[str, list[str]] = Field(default_factory=dict)
+    row_filters: dict[str, str] = Field(default_factory=dict)
+    masking_rules: dict[str, dict[str, MaskingRule]] = Field(default_factory=dict)
+
+    def summary(self) -> dict:
+        return {
+            "allowed_columns": {
+                table: sorted(columns) for table, columns in self.allowed_columns.items()
+            },
+            "blocked_columns": {
+                table: sorted(columns) for table, columns in self.blocked_columns.items()
+            },
+            "row_filters": dict(self.row_filters),
+            "masked_columns": {
+                f"{table}.{column}": rule.mode
+                for table, rules in self.masking_rules.items()
+                for column, rule in rules.items()
+            },
+        }

querypilot/adapters/__init__.py

@@ -0,0 +1 @@
+"""Agent tool adapters."""

querypilot/adapters/anthropic.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def anthropic_tools() -> list[dict[str, Any]]:
+    return [
+        {
+            "name": "ask_database",
+            "description": "Ask a natural-language question of the database using QueryPilot safety checks.",
+            "input_schema": _object_schema(
+                {"question": {"type": "string", "description": "Question to answer."}},
+                ["question"],
+            ),
+        },
+        {
+            "name": "search_schema",
+            "description": "Search tables and columns relevant to a natural-language query.",
+            "input_schema": _object_schema(
+                {"query": {"type": "string", "description": "Schema search query."}},
+                ["query"],
+            ),
+        },
+        {
+            "name": "validate_sql",
+            "description": "Validate SQL and return the safe rewritten form when allowed.",
+            "input_schema": _object_schema(
+                {"sql": {"type": "string", "description": "SQL to validate."}},
+                ["sql"],
+            ),
+        },
+        {
+            "name": "execute_sql",
+            "description": "Validate and execute read-only SQL safely.",
+            "input_schema": _object_schema(
+                {"sql": {"type": "string", "description": "SQL to execute."}},
+                ["sql"],
+            ),
+        },
+    ]
+
+
+def _object_schema(properties: dict[str, Any], required: list[str]) -> dict[str, Any]:
+    return {
+        "type": "object",
+        "properties": properties,
+        "required": required,
+        "additionalProperties": False,
+    }

querypilot/adapters/openai.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def openai_tools() -> list[dict[str, Any]]:
+    return [
+        _function_tool(
+            "ask_database",
+            "Ask a natural-language question of the database using QueryPilot safety checks.",
+            {"question": {"type": "string", "description": "Question to answer."}},
+            ["question"],
+        ),
+        _function_tool(
+            "search_schema",
+            "Search tables and columns relevant to a natural-language query.",
+            {"query": {"type": "string", "description": "Schema search query."}},
+            ["query"],
+        ),
+        _function_tool(
+            "validate_sql",
+            "Validate SQL and return the safe rewritten form when allowed.",
+            {"sql": {"type": "string", "description": "SQL to validate."}},
+            ["sql"],
+        ),
+        _function_tool(
+            "execute_sql",
+            "Validate and execute read-only SQL safely.",
+            {"sql": {"type": "string", "description": "SQL to execute."}},
+            ["sql"],
+        ),
+    ]
+
+
+def _function_tool(
+    name: str,
+    description: str,
+    properties: dict[str, Any],
+    required: list[str],
+) -> dict[str, Any]:
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": description,
+            "parameters": {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+                "additionalProperties": False,
+            },
+        },
+    }

querypilot/audit/__init__.py

@@ -0,0 +1,10 @@
+from querypilot.audit.sinks import AuditSink, InMemoryAuditSink, JSONLAuditSink
+from querypilot.audit.types import AuditMetadata, QueryAuditRecord
+
+__all__ = [
+    "AuditMetadata",
+    "AuditSink",
+    "InMemoryAuditSink",
+    "JSONLAuditSink",
+    "QueryAuditRecord",
+]

querypilot/audit/sinks.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol
+
+from querypilot.audit.types import QueryAuditRecord
+
+
+class AuditSink(Protocol):
+    def write(self, record: QueryAuditRecord) -> None:
+        ...
+
+    def recent(self, limit: int = 100) -> list[QueryAuditRecord]:
+        ...
+
+
+class InMemoryAuditSink:
+    def __init__(self) -> None:
+        self.records: list[QueryAuditRecord] = []
+
+    def write(self, record: QueryAuditRecord) -> None:
+        self.records.append(record)
+
+    def recent(self, limit: int = 100) -> list[QueryAuditRecord]:
+        return list(reversed(self.records[-limit:]))
+
+
+class JSONLAuditSink:
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self._records: list[QueryAuditRecord] = []
+
+    def write(self, record: QueryAuditRecord) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        with self.path.open("a", encoding="utf-8") as handle:
+            handle.write(record.model_dump_json() + "\n")
+        self._records.append(record)
+
+    def recent(self, limit: int = 100) -> list[QueryAuditRecord]:
+        if self._records:
+            return list(reversed(self._records[-limit:]))
+        if not self.path.exists():
+            return []
+        lines = self.path.read_text(encoding="utf-8").splitlines()[-limit:]
+        records = [QueryAuditRecord.model_validate_json(line) for line in lines]
+        return list(reversed(records))

querypilot/audit/types.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from querypilot.core.types import ValidationResult
+
+
+class AuditMetadata(BaseModel):
+    actor: str | None = None
+    session_id: str | None = None
+    app_name: str | None = None
+    trace_id: str | None = None
+
+
+class QueryAuditRecord(BaseModel):
+    audit_id: str = Field(default_factory=lambda: str(uuid4()))
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    operation: str
+    question: str | None = None
+    sql: str | None = None
+    rewritten_sql: str | None = None
+    validation: ValidationResult | None = None
+    valid: bool | None = None
+    executed: bool = False
+    row_count: int | None = None
+    execution_time_ms: int | None = None
+    error: str | None = None
+    access_policy: dict = Field(default_factory=dict)
+    actor: str | None = None
+    session_id: str | None = None
+    app_name: str | None = None
+    trace_id: str | None = None
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        operation: str,
+        metadata: AuditMetadata | None = None,
+        **kwargs,
+    ) -> "QueryAuditRecord":
+        metadata = metadata or AuditMetadata()
+        return cls(
+            operation=operation,
+            actor=metadata.actor,
+            session_id=metadata.session_id,
+            app_name=metadata.app_name,
+            trace_id=metadata.trace_id,
+            **kwargs,
+        )

querypilot/cli.py

@@ -0,0 +1,401 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+from querypilot.access import AccessPolicy
+from querypilot import QueryPilot
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="querypilot")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    serve_parser = subparsers.add_parser("serve", help="Run the FastAPI QueryPilot server.")
+    _add_runtime_args(serve_parser)
+    serve_parser.add_argument("--host", default=os.getenv("QUERYPILOT_HOST", "127.0.0.1"))
+    serve_parser.add_argument("--port", type=int, default=int(os.getenv("QUERYPILOT_PORT", "8000")))
+
+    mcp_parser = subparsers.add_parser("mcp", help="Run the QueryPilot MCP server.")
+    _add_runtime_args(mcp_parser)
+    mcp_parser.add_argument(
+        "--transport",
+        default=os.getenv("QUERYPILOT_MCP_TRANSPORT", "stdio"),
+        choices=["stdio", "streamable-http", "sse"],
+    )
+
+    eval_parser = subparsers.add_parser("eval", help="Eval-driven SQL reliability harness.")
+    eval_subparsers = eval_parser.add_subparsers(dest="eval_command", required=True)
+
+    eval_run_parser = eval_subparsers.add_parser(
+        "run",
+        help="Run a benchmark suite end-to-end and emit a SuiteReport.",
+    )
+    _add_eval_run_args(eval_run_parser)
+
+    eval_replay_parser = eval_subparsers.add_parser(
+        "replay",
+        help="Materialize a regression suite from an audit log.",
+    )
+    _add_eval_replay_args(eval_replay_parser)
+
+    eval_check_parser = eval_subparsers.add_parser(
+        "check",
+        help="Compare a SuiteReport JSON against thresholds and a baseline.",
+    )
+    _add_eval_check_args(eval_check_parser)
+
+    eval_init_parser = eval_subparsers.add_parser(
+        "init",
+        help="Scaffold suites/ and .eval/ in the current directory.",
+    )
+    eval_init_parser.add_argument(
+        "--target",
+        default=".",
+        help="Directory to scaffold into (default: current working directory).",
+    )
+    eval_init_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite existing files.",
+    )
+
+    args = parser.parse_args(argv)
+
+    if args.command == "serve":
+        _serve(args)
+        return 0
+    if args.command == "mcp":
+        _mcp(args)
+        return 0
+    if args.command == "eval":
+        if args.eval_command == "run":
+            return _eval_run(args)
+        if args.eval_command == "replay":
+            return _eval_replay(args)
+        if args.eval_command == "check":
+            return _eval_check(args)
+        if args.eval_command == "init":
+            return _eval_init(args)
+        raise SystemExit(f"Unknown eval subcommand: {args.eval_command}")
+    return 0
+
+
+def _add_runtime_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--database-url", default=os.getenv("QUERYPILOT_DATABASE_URL"))
+    parser.add_argument("--dialect", default=os.getenv("QUERYPILOT_DIALECT", "sqlite"))
+    parser.add_argument("--max-rows", type=int, default=int(os.getenv("QUERYPILOT_MAX_ROWS", "100")))
+    parser.add_argument(
+        "--timeout-seconds",
+        type=int,
+        default=int(os.getenv("QUERYPILOT_TIMEOUT_SECONDS", "10")),
+    )
+    parser.add_argument(
+        "--access-policy-json",
+        default=os.getenv("QUERYPILOT_ACCESS_POLICY_JSON"),
+        help="JSON object for AccessPolicy configuration.",
+    )
+
+
+def _add_eval_replay_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--audit-jsonl",
+        required=True,
+        help="Path to a JSONL audit log written by JSONLAuditSink.",
+    )
+    parser.add_argument(
+        "--fixture-db",
+        required=True,
+        help="Database URL to attach to every replayed case (e.g. sqlite:///fixtures/demo.db).",
+    )
+    parser.add_argument(
+        "--fixture-dialect",
+        default=None,
+        help="Override the dialect inferred from --fixture-db (sqlite/postgres/mysql/snowflake/bigquery/redshift).",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Where to write the replayed suite (.yaml/.yml/.json).",
+    )
+    parser.add_argument(
+        "--name",
+        default="audit_replay",
+        help="Suite name to embed in the output file.",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=1000,
+        help="Maximum number of cases to materialize.",
+    )
+    parser.add_argument(
+        "--include-masked",
+        action="store_true",
+        help="Include cases whose audit record applied a non-empty access policy.",
+    )
+    parser.add_argument(
+        "--include-failures",
+        action="store_true",
+        help="Include cases that failed validation, execution, or had errors.",
+    )
+    parser.add_argument(
+        "--include-empty",
+        action="store_true",
+        help="Include cases that returned zero rows.",
+    )
+    parser.add_argument(
+        "--tag",
+        action="append",
+        default=[],
+        help="Extra tag to attach to every replayed case (repeatable).",
+    )
+
+
+def _add_eval_check_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--report",
+        required=True,
+        help="Path to a SuiteReport JSON written by `querypilot eval run --report ...`.",
+    )
+    parser.add_argument(
+        "--baseline",
+        default=None,
+        help="Optional baseline SuiteReport JSON for regression comparison.",
+    )
+    parser.add_argument(
+        "--threshold",
+        type=float,
+        default=None,
+        help="Minimum overall pass_rate (0..1).",
+    )
+    parser.add_argument(
+        "--max-p95-ms",
+        type=int,
+        default=None,
+        help="Maximum p95 case latency in milliseconds.",
+    )
+    parser.add_argument(
+        "--require-safety",
+        type=float,
+        default=None,
+        help="Minimum safety_pass_rate (0..1).",
+    )
+    parser.add_argument(
+        "--require-correctness",
+        type=float,
+        default=None,
+        help="Minimum correctness_rate (0..1).",
+    )
+    parser.add_argument(
+        "--outcome-json",
+        default=None,
+        help="Optional path to write the structured CheckOutcome JSON.",
+    )
+
+
+def _add_eval_run_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--suite",
+        required=True,
+        help="Path to a suite YAML/JSON file or a directory of suite files.",
+    )
+    parser.add_argument(
+        "--database-url",
+        default=os.getenv("QUERYPILOT_DATABASE_URL"),
+        help="Default database URL when a case does not set fixture_db.",
+    )
+    parser.add_argument("--dialect", default=os.getenv("QUERYPILOT_DIALECT", "sqlite"))
+    parser.add_argument(
+        "--generator",
+        default="demo",
+        choices=("demo", "openai", "anthropic"),
+        help="SQL generator to evaluate. Use 'demo' for the offline deterministic generator.",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        help="Override the default model for openai/anthropic generators.",
+    )
+    parser.add_argument(
+        "--report",
+        default=None,
+        help="Path to write the JSON SuiteReport. Terminal output is always printed.",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Parallel workers (>1 uses ThreadPoolExecutor; LLM rate limits apply).",
+    )
+    parser.add_argument(
+        "--max-rows",
+        type=int,
+        default=int(os.getenv("QUERYPILOT_MAX_ROWS", "100")),
+    )
+    parser.add_argument(
+        "--timeout-seconds",
+        type=int,
+        default=int(os.getenv("QUERYPILOT_TIMEOUT_SECONDS", "10")),
+    )
+    parser.add_argument("--no-color", action="store_true", help="Disable ANSI colors in terminal output.")
+
+
+def _serve(args: argparse.Namespace) -> None:
+    if not args.database_url:
+        raise SystemExit("--database-url or QUERYPILOT_DATABASE_URL is required.")
+    try:
+        import uvicorn
+    except ImportError as exc:
+        raise SystemExit("Install querypilot[server] to use `querypilot serve`.") from exc
+
+    from querypilot.server.app import create_app
+
+    app = create_app(
+        database_url=args.database_url,
+        dialect=args.dialect,
+        max_rows=args.max_rows,
+        timeout_seconds=args.timeout_seconds,
+        access_policy=_access_policy_from_args(args),
+    )
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+def _mcp(args: argparse.Namespace) -> None:
+    if not args.database_url:
+        raise SystemExit("--database-url or QUERYPILOT_DATABASE_URL is required.")
+
+    from querypilot.mcp.server import create_mcp_server
+
+    qp = QueryPilot.connect(
+        database_url=args.database_url,
+        dialect=args.dialect,
+        max_rows=args.max_rows,
+        timeout_seconds=args.timeout_seconds,
+        access_policy=_access_policy_from_args(args),
+    )
+    server = create_mcp_server(qp)
+    server.run(transport=args.transport)
+
+
+def _eval_run(args: argparse.Namespace) -> int:
+    from querypilot.evals.factory import (
+        build_cost_tracker_factory,
+        build_generator,
+        build_qp_factory,
+        load_suite_or_dir,
+    )
+    from querypilot.evals.report import render_terminal, write_json
+    from querypilot.evals.suite_runner import run_suite
+
+    try:
+        suite = load_suite_or_dir(args.suite)
+    except Exception as exc:
+        raise SystemExit(f"Failed to load suite at {args.suite}: {exc}") from exc
+
+    if not args.database_url and suite.fixture_db is None:
+        raise SystemExit(
+            "--database-url is required when the suite does not declare fixture_db."
+        )
+
+    generator = build_generator(args.generator, model=args.model)
+    tracker_factory = build_cost_tracker_factory(args.generator)
+    qp_factory = build_qp_factory(
+        database_url=args.database_url or suite.fixture_db or "",
+        dialect=args.dialect,
+        generator=generator,
+        max_rows=args.max_rows,
+        timeout_seconds=args.timeout_seconds,
+    )
+
+    report = run_suite(
+        suite,
+        qp_factory=qp_factory,
+        cost_tracker_factory=tracker_factory,
+        max_workers=args.workers,
+        generator_name=args.generator,
+        model_name=args.model,
+        database_url=args.database_url,
+    )
+
+    if args.report:
+        write_json(report, args.report)
+
+    print(render_terminal(report, color=not args.no_color))
+    return 0
+
+
+def _eval_check(args: argparse.Namespace) -> int:
+    from querypilot.evals.check import (
+        check_report,
+        format_outcome,
+        load_report,
+        write_outcome,
+    )
+
+    report = load_report(args.report)
+    baseline = load_report(args.baseline) if args.baseline else None
+
+    outcome = check_report(
+        report,
+        baseline=baseline,
+        threshold=args.threshold,
+        max_p95_ms=args.max_p95_ms,
+        require_safety_pass_rate=args.require_safety,
+        require_correctness_rate=args.require_correctness,
+    )
+
+    if args.outcome_json:
+        write_outcome(outcome, args.outcome_json)
+
+    print(format_outcome(outcome))
+    return 0 if outcome.ok else 1
+
+
+def _eval_init(args: argparse.Namespace) -> int:
+    from querypilot.evals.init import scaffold
+
+    target = Path(args.target).resolve()
+    written = scaffold(target, force=args.force)
+    print(f"Scaffolded {len(written)} files in {target}")
+    for path in written:
+        print(f"  {path.relative_to(target)}")
+    return 0
+
+
+def _eval_replay(args: argparse.Namespace) -> int:
+    from querypilot.evals.loader import write_suite
+    from querypilot.evals.replay import replay_from_jsonl
+
+    suite = replay_from_jsonl(
+        args.audit_jsonl,
+        fixture_db=args.fixture_db,
+        fixture_dialect=args.fixture_dialect,
+        suite_name=args.name,
+        only_successful=not args.include_failures,
+        skip_masked=not args.include_masked,
+        skip_empty_results=not args.include_empty,
+        limit=args.limit,
+        extra_tags=args.tag,
+    )
+
+    target = write_suite(suite, args.output)
+    print(f"Wrote {len(suite.cases)} cases to {target}")
+    return 0
+
+
+def _access_policy_from_args(args: argparse.Namespace) -> AccessPolicy | None:
+    if not args.access_policy_json:
+        return None
+    try:
+        payload = json.loads(args.access_policy_json)
+    except json.JSONDecodeError as exc:
+        raise SystemExit(f"Invalid --access-policy-json: {exc}") from exc
+    return AccessPolicy.model_validate(payload)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

querypilot/connectors/__init__.py

@@ -0,0 +1,5 @@
+from querypilot.connectors.base import BaseConnector
+from querypilot.connectors.postgres import PostgresConnector
+from querypilot.connectors.sqlite import SQLiteConnector
+
+__all__ = ["BaseConnector", "PostgresConnector", "SQLiteConnector"]