querypilot 0.1.0__py3-none-any.whl

querypilot/execution/formatter.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+
+def explain_result(question: str, sql: str, rows: list[dict]) -> str:
+    if not rows:
+        return f"The validated query returned no rows for: {question}"
+    columns = ", ".join(rows[0].keys())
+    return (
+        f"The validated query answers the question using columns {columns}. "
+        f"It returned {len(rows)} row(s) after QueryPilot safety checks."
+    )

querypilot/execution/runner.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import time
+
+from querypilot.connectors.base import BaseConnector
+from querypilot.core.types import QueryResult
+
+
+def execute(connector: BaseConnector, sql: str) -> QueryResult:
+    start = time.perf_counter()
+    rows, row_count = connector.execute_readonly(sql)
+    elapsed_ms = int((time.perf_counter() - start) * 1000)
+    return QueryResult(sql=sql, rows=rows, row_count=row_count, execution_time_ms=elapsed_ms)

querypilot/generation/__init__.py

@@ -0,0 +1,9 @@
+from querypilot.generation.llm import AnthropicSQLGenerator, OpenAISQLGenerator
+from querypilot.generation.sql_generator import DemoSQLGenerator, SQLGenerator
+
+__all__ = [
+    "AnthropicSQLGenerator",
+    "DemoSQLGenerator",
+    "OpenAISQLGenerator",
+    "SQLGenerator",
+]

querypilot/generation/llm.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+from querypilot.core.types import DatabaseSchema, GeneratedSQL, ValidationResult
+from querypilot.generation.prompt_builder import build_sql_generation_prompt
+
+
+SQL_JSON_SCHEMA: dict[str, Any] = {
+    "name": "querypilot_sql",
+    "schema": {
+        "type": "object",
+        "properties": {
+            "sql": {"type": "string"},
+            "explanation": {"type": "string"},
+        },
+        "required": ["sql", "explanation"],
+        "additionalProperties": False,
+    },
+    "strict": True,
+}
+
+
+class OpenAISQLGenerator:
+    def __init__(
+        self,
+        client: Any | None = None,
+        model: str = "gpt-5.1",
+        max_output_tokens: int = 800,
+    ) -> None:
+        self.client = client or _default_openai_client()
+        self.model = model
+        self.max_output_tokens = max_output_tokens
+
+    def generate(self, question: str, schema: DatabaseSchema, max_rows: int) -> GeneratedSQL:
+        prompt = build_sql_generation_prompt(question, schema, max_rows)
+        return self._create(question, prompt.instructions, prompt.user_prompt)
+
+    def repair(
+        self,
+        question: str,
+        schema: DatabaseSchema,
+        max_rows: int,
+        previous_sql: str,
+        validation: ValidationResult,
+    ) -> GeneratedSQL:
+        prompt = build_sql_generation_prompt(
+            question,
+            schema,
+            max_rows,
+            validation_errors=validation.errors,
+            previous_sql=previous_sql,
+        )
+        return self._create(question, prompt.instructions, prompt.user_prompt)
+
+    def _create(self, question: str, instructions: str, user_prompt: str) -> GeneratedSQL:
+        response = self.client.responses.create(
+            model=self.model,
+            instructions=instructions,
+            input=user_prompt,
+            max_output_tokens=self.max_output_tokens,
+            text={"format": {"type": "json_schema", **SQL_JSON_SCHEMA}},
+        )
+        return _generated_from_text(question, _extract_openai_text(response))
+
+
+class AnthropicSQLGenerator:
+    def __init__(
+        self,
+        client: Any | None = None,
+        model: str = "claude-sonnet-4-20250514",
+        max_tokens: int = 800,
+    ) -> None:
+        self.client = client or _default_anthropic_client()
+        self.model = model
+        self.max_tokens = max_tokens
+
+    def generate(self, question: str, schema: DatabaseSchema, max_rows: int) -> GeneratedSQL:
+        prompt = build_sql_generation_prompt(question, schema, max_rows)
+        return self._create(question, prompt.instructions, prompt.user_prompt)
+
+    def repair(
+        self,
+        question: str,
+        schema: DatabaseSchema,
+        max_rows: int,
+        previous_sql: str,
+        validation: ValidationResult,
+    ) -> GeneratedSQL:
+        prompt = build_sql_generation_prompt(
+            question,
+            schema,
+            max_rows,
+            validation_errors=validation.errors,
+            previous_sql=previous_sql,
+        )
+        return self._create(question, prompt.instructions, prompt.user_prompt)
+
+    def _create(self, question: str, system_prompt: str, user_prompt: str) -> GeneratedSQL:
+        message = self.client.messages.create(
+            model=self.model,
+            max_tokens=self.max_tokens,
+            system=system_prompt,
+            messages=[{"role": "user", "content": user_prompt}],
+        )
+        return _generated_from_text(question, _extract_anthropic_text(message))
+
+
+def _generated_from_text(question: str, text: str) -> GeneratedSQL:
+    try:
+        payload = json.loads(_strip_code_fence(text))
+    except json.JSONDecodeError as exc:
+        return GeneratedSQL(
+            question=question,
+            sql=None,
+            errors=[f"LLM response was not valid JSON: {exc}"],
+        )
+
+    sql = payload.get("sql")
+    if not isinstance(sql, str) or not sql.strip():
+        return GeneratedSQL(question=question, sql=None, errors=["LLM response did not include SQL."])
+    explanation = payload.get("explanation")
+    return GeneratedSQL(
+        question=question,
+        sql=sql.strip(),
+        explanation=explanation if isinstance(explanation, str) else None,
+    )
+
+
+def _strip_code_fence(text: str) -> str:
+    stripped = text.strip()
+    match = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", stripped, flags=re.DOTALL)
+    return match.group(1).strip() if match else stripped
+
+
+def _extract_openai_text(response: Any) -> str:
+    output_text = getattr(response, "output_text", None)
+    if isinstance(output_text, str):
+        return output_text
+    output = getattr(response, "output", None)
+    if isinstance(output, list):
+        for item in output:
+            for content in getattr(item, "content", []) or []:
+                text = getattr(content, "text", None)
+                if isinstance(text, str):
+                    return text
+    return str(response)
+
+
+def _extract_anthropic_text(message: Any) -> str:
+    content = getattr(message, "content", None)
+    if isinstance(content, list):
+        chunks: list[str] = []
+        for block in content:
+            text = getattr(block, "text", None)
+            if isinstance(text, str):
+                chunks.append(text)
+            elif isinstance(block, dict) and isinstance(block.get("text"), str):
+                chunks.append(block["text"])
+        if chunks:
+            return "".join(chunks)
+    return str(message)
+
+
+def _default_openai_client() -> Any:
+    try:
+        from openai import OpenAI
+    except ImportError as exc:
+        raise ImportError("Install querypilot[openai] to use OpenAISQLGenerator.") from exc
+    return OpenAI()
+
+
+def _default_anthropic_client() -> Any:
+    try:
+        import anthropic
+    except ImportError as exc:
+        raise ImportError("Install querypilot[anthropic] to use AnthropicSQLGenerator.") from exc
+    return anthropic.Anthropic()

querypilot/generation/prompt_builder.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from querypilot.core.types import DatabaseSchema
+
+
+class SQLGenerationPrompt(BaseModel):
+    instructions: str
+    user_prompt: str
+
+
+def build_sql_generation_prompt(
+    question: str,
+    schema: DatabaseSchema,
+    max_rows: int,
+    validation_errors: list[str] | None = None,
+    previous_sql: str | None = None,
+) -> SQLGenerationPrompt:
+    instructions = (
+        "You generate safe SQL for QueryPilot. Return only JSON with keys "
+        "`sql` and `explanation`. Generate SELECT-only SQL. Use only tables "
+        "and columns present in the provided schema. Do not use mutations, DDL, "
+        "multiple statements, comments, or vendor-specific unsafe commands. "
+        f"The maximum row limit is {max_rows}; include a LIMIT no larger than that "
+        "unless the query is an aggregate that returns one row."
+    )
+    if validation_errors:
+        instructions += " Repair the SQL using the validation errors."
+
+    parts = [
+        f"Question: {question}",
+        "",
+        "Schema:",
+        build_schema_context(schema),
+    ]
+    if previous_sql:
+        parts.extend(["", f"Previous SQL: {previous_sql}"])
+    if validation_errors:
+        parts.extend(["", "Validation errors:", "\n".join(f"- {error}" for error in validation_errors)])
+    return SQLGenerationPrompt(instructions=instructions, user_prompt="\n".join(parts))
+
+
+def build_schema_context(schema: DatabaseSchema) -> str:
+    lines: list[str] = []
+    for table in schema.tables:
+        columns = ", ".join(f"{column.name} {column.type}" for column in table.columns)
+        lines.append(f"{table.name}({columns})")
+    return "\n".join(lines)

querypilot/generation/sql_generator.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import re
+from typing import Protocol
+
+from querypilot.core.types import DatabaseSchema, GeneratedSQL
+
+
+class SQLGenerator(Protocol):
+    def generate(self, question: str, schema: DatabaseSchema, max_rows: int) -> GeneratedSQL:
+        ...
+
+
+class DemoSQLGenerator:
+    def generate(self, question: str, schema: DatabaseSchema, max_rows: int) -> GeneratedSQL:
+        normalized = question.lower()
+        explicit_limit = _extract_limit(normalized)
+        limit = min(explicit_limit or max_rows, max_rows)
+
+        if "count" in normalized and schema.tables:
+            table = _best_table(normalized, schema) or schema.tables[0]
+            return GeneratedSQL(
+                question=question,
+                sql=f"SELECT COUNT(*) AS count FROM {table.name} LIMIT {limit}",
+                explanation=f"Counts rows in {table.name}.",
+            )
+
+        customer_table = schema.get_table("customers")
+        if customer_table and any(term in normalized for term in ["customer", "customers"]):
+            if "revenue" in normalized and customer_table.get_column("revenue"):
+                return GeneratedSQL(
+                    question=question,
+                    sql=(
+                        "SELECT customer_name, revenue FROM customers "
+                        f"ORDER BY revenue DESC LIMIT {limit}"
+                    ),
+                    explanation="Ranks customers by revenue.",
+                )
+            if "arr" in normalized and customer_table.get_column("arr"):
+                return GeneratedSQL(
+                    question=question,
+                    sql=f"SELECT customer_name, arr FROM customers ORDER BY arr DESC LIMIT {limit}",
+                    explanation="Ranks customers by ARR.",
+                )
+            return GeneratedSQL(
+                question=question,
+                sql=f"SELECT * FROM customers LIMIT {limit}",
+                explanation="Returns customer records.",
+            )
+
+        table = _best_table(normalized, schema)
+        if table and any(term in normalized for term in ["show", "list", "top"]):
+            return GeneratedSQL(
+                question=question,
+                sql=f"SELECT * FROM {table.name} LIMIT {limit}",
+                explanation=f"Returns rows from {table.name}.",
+            )
+
+        return GeneratedSQL(
+            question=question,
+            sql=None,
+            explanation=None,
+            errors=[
+                "Offline generator could not map this question to a safe SQL template. "
+                "Configure a SQLGenerator provider for broader natural-language support."
+            ],
+        )
+
+
+def _extract_limit(question: str) -> int | None:
+    match = re.search(r"\btop\s+(\d+)\b|\blimit\s+(\d+)\b", question)
+    if not match:
+        return None
+    value = match.group(1) or match.group(2)
+    return int(value)
+
+
+def _best_table(question: str, schema: DatabaseSchema):
+    for table in schema.tables:
+        if table.name.lower() in question or table.name.lower().rstrip("s") in question:
+            return table
+    return None

querypilot/mcp/__init__.py

@@ -0,0 +1,3 @@
+from querypilot.mcp.server import create_mcp_server
+
+__all__ = ["create_mcp_server"]

querypilot/mcp/server.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from querypilot import QueryPilot
+
+
+def create_mcp_server(
+    querypilot: QueryPilot,
+    *,
+    fastmcp_cls: type | None = None,
+):
+    fastmcp = fastmcp_cls or _load_fastmcp()
+    server = fastmcp("QueryPilot", stateless_http=True, json_response=True)
+
+    @server.tool()
+    def ask_database(question: str) -> dict[str, Any]:
+        """Ask a database question through QueryPilot's safe SQL flow."""
+        return _safe_call(lambda: querypilot.ask(question).model_dump())
+
+    @server.tool()
+    def search_schema(query: str) -> list[dict[str, Any]] | dict[str, str]:
+        """Search tables and columns relevant to a question."""
+        return _safe_call(lambda: [match.model_dump() for match in querypilot.search_schema(query)])
+
+    @server.tool()
+    def validate_sql(sql: str) -> dict[str, Any]:
+        """Validate SQL and return policy metadata."""
+        return _safe_call(lambda: querypilot.validate_sql(sql).model_dump())
+
+    @server.tool()
+    def execute_sql(sql: str) -> dict[str, Any]:
+        """Validate and execute read-only SQL safely."""
+        return _safe_call(lambda: querypilot.execute_sql(sql).model_dump())
+
+    return server
+
+
+def _safe_call(callback: Callable[[], Any]) -> Any:
+    try:
+        return callback()
+    except ValueError as exc:
+        return {"error": str(exc)}
+
+
+def _load_fastmcp():
+    try:
+        from mcp.server.fastmcp import FastMCP
+    except ImportError as exc:
+        raise ImportError("Install querypilot[mcp] to run the MCP server.") from exc
+    return FastMCP

querypilot/schema/__init__.py

@@ -0,0 +1,3 @@
+from querypilot.schema.search import search_schema
+
+__all__ = ["search_schema"]

querypilot/schema/introspector.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from querypilot.connectors.base import BaseConnector
+from querypilot.core.types import DatabaseSchema
+
+
+def introspect(connector: BaseConnector) -> DatabaseSchema:
+    return connector.get_schema()

querypilot/schema/search.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import re
+
+from querypilot.core.types import DatabaseSchema, SchemaMatch
+
+
+def search_schema(schema: DatabaseSchema, query: str) -> list[SchemaMatch]:
+    tokens = set(re.findall(r"[a-z0-9_]+", query.lower()))
+    matches: list[SchemaMatch] = []
+    for table in schema.tables:
+        table_tokens = _name_tokens(table.name)
+        table_score = len(tokens & table_tokens) * 3
+        column_hits: list[str] = []
+        for column in table.columns:
+            column_score = len(tokens & _name_tokens(column.name))
+            if column_score:
+                column_hits.append(column.name)
+                table_score += column_score
+        if table_score:
+            matches.append(
+                SchemaMatch(
+                    table=table.name,
+                    column=", ".join(column_hits) if column_hits else None,
+                    score=table_score,
+                    reason="Matched table or column name",
+                )
+            )
+    return sorted(matches, key=lambda match: (-match.score, match.table))
+
+
+def _name_tokens(name: str) -> set[str]:
+    return set(re.findall(r"[a-z0-9]+", name.lower()))

querypilot/server/__init__.py

@@ -0,0 +1,3 @@
+from querypilot.server.app import create_app
+
+__all__ = ["create_app"]

querypilot/server/app.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from querypilot import QueryPilot
+from querypilot.access import AccessPolicy
+from querypilot.audit import AuditMetadata
+from querypilot.core.config import SafetyPolicy
+from querypilot.evals.pipeline import run_case
+from querypilot.evals.suite import BenchmarkCase, ComparisonConfig
+
+
+class QuestionRequest(BaseModel):
+    question: str
+    metadata: AuditMetadata | None = None
+
+
+class SQLRequest(BaseModel):
+    sql: str
+    metadata: AuditMetadata | None = None
+
+
+class SchemaSearchRequest(BaseModel):
+    query: str
+    metadata: AuditMetadata | None = None
+
+
+class EvalRunRequest(BaseModel):
+    cases: list[BenchmarkCase] = Field(default_factory=list)
+    comparison: ComparisonConfig | None = None
+
+
+def create_app(
+    querypilot: QueryPilot | None = None,
+    *,
+    database_url: str | None = None,
+    dialect: str = "sqlite",
+    readonly: bool = True,
+    max_rows: int = 100,
+    timeout_seconds: int = 10,
+    max_generation_attempts: int = 2,
+    allowed_tables: list[str] | None = None,
+    blocked_tables: list[str] | None = None,
+    safety_policy: SafetyPolicy | None = None,
+    access_policy: AccessPolicy | None = None,
+) -> FastAPI:
+    qp = querypilot or _connect_querypilot(
+        database_url=database_url,
+        dialect=dialect,
+        readonly=readonly,
+        max_rows=max_rows,
+        timeout_seconds=timeout_seconds,
+        max_generation_attempts=max_generation_attempts,
+        allowed_tables=allowed_tables,
+        blocked_tables=blocked_tables,
+        safety_policy=safety_policy,
+        access_policy=access_policy,
+    )
+
+    app = FastAPI(
+        title="QueryPilot",
+        description="Safe SQL tool layer for AI agents.",
+        version="0.1.0",
+    )
+    app.state.querypilot = qp
+
+    @app.get("/health")
+    def health() -> dict[str, str]:
+        return {"status": "ok"}
+
+    @app.get("/schema")
+    def get_schema() -> dict[str, Any]:
+        return qp.get_schema().model_dump()
+
+    @app.post("/search-schema")
+    def search_schema(request: SchemaSearchRequest) -> list[dict[str, Any]]:
+        scoped_qp = qp.with_audit_metadata(request.metadata)
+        return [match.model_dump() for match in scoped_qp.search_schema(request.query)]
+
+    @app.post("/ask")
+    def ask(request: QuestionRequest) -> dict[str, Any]:
+        scoped_qp = qp.with_audit_metadata(request.metadata)
+        try:
+            return scoped_qp.ask(request.question).model_dump()
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    @app.post("/generate-sql")
+    def generate_sql(request: QuestionRequest) -> dict[str, Any]:
+        scoped_qp = qp.with_audit_metadata(request.metadata)
+        return scoped_qp.generate_sql(request.question).model_dump()
+
+    @app.post("/validate-sql")
+    def validate_sql(request: SQLRequest) -> dict[str, Any]:
+        scoped_qp = qp.with_audit_metadata(request.metadata)
+        return scoped_qp.validate_sql(request.sql).model_dump()
+
+    @app.post("/execute-sql")
+    def execute_sql(request: SQLRequest) -> dict[str, Any]:
+        scoped_qp = qp.with_audit_metadata(request.metadata)
+        try:
+            return scoped_qp.execute_sql(request.sql).model_dump()
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    @app.post("/evals/run")
+    def run_evals(request: EvalRunRequest) -> dict[str, Any]:
+        comparison = request.comparison or ComparisonConfig()
+        results = [
+            run_case(case, lambda _: qp, comparison=comparison).model_dump(mode="json")
+            for case in request.cases
+        ]
+        passed = sum(1 for r in results if r["passed"])
+        return {
+            "total": len(results),
+            "passed": passed,
+            "failed": len(results) - passed,
+            "case_results": results,
+        }
+
+    @app.get("/audit/recent")
+    def recent_audit(limit: int = 100) -> list[dict[str, Any]]:
+        return [record.model_dump() for record in qp.get_audit_records(limit)]
+
+    return app
+
+
+def _connect_querypilot(
+    *,
+    database_url: str | None,
+    dialect: str,
+    readonly: bool,
+    max_rows: int,
+    timeout_seconds: int,
+    max_generation_attempts: int,
+    allowed_tables: list[str] | None,
+    blocked_tables: list[str] | None,
+    safety_policy: SafetyPolicy | None,
+    access_policy: AccessPolicy | None,
+) -> QueryPilot:
+    if database_url is None:
+        raise ValueError("database_url is required when querypilot is not provided.")
+    return QueryPilot.connect(
+        database_url=database_url,
+        dialect=dialect,
+        readonly=readonly,
+        max_rows=max_rows,
+        timeout_seconds=timeout_seconds,
+        max_generation_attempts=max_generation_attempts,
+        allowed_tables=allowed_tables,
+        blocked_tables=blocked_tables,
+        safety_policy=safety_policy,
+        access_policy=access_policy,
+    )

querypilot/validation/__init__.py

@@ -0,0 +1,3 @@
+from querypilot.validation.validator import SQLValidator
+
+__all__ = ["SQLValidator"]

querypilot/validation/policies.py

@@ -0,0 +1,3 @@
+from querypilot.core.safety_defaults import BLOCKED_POSTGRES_FUNCTIONS
+
+__all__ = ["BLOCKED_POSTGRES_FUNCTIONS"]