pretensor 0.1.0__py3-none-any.whl

pretensor/benchmark/l3/agent.py

@@ -0,0 +1,316 @@
+"""LLM tool-use agent loop for the L3 pretensor runner.
+
+The runner asks the LLM to translate one NL question into PostgreSQL,
+giving the model the MCP tool set discovered from ``pretensor serve``.
+The loop here is provider-agnostic: it owns the message history, calls
+the LLM (via :class:`AgentLlmClient`), routes tool invocations through
+a caller-supplied callback, and stops as soon as the model returns
+final text (the SQL).
+
+A hard cap on iterations keeps a confused model from running away
+forever. Per-call telemetry (latency, tokens, tool trace) is returned
+so the runner can populate the JSON envelope without a second pass.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Literal, Protocol
+
+from pretensor.errors import PretensorError
+
+__all__ = [
+    "AgentLlmClient",
+    "AgentLoopError",
+    "AgentLoopResult",
+    "AgentMessage",
+    "AgentStep",
+    "AgentTool",
+    "AgentToolCall",
+    "AgentToolResult",
+    "DEFAULT_MAX_ITERATIONS",
+    "ToolCallTraceEntry",
+    "ToolInvocationOutcome",
+    "ToolInvoker",
+    "run_agent_loop",
+]
+
+
+DEFAULT_MAX_ITERATIONS = 16
+"""Cap on the LLM ↔ tool round-trips per question.
+
+Pagila's hardest gold question solves in ≤4 round-trips today; 16
+leaves headroom for messier schemas without letting a confused model
+spin forever. Hitting the cap is recorded in ``notes[]`` and the
+question is marked failed — never silently accepted.
+"""
+
+
+_TOOL_RESULT_TRUNCATE_BYTES = 8000
+"""Soft cap on the tool result handed back to the LLM, in UTF-8 bytes.
+
+The MCP server can return large payloads (e.g. ``cypher`` over a
+big graph). Trimming protects the next request's prompt budget; the
+tool trace stores the *full* ``response_size`` (also in bytes) so
+an auditor can see when truncation kicked in. Both values share the
+same unit so comparing them is unambiguous.
+"""
+
+_TRUNCATION_SUFFIX = "…[truncated]"
+
+
+@dataclass(frozen=True, slots=True)
+class AgentTool:
+    """An MCP tool exposed to the agent.
+
+    Mirrors :class:`mcp.types.Tool` but lives in the L3 namespace so
+    callers don't have to import the MCP package to read the trace.
+    """
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class AgentToolCall:
+    """One tool invocation requested by the LLM.
+
+    ``id`` is the provider's identifier (Anthropic ``tool_use_id``,
+    OpenAI ``tool_call_id``); we forward it back unchanged so the
+    next request can pair tool results with their calls.
+    """
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class AgentToolResult:
+    """Result of executing one :class:`AgentToolCall`."""
+
+    id: str
+    content: str
+    is_error: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class AgentMessage:
+    """One turn of the LLM conversation, provider-agnostic.
+
+    Role is either ``"assistant"`` (LLM produced this turn) or
+    ``"user"`` (we produced this turn — initial question, or a batch
+    of tool results). A single assistant turn may contain text and
+    one or more tool calls; a single user turn either holds the
+    initial prompt text or a batch of tool results.
+    """
+
+    role: Literal["user", "assistant"]
+    text: str | None = None
+    tool_calls: tuple[AgentToolCall, ...] = ()
+    tool_results: tuple[AgentToolResult, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AgentStep:
+    """One round-trip with the LLM.
+
+    Either ``text`` is set (the model finished and returned a final
+    answer — for L3, the SQL) or ``tool_calls`` is non-empty (the
+    model wants to call tools before producing text). Token counts
+    are optional; some providers omit them.
+    """
+
+    text: str | None
+    tool_calls: tuple[AgentToolCall, ...]
+    prompt_tokens: int | None
+    completion_tokens: int | None
+    latency_ms: int
+
+
+class AgentLlmClient(Protocol):
+    """Single-step tool-use surface for the L3 pretensor runner.
+
+    Implementations translate :class:`AgentMessage` history to the
+    provider's wire format (Anthropic content blocks, OpenAI
+    ``tool_calls``, etc.) and back. The agent loop does not see the
+    provider-specific shape.
+    """
+
+    def agent_complete(
+        self,
+        *,
+        system: str,
+        messages: list[AgentMessage],
+        tools: list[AgentTool],
+        model: str,
+        temperature: float,
+    ) -> AgentStep: ...
+
+
+class AgentLoopError(PretensorError, RuntimeError):
+    """Raised when the loop cannot make further progress.
+
+    Specifically: the iteration cap was hit without the LLM returning
+    final text, OR the LLM produced an empty step (no text and no
+    tool calls — providers shouldn't but it's not impossible).
+    """
+
+
+@dataclass(frozen=True, slots=True)
+class ToolCallTraceEntry:
+    """One row of the per-question tool-call trace.
+
+    Mirrors AC-mandated shape: ``{tool, args, response_size}``.
+    ``is_error`` flags MCP-level failures (subprocess error, unknown
+    tool, handler exception); the runner uses it to surface failed
+    runs without losing the rest of the trace.
+    """
+
+    tool: str
+    args: dict[str, Any]
+    response_size: int
+    is_error: bool
+
+
+@dataclass(slots=True)
+class AgentLoopResult:
+    """Final outcome of the loop for one question.
+
+    ``text`` is the LLM's last textual reply (the SQL, with no fence
+    cleanup yet — the runner does that). ``trace`` is the ordered list
+    of tool calls; ``iterations`` is how many LLM round-trips ran. Token
+    and latency totals are sums across all steps.
+    """
+
+    text: str
+    trace: list[ToolCallTraceEntry] = field(default_factory=list)
+    iterations: int = 0
+    total_prompt_tokens: int | None = None
+    total_completion_tokens: int | None = None
+    total_llm_latency_ms: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class ToolInvocationOutcome:
+    """Outcome of one tool invocation, ready to attach back to the LLM.
+
+    Plain ``content`` + ``is_error`` so the runner doesn't have to know
+    about provider-specific tool_use ids — the loop pairs the outcome
+    with the originating call when it builds the next user turn.
+    """
+
+    content: str
+    is_error: bool = False
+
+
+ToolInvoker = Callable[[str, dict[str, Any]], ToolInvocationOutcome]
+"""Sync callable the loop uses to run one tool.
+
+The runner provides this — typically wrapping an
+:class:`McpClient.call_tool` plus serialisation. Returning a
+:class:`ToolInvocationOutcome` (rather than a raw dict) lets the runner
+mark tool-level errors that the LLM should still see.
+"""
+
+
+def run_agent_loop(
+    *,
+    client: AgentLlmClient,
+    system: str,
+    user: str,
+    tools: list[AgentTool],
+    invoke_tool: ToolInvoker,
+    model: str,
+    temperature: float,
+    max_iterations: int = DEFAULT_MAX_ITERATIONS,
+) -> AgentLoopResult:
+    """Run one question to completion against ``client`` and ``invoke_tool``.
+
+    Each iteration sends the full conversation history to the LLM. If
+    the model returns text, the loop exits with that text. If the model
+    returns tool calls, each call is executed (via ``invoke_tool``),
+    the results are appended as a user turn, and the loop continues.
+    """
+    messages: list[AgentMessage] = [AgentMessage(role="user", text=user)]
+    trace: list[ToolCallTraceEntry] = []
+    total_prompt_tokens: int | None = None
+    total_completion_tokens: int | None = None
+    total_latency_ms = 0
+
+    for iteration in range(1, max_iterations + 1):
+        step = client.agent_complete(
+            system=system,
+            messages=messages,
+            tools=tools,
+            model=model,
+            temperature=temperature,
+        )
+        total_latency_ms += step.latency_ms
+        if step.prompt_tokens is not None:
+            total_prompt_tokens = (total_prompt_tokens or 0) + step.prompt_tokens
+        if step.completion_tokens is not None:
+            total_completion_tokens = (
+                total_completion_tokens or 0
+            ) + step.completion_tokens
+
+        if step.text is not None and not step.tool_calls:
+            return AgentLoopResult(
+                text=step.text,
+                trace=trace,
+                iterations=iteration,
+                total_prompt_tokens=total_prompt_tokens,
+                total_completion_tokens=total_completion_tokens,
+                total_llm_latency_ms=total_latency_ms,
+            )
+
+        if not step.tool_calls:
+            raise AgentLoopError(
+                f"LLM returned an empty step at iteration {iteration} "
+                "(no text and no tool calls)."
+            )
+
+        # Record the assistant turn (text + tool calls) so the next request
+        # carries the full history. ``step.text`` may be non-empty when the
+        # model "thinks out loud" before calling a tool — keep it.
+        messages.append(
+            AgentMessage(
+                role="assistant",
+                text=step.text,
+                tool_calls=step.tool_calls,
+            )
+        )
+
+        results: list[AgentToolResult] = []
+        for call in step.tool_calls:
+            outcome = invoke_tool(call.name, call.arguments)
+            content_bytes = outcome.content.encode("utf-8")
+            trace.append(
+                ToolCallTraceEntry(
+                    tool=call.name,
+                    args=dict(call.arguments),
+                    response_size=len(content_bytes),
+                    is_error=outcome.is_error,
+                )
+            )
+            content = outcome.content
+            if len(content_bytes) > _TOOL_RESULT_TRUNCATE_BYTES:
+                # Slice on bytes to keep the truncation cap unambiguous;
+                # decode with ``errors="ignore"`` so a multibyte
+                # codepoint straddling the boundary doesn't raise.
+                content = (
+                    content_bytes[:_TOOL_RESULT_TRUNCATE_BYTES].decode(
+                        "utf-8", errors="ignore"
+                    )
+                    + _TRUNCATION_SUFFIX
+                )
+            results.append(
+                AgentToolResult(id=call.id, content=content, is_error=outcome.is_error)
+            )
+
+        messages.append(AgentMessage(role="user", tool_results=tuple(results)))
+
+    raise AgentLoopError(
+        f"Agent did not produce final text within {max_iterations} iterations."
+    )

pretensor/benchmark/l3/db.py

@@ -0,0 +1,188 @@
+"""Read-only SQL execution helpers for the L3 benchmark runners.
+
+The runners need to execute both the gold SQL and the agent-emitted SQL
+against a real PostgreSQL instance, then compare the resulting rows. This
+module is the only place that opens DB connections in the L3 layer, so all
+the safety guardrails live here:
+
+* **Database URL is resolved from a per-dataset env var.** No CLI flag —
+  ``--dataset`` already pins which DB to talk to, and growing the CLI
+  surface fragments the contract.
+* **Every query runs in a transaction that is rolled back at the end.**
+  Defense-in-depth against a misbehaving agent that emits ``DROP TABLE``;
+  the read-only intent of the run is preserved even if the syntactic
+  guard below misses a corner case.
+* **Statement-level timeout** (30s by default) prevents pathological
+  agent SQL (cross joins, accidental cartesian products) from hanging
+  the runner.
+* **Non-SELECT guard** rejects any agent SQL whose first non-comment,
+  non-whitespace token is not ``SELECT``, ``WITH``, or ``(``. Together
+  with the rolled-back transaction this is two independent layers of
+  protection — a slip in either is contained by the other.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Any
+
+import sqlalchemy
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import SQLAlchemyError
+
+from pretensor.benchmark.runner import Dataset
+from pretensor.errors import PretensorError
+
+__all__ = [
+    "DEFAULT_STATEMENT_TIMEOUT_MS",
+    "QueryExecutionError",
+    "RefusedNonSelectError",
+    "execute_query",
+    "is_select_only",
+    "resolve_database_url",
+]
+
+
+DEFAULT_STATEMENT_TIMEOUT_MS = 30_000
+
+
+_DATASET_TO_ENV_VAR: dict[Dataset, str] = {
+    Dataset.PAGILA: "PAGILA_DATABASE_URL",
+    Dataset.TPCH: "TPCH_DATABASE_URL",
+    Dataset.ADVENTUREWORKS: "ADVENTUREWORKS_DATABASE_URL",
+}
+"""L3 baseline supports the three OSS-bundled datasets that ship with DDL.
+
+The synthetic fixtures (``analytics_dwh``, ``saas_multitenant``,
+``adversarial``) have no DDL dump and no live DB they correspond to,
+so they are deliberately absent here — calling ``resolve_database_url``
+on them raises a clear error.
+"""
+
+
+class QueryExecutionError(PretensorError, RuntimeError):
+    """Raised when a SQL statement fails to execute (parse error, missing
+    object, runtime fault, statement timeout, etc.)."""
+
+
+class RefusedNonSelectError(QueryExecutionError):
+    """Raised when an agent SQL statement is rejected by the SELECT-only guard.
+
+    A subclass of ``QueryExecutionError`` so callers that want to record
+    "could not execute" can use a single ``except`` and still inspect the
+    type to distinguish "refused before execution" from "DB error".
+    """
+
+
+def resolve_database_url(dataset: Dataset, env: dict[str, str] | None = None) -> str:
+    """Read the per-dataset DSN env var; raise if unset or dataset unsupported.
+
+    ``env`` is for tests (pass an explicit dict to bypass ``os.environ``);
+    production callers leave it ``None`` and the function reads the
+    process environment.
+    """
+    if env is None:
+        env = dict(os.environ)
+    try:
+        var_name = _DATASET_TO_ENV_VAR[dataset]
+    except KeyError as exc:
+        raise LookupError(
+            f"L3 baseline runner does not support dataset {dataset.value!r}: "
+            f"only {[d.value for d in _DATASET_TO_ENV_VAR]} have DDL bundles."
+        ) from exc
+    value = env.get(var_name)
+    if not value:
+        raise LookupError(
+            f"Set {var_name} to the PostgreSQL DSN of a {dataset.value} "
+            f"database (e.g. {var_name}=postgresql://user@host:5432/{dataset.value})."
+        )
+    return value
+
+
+# Strip line comments (``--`` to end of line) and block comments (``/* ... */``);
+# whitespace; then peek at the leading token. CTEs (``WITH ...``) and
+# parenthesised SELECTs (``(SELECT ...) UNION ...``) are allowed alongside
+# bare ``SELECT``.
+_LINE_COMMENT_RE = re.compile(r"--[^\n]*")
+_BLOCK_COMMENT_RE = re.compile(r"/\*.*?\*/", re.DOTALL)
+_ALLOWED_LEAD_TOKENS = ("SELECT", "WITH", "TABLE", "VALUES", "(")
+
+
+def is_select_only(sql: str) -> bool:
+    """Return ``True`` iff ``sql`` begins with a read-only top-level keyword.
+
+    ``TABLE`` and ``VALUES`` are PostgreSQL read forms; ``WITH`` covers
+    CTEs whose final statement is a ``SELECT`` (we don't deep-parse the
+    CTE — combined with the rolled-back transaction in
+    :func:`execute_query`, that's enough). A bare ``(`` allows
+    ``(SELECT ...) UNION ...`` style queries.
+
+    A statement that *contains* DDL/DML keywords later (in a string
+    literal, in a column comment, etc.) is not blocked here — that's a
+    job for the database's own parser. The only goal of this guard is
+    "first thing the LLM emits must be a read".
+    """
+    stripped = _BLOCK_COMMENT_RE.sub(" ", sql)
+    stripped = _LINE_COMMENT_RE.sub(" ", stripped).strip()
+    if not stripped:
+        return False
+    upper = stripped.upper()
+    for tok in _ALLOWED_LEAD_TOKENS:
+        if tok == "(" and upper.startswith("("):
+            return True
+        if upper.startswith(tok) and (
+            len(upper) == len(tok) or not upper[len(tok)].isalnum()
+        ):
+            return True
+    return False
+
+
+def execute_query(
+    dsn: str,
+    sql: str,
+    *,
+    enforce_select_only: bool = False,
+    statement_timeout_ms: int = DEFAULT_STATEMENT_TIMEOUT_MS,
+) -> tuple[list[tuple[Any, ...]], list[str]]:
+    """Run ``sql`` against ``dsn`` inside a rolled-back transaction.
+
+    Returns ``(rows, column_names)``. ``enforce_select_only=True`` runs
+    the syntactic guard first and raises :class:`RefusedNonSelectError`
+    when it fails — used for agent-emitted SQL. Gold SQL skips the guard
+    (it's authored by us and may use any read form, but we still wrap
+    in a rolled-back transaction).
+
+    Any DB error (parse failure, missing object, statement timeout, etc.)
+    is wrapped in :class:`QueryExecutionError` so callers can record it
+    without having to know the SQLAlchemy / driver-specific exception
+    hierarchy.
+    """
+    if enforce_select_only and not is_select_only(sql):
+        raise RefusedNonSelectError(
+            "agent SQL refused: first non-comment token is not SELECT/WITH/(/VALUES."
+        )
+    # A new Engine per call is intentional: L3 only runs ~10–25 questions
+    # per dataset (so 2N ≤ ~50 engines), each query owns its own pool /
+    # transaction lifecycle, and the dispose() in the finally block keeps
+    # connection state from leaking across questions if anything goes wrong.
+    engine: Engine = sqlalchemy.create_engine(dsn)
+    try:
+        with engine.connect() as conn:
+            trans = conn.begin()
+            try:
+                conn.execute(
+                    sqlalchemy.text(
+                        f"SET LOCAL statement_timeout = {int(statement_timeout_ms)}"
+                    )
+                )
+                cursor = conn.execute(sqlalchemy.text(sql))
+                rows = [tuple(row) for row in cursor.fetchall()]
+                column_names = list(cursor.keys())
+                return rows, column_names
+            except SQLAlchemyError as exc:
+                raise QueryExecutionError(str(exc)) from exc
+            finally:
+                trans.rollback()
+    finally:
+        engine.dispose()

pretensor/benchmark/l3/gold.py

@@ -0,0 +1,85 @@
+"""L3 gold-question loader.
+
+The L3 gold corpus lives at ``scripts/data/<dataset>_nl2sql_bench.json`` —
+the same file L2's ``query_recall`` reads. The L3 baseline runner only
+needs three fields per question (``id``, ``question``, ``expected_sql``),
+so this loader strips the rest and returns a sorted list for byte-stable
+output ordering.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+from pretensor.benchmark.fixtures import Fixture
+
+__all__ = ["L3GoldEntry", "load_l3_gold"]
+
+
+@dataclass(frozen=True, slots=True)
+class L3GoldEntry:
+    """One NL-to-SQL question to grade.
+
+    ``id`` is the stable identifier used in the JSON envelope's
+    ``per_item`` list. ``question`` is the natural-language prompt the
+    runner sends to the LLM. ``expected_sql`` is the human-authored gold
+    query whose result rows define the correct answer.
+    """
+
+    id: str
+    question: str
+    expected_sql: str
+
+
+def load_l3_gold(fixture: Fixture) -> tuple[Path, bytes, list[L3GoldEntry]]:
+    """Read and parse the dataset's NL-to-SQL gold file.
+
+    Returns ``(questions_path, raw_bytes, entries)``. ``raw_bytes`` is
+    the file contents exactly as read from disk — exposed so callers can
+    fingerprint the fixture (e.g. ``hashlib.sha256``) without re-reading
+    the file from a second I/O. ``questions_path`` makes
+    ``fixture.questions_path``'s resolved value part of the function's
+    contract so callers don't need ``Optional`` narrowing.
+
+    Entries are sorted by ``id`` so the runner's per-item output stays
+    byte-stable across re-runs even if the source JSON's array order
+    shifts. Raises :class:`FileNotFoundError` if the dataset has no
+    question set on disk; raises :class:`ValueError` for malformed
+    records.
+    """
+    if fixture.questions_path is None:
+        raise FileNotFoundError(
+            f"Dataset {fixture.name.value!r} has no NL-to-SQL gold "
+            "questions file checked in (looked for "
+            f"scripts/data/{fixture.name.value}_nl2sql_bench.json)."
+        )
+    questions_path = fixture.questions_path
+    raw_bytes = questions_path.read_bytes()
+    raw = json.loads(raw_bytes.decode("utf-8"))
+    if not isinstance(raw, list):
+        raise ValueError(
+            f"{questions_path}: expected a JSON array, got {type(raw).__name__}."
+        )
+    entries: list[L3GoldEntry] = []
+    for index, record in enumerate(raw):
+        if not isinstance(record, dict):
+            raise ValueError(
+                f"{questions_path}[{index}]: expected an object, got "
+                f"{type(record).__name__}."
+            )
+        try:
+            entries.append(
+                L3GoldEntry(
+                    id=str(record["id"]),
+                    question=str(record["question"]),
+                    expected_sql=str(record["expected_sql"]),
+                )
+            )
+        except KeyError as exc:
+            raise ValueError(
+                f"{questions_path}[{index}]: missing field {exc.args[0]!r}."
+            ) from exc
+    entries.sort(key=lambda e: e.id)
+    return questions_path, raw_bytes, entries