llm/structured_schema.py
Rama main · réplica de lectura del repositorio original.
# SPDX-License-Identifier: BUSL-1.1
# Copyright (c) 2024-2026 Ali Sadeghi Aghili
"""Pydantic-backed schema for Phase 2 task 3's constrained SQL-generation output.
With deterministic decoding in place (Phase 2 task 2), the next latency/
quality lever is asking the model for one schema-shaped JSON object
instead of free text plus :func:`security.sql_guard.clean_sql`'s
fence-stripping / preamble-removal regex surgery and the ``OUT_OF_SCOPE``
string sentinel. Whether that trade is actually *worth* it is an empirical
question the golden-set evaluation answers (see ``eval/cli.py``'s
``--structured`` flag) — this module only defines the schema and the
small amount of glue needed to evaluate both paths; it does not decide
which one wins, and the feature stays off by default (see
``config.Settings.llm_structured_output``) until that evaluation says
otherwise.
One model, not two definitions of the same contract
------------------------------------------------------
:data:`SqlGeneration` is the single source of truth: the same Pydantic
model both generates the JSON Schema sent to the endpoint
(:data:`SQL_GENERATION_SCHEMA`, via :meth:`~pydantic.BaseModel.model_json_schema`)
and validates what comes back (:func:`sql_from_structured`, via
:meth:`~pydantic.BaseModel.model_validate`). A hand-written schema dict
alongside separate hand-written parsing code can drift silently — the
model returns something the schema permitted and the parser then
disagrees about it — which is exactly the failure mode a single
generate-and-validate model closes by construction.
Two details in :func:`_to_strict_json_schema` matter enough to call out
explicitly:
1. OpenAI-compatible *strict* ``json_schema`` mode requires
``additionalProperties: false`` on every object **and** every property
listed in ``required`` — Pydantic's ``model_config = {"extra":
"forbid"}`` already emits the former, but Pydantic only lists a field
in ``required`` when it has no default, so a model whose fields all
have defaults (this one) needs its ``required`` list built explicitly
from every property name, not taken from Pydantic's own output as-is.
2. Pydantic emits ``$ref``/``$defs`` for nested models; grammar backends
behind an OpenAI-compatible server vary in whether they resolve those
(llama.cpp's GBNF conversion in particular has historically not).
:data:`SqlGeneration` is kept deliberately flat (no nested model
fields) to sidestep the question entirely rather than inline
``$defs`` after the fact.
"""
from __future__ import annotations
from typing import Any
from pydantic import BaseModel, ConfigDict, Field, ValidationError
_OUT_OF_SCOPE_SENTINEL = "OUT_OF_SCOPE"
class SqlGeneration(BaseModel):
"""The single source of truth for the constrained SQL-generation output.
Deliberately flat — no nested model fields — so the JSON Schema this
generates has no ``$defs``/``$ref`` for a grammar backend to (maybe)
not resolve. See the module docstring.
Examples
--------
>>> SqlGeneration(sql="SELECT 1", out_of_scope=False).sql
'SELECT 1'
>>> SqlGeneration().out_of_scope
False
"""
model_config = ConfigDict(extra="forbid")
sql: str = Field(
default="", description="The generated T-SQL query. Empty when out_of_scope is true."
)
out_of_scope: bool = Field(
default=False, description="True if the question cannot be answered from the known schema."
)
confidence: float = Field(
default=1.0, description="The model's own confidence in this SQL, 0.0-1.0."
)
assumptions: list[str] = Field(
default_factory=list,
description="Any assumptions the model made to resolve ambiguity in the question.",
)
def _to_strict_json_schema(model: type[BaseModel]) -> dict[str, Any]:
"""Post-process ``model.model_json_schema()`` for OpenAI strict mode.
Strict ``json_schema`` mode requires ``additionalProperties: false``
(Pydantic's ``extra="forbid"`` already emits this) **and** every
property name listed under ``required``, even ones with a default —
Pydantic only lists genuinely-mandatory fields there by default, so
this rebuilds ``required`` from every property name explicitly rather
than trusting Pydantic's own list.
Examples
--------
>>> schema = _to_strict_json_schema(SqlGeneration)
>>> schema["additionalProperties"]
False
>>> sorted(schema["required"]) == sorted(schema["properties"])
True
>>> "$defs" in schema
False
"""
schema = model.model_json_schema()
schema.pop("$defs", None) # belt-and-braces: SqlGeneration is flat by design
# The class docstring ends up in "description" and every field's own
# name in "title" -- neither is needed by the endpoint (the field
# "description"s already carry the guidance that matters) and both
# just add tokens to a payload sent on every call.
schema.pop("description", None)
schema.pop("title", None)
for prop in schema.get("properties", {}).values():
prop.pop("title", None)
schema["additionalProperties"] = False
schema["required"] = list(schema.get("properties", {}))
return schema
#: The schema every structured SQL-generation call asks for — generated
#: from, and validated against, :class:`SqlGeneration`.
SQL_GENERATION_SCHEMA: dict[str, Any] = _to_strict_json_schema(SqlGeneration)
class SchemaViolationError(ValueError):
"""A response was valid JSON matching the wire schema but failed the model.
This is a real, expected outcome — not a bug — for a structured-output
call: an endpoint's constrained decoding guarantees the *shape* named
in the JSON Schema (right keys, right JSON types), not the additional
constraints :class:`SqlGeneration` enforces via Pydantic (e.g.
``extra="forbid"`` rejecting a stray key the schema's own
``additionalProperties: false`` should have already blocked, but a
lenient backend didn't). A caller building the ``llm`` status block
(``docs/api-contract-v2.md`` §6) should map this to
``finish_reason: "schema_violation"`` — the contract already names
that value for exactly this state — rather than letting the
underlying :class:`pydantic.ValidationError` escape untranslated.
"""
def sql_from_structured(obj: dict[str, Any]) -> str:
"""Validate *obj* against :class:`SqlGeneration` and extract its SQL.
Parameters
----------
obj:
A dict shaped like :data:`SQL_GENERATION_SCHEMA`, e.g. from
:meth:`~llm.base.LLMBackend.generate_structured`.
Returns
-------
str
The validated ``sql`` field.
Raises
------
ValueError("OUT_OF_SCOPE")
If ``out_of_scope`` is true — the same sentinel
:class:`~llm.base.LLMBackend` raises on the text path, so callers
(:class:`~llm.sql_agent.SQLAgent`, ``eval/runner.py``) don't need a
second out-of-scope contract to handle.
SchemaViolationError
If *obj* does not validate against :class:`SqlGeneration` (a
missing/extra/wrong-typed field) — a real, distinguishable state,
not a :class:`KeyError` or a bare :class:`pydantic.ValidationError`
escaping untranslated. See that class's docstring.
Examples
--------
>>> sql_from_structured({"sql": "SELECT 1", "out_of_scope": False})
'SELECT 1'
>>> sql_from_structured({"sql": "", "out_of_scope": True})
Traceback (most recent call last):
...
ValueError: OUT_OF_SCOPE
>>> sql_from_structured( # doctest: +ELLIPSIS
... {"sql": "SELECT 1", "out_of_scope": False, "extra_key": 1}
... )
Traceback (most recent call last):
...
llm.structured_schema.SchemaViolationError: ...
"""
try:
parsed = SqlGeneration.model_validate(obj)
except ValidationError as exc:
raise SchemaViolationError(
f"Structured response did not match SqlGeneration: {exc}"
) from exc
if parsed.out_of_scope:
raise ValueError(_OUT_OF_SCOPE_SENTINEL)
return parsed.sql