text2sql-engine 0.1.0__py3-none-any.whl

text2sql/pipeline/linking.py

@@ -0,0 +1,206 @@
+"""Stage 1 — schema linking.
+
+Reduces the full schema to the part relevant to the request. Uses an
+intermediate LLM with forced structured (JSON) output.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..config import Settings
+from ..logging_utils import get_logger
+from ..prompts import PromptTemplates
+from ..providers.base import LLMProvider
+from ..schema import Schema, schema_to_prompt
+from ..types import (
+    LinkedColumn,
+    LinkedJoin,
+    LinkedSchema,
+    LLMResponse,
+    StageMetadata,
+    _Timer,
+)
+from .json_utils import extract_json, with_retry
+
+logger = get_logger("pipeline.linking")
+
+#: Name for the structured-output tool / response_format schema.
+LINKING_SCHEMA_NAME = "linked_schema"
+
+#: JSON schema forced on the Stage 1 output.
+LINKING_JSON_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "additionalProperties": False,
+    "properties": {
+        "tables": {"type": "array", "items": {"type": "string"}},
+        "columns": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "additionalProperties": False,
+                "properties": {
+                    "table": {"type": "string"},
+                    "column": {"type": "string"},
+                    "reason": {"type": "string"},
+                },
+                "required": ["table", "column", "reason"],
+            },
+        },
+        "joins": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "additionalProperties": False,
+                "properties": {
+                    "left_table": {"type": "string"},
+                    "left_column": {"type": "string"},
+                    "right_table": {"type": "string"},
+                    "right_column": {"type": "string"},
+                },
+                "required": [
+                    "left_table",
+                    "left_column",
+                    "right_table",
+                    "right_column",
+                ],
+            },
+        },
+        "entities": {"type": "array", "items": {"type": "string"}},
+        "filters": {"type": "array", "items": {"type": "string"}},
+        "rationale": {"type": "string"},
+    },
+    "required": ["tables", "columns", "joins", "entities", "filters", "rationale"],
+}
+
+
+class LinkingStage:
+    """Runs schema linking (Stage 1)."""
+
+    def __init__(
+        self,
+        provider: LLMProvider,
+        settings: Settings,
+        prompts: PromptTemplates,
+    ) -> None:
+        self._provider = provider
+        self._settings = settings
+        self._prompts = prompts
+
+    def run(self, request: str, schema: Schema) -> tuple[LinkedSchema, StageMetadata]:
+        """Link the request against ``schema``.
+
+        Returns the reduced schema and the stage metadata.
+        """
+        stage_cfg = self._settings.linking
+        template = self._prompts.load(stage_cfg.prompt_template)
+        system, user = template.render(
+            schema=schema_to_prompt(schema),
+            request=request,
+            dialect=self._settings.general.sql_dialect,
+            max_tables=stage_cfg.max_tables,
+        )
+
+        attempts = {"n": 0}
+
+        def attempt() -> tuple[LinkedSchema, LLMResponse]:
+            attempts["n"] += 1
+            response = self._provider.complete_json(
+                system=system,
+                user=user,
+                json_schema=LINKING_JSON_SCHEMA,
+                schema_name=LINKING_SCHEMA_NAME,
+                temperature=stage_cfg.temperature,
+                top_p=stage_cfg.top_p,
+                max_tokens=stage_cfg.max_tokens,
+            )
+            data = extract_json(response.text, strict=self._settings.general.strict_json)
+            linked = self._to_linked_schema(data, schema, stage_cfg.max_tables)
+            return linked, response
+
+        with _Timer() as timer:
+            linked, response = with_retry(
+                attempt,
+                retry=self._settings.retry,
+                description="Schema linking",
+            )
+
+        metadata = StageMetadata(
+            provider=self._provider.name,
+            model=self._provider.model,
+            usage=response.usage,
+            latency_seconds=timer.elapsed,
+            attempts=attempts["n"],
+        )
+        logger.info(
+            "Linking picked %d table(s): %s",
+            len(linked.tables),
+            ", ".join(linked.tables) or "(none)",
+        )
+        return linked, metadata
+
+    @staticmethod
+    def _to_linked_schema(
+        data: dict[str, Any],
+        schema: Schema,
+        max_tables: int,
+    ) -> LinkedSchema:
+        """Check the parsed JSON against the real schema and build a LinkedSchema.
+
+        Tables/columns not in the source schema are dropped. This guards against
+        hallucinated names. Tables are capped at ``max_tables``.
+        """
+        valid_tables: list[str] = []
+        for name in data.get("tables", []):
+            tbl = schema.table(str(name))
+            if tbl is not None and tbl.name not in valid_tables:
+                valid_tables.append(tbl.name)
+            elif tbl is None:
+                logger.debug("Dropping hallucinated table '%s'.", name)
+        valid_tables = valid_tables[:max_tables]
+        valid_set = {t.lower() for t in valid_tables}
+
+        columns: list[LinkedColumn] = []
+        for col in data.get("columns", []):
+            if not isinstance(col, dict):
+                continue
+            table_name = str(col.get("table", ""))
+            column_name = str(col.get("column", ""))
+            tbl = schema.table(table_name)
+            if tbl is None or tbl.name.lower() not in valid_set:
+                continue
+            if tbl.column(column_name) is None:
+                logger.debug("Dropping hallucinated column '%s.%s'.", table_name, column_name)
+                continue
+            columns.append(
+                LinkedColumn(
+                    table=tbl.name,
+                    column=column_name,
+                    reason=str(col.get("reason", "")),
+                )
+            )
+
+        joins: list[LinkedJoin] = []
+        for j in data.get("joins", []):
+            if not isinstance(j, dict):
+                continue
+            joins.append(
+                LinkedJoin(
+                    left_table=str(j.get("left_table", "")),
+                    left_column=str(j.get("left_column", "")),
+                    right_table=str(j.get("right_table", "")),
+                    right_column=str(j.get("right_column", "")),
+                )
+            )
+
+        entities = [str(e) for e in data.get("entities", []) if str(e).strip()]
+        filters = [str(f) for f in data.get("filters", []) if str(f).strip()]
+
+        return LinkedSchema(
+            tables=valid_tables,
+            columns=columns,
+            joins=joins,
+            entities=entities,
+            filters=filters,
+            rationale=str(data.get("rationale", "")),
+        )

text2sql/pipeline/pipeline.py

@@ -0,0 +1,156 @@
+"""Public API: the :class:`Text2SQL` orchestrator.
+
+Combines Stage 1 (linking) and Stage 2 (generation) in one ``run`` call.
+Everything a caller may want to swap — settings, schema provider, prompt
+templates, and the per-stage LLM providers — is injectable.
+"""
+
+from __future__ import annotations
+
+from ..config import Settings, load_settings
+from ..errors import EmptyLinkingError, SchemaError
+from ..logging_utils import configure_logging, get_logger
+from ..prompts import PromptTemplates
+from ..providers.base import LLMProvider
+from ..providers.registry import build_provider
+from ..schema import FileSchemaProvider, Schema
+from ..schema.provider import SchemaProvider
+from ..types import RunMetadata, Text2SQLResult, _Timer
+from .generation import GenerationStage
+from .linking import LinkingStage
+
+logger = get_logger("pipeline")
+
+
+class Text2SQL:
+    """Two-stage, schema-agnostic Text-to-SQL engine.
+
+    Use :meth:`from_config` for the config-driven path, or the constructor
+    directly to inject custom parts.
+
+    Args:
+        settings: Full settings.
+        schema_provider: Gives the schema metadata.
+        prompts: Prompt template loader.
+        linking_provider: LLM provider for Stage 1.
+        generation_provider: LLM provider for Stage 2.
+    """
+
+    def __init__(
+        self,
+        settings: Settings,
+        schema_provider: SchemaProvider,
+        prompts: PromptTemplates,
+        linking_provider: LLMProvider,
+        generation_provider: LLMProvider,
+    ) -> None:
+        self._settings = settings
+        self._schema_provider = schema_provider
+        self._prompts = prompts
+        self._linking = LinkingStage(linking_provider, settings, prompts)
+        self._generation = GenerationStage(generation_provider, settings, prompts)
+
+        # Load the schema once, at construction, so problems show up early.
+        self._schema: Schema = schema_provider.load()
+        if not self._schema.tables:
+            raise SchemaError("Loaded schema has no tables.")
+
+    # Construction
+    @classmethod
+    def from_config(
+        cls,
+        *,
+        env_file: str | None = None,
+        config_path: str | None = None,
+        schema_provider: SchemaProvider | None = None,
+        prompts: PromptTemplates | None = None,
+        linking_provider: LLMProvider | None = None,
+        generation_provider: LLMProvider | None = None,
+    ) -> "Text2SQL":
+        """Build an engine from ``.env`` + ``config.toml``.
+
+        Any part can be overridden with a keyword argument to inject a custom
+        one. Anything left as ``None`` is built from config.
+
+        Raises:
+            ConfigError: If config is missing or invalid.
+            SchemaError: If the schema cannot be loaded.
+        """
+        settings = load_settings(env_file=env_file, config_path=config_path)
+        configure_logging(settings.general.log_level)
+
+        schema_provider = schema_provider or FileSchemaProvider(settings.schema_path)
+        # settings.prompt_dir may be None -> PromptTemplates uses packaged files.
+        prompts = prompts or PromptTemplates(settings.prompt_dir)
+
+        linking_provider = linking_provider or build_provider(
+            settings.linking_selection, timeout=settings.linking.timeout_seconds
+        )
+        generation_provider = generation_provider or build_provider(
+            settings.sql_selection, timeout=settings.generation.timeout_seconds
+        )
+
+        logger.debug(
+            "Text2SQL ready: linking=%s/%s generation=%s/%s dialect=%s",
+            linking_provider.name,
+            linking_provider.model,
+            generation_provider.name,
+            generation_provider.model,
+            settings.general.sql_dialect,
+        )
+        return cls(
+            settings=settings,
+            schema_provider=schema_provider,
+            prompts=prompts,
+            linking_provider=linking_provider,
+            generation_provider=generation_provider,
+        )
+
+    # Execution
+    def run(self, request: str) -> Text2SQLResult:
+        """Run the full pipeline for a natural-language ``request``.
+
+        Returns:
+            A :class:`~text2sql.types.Text2SQLResult` with the SQL string, the
+            Stage 1 linked schema, and run metadata.
+
+        Raises:
+            EmptyLinkingError: If Stage 1 links no tables/columns.
+            StructuredOutputError: If a stage's output cannot be parsed after
+                all retries.
+            ProviderError: If an LLM call fails.
+        """
+        if not request or not request.strip():
+            raise ValueError("request must be a non-empty string.")
+
+        logger.info("Running Text2SQL for request: %r", request)
+        with _Timer() as total_timer:
+            linked, linking_meta = self._linking.run(request, self._schema)
+
+            if linked.is_empty():
+                raise EmptyLinkingError(
+                    "Schema linking found no relevant table or column for the "
+                    "request. The request may not match the loaded schema."
+                )
+
+            sql, explanation, generation_meta = self._generation.run(
+                request, linked, self._schema
+            )
+
+        metadata = RunMetadata(
+            linking=linking_meta,
+            generation=generation_meta,
+            total_latency_seconds=total_timer.elapsed,
+        )
+        return Text2SQLResult(
+            sql=sql,
+            explanation=explanation,
+            linked_schema=linked,
+            metadata=metadata,
+            request=request,
+        )
+
+    @property
+    def schema(self) -> Schema:
+        """The loaded schema. Read-only access for callers."""
+        return self._schema

text2sql/prompts/__init__.py

@@ -0,0 +1,96 @@
+"""Prompt template loading.
+
+Templates are editable ``.txt`` files, kept outside the code. Each file has a
+``SYSTEM:`` section and a ``USER:`` section. Both support ``str.format``-style
+``{placeholder}`` substitution.
+"""
+
+from __future__ import annotations
+
+from importlib import resources
+from pathlib import Path
+
+from ..errors import ConfigError
+
+#: Package that holds the built-in prompt templates.
+_PACKAGE = "text2sql.prompts"
+
+_SYSTEM_MARKER = "SYSTEM:"
+_USER_MARKER = "USER:"
+
+
+class PromptTemplate:
+    """A parsed template with a system and a user section."""
+
+    def __init__(self, system: str, user: str) -> None:
+        self.system = system
+        self.user = user
+
+    def render(self, **kwargs: object) -> tuple[str, str]:
+        """Return ``(system, user)`` with placeholders filled.
+
+        Raises:
+            ConfigError: If the template uses a placeholder that is not given.
+        """
+        try:
+            return (
+                self.system.format(**kwargs),
+                self.user.format(**kwargs),
+            )
+        except KeyError as exc:
+            raise ConfigError(
+                f"Prompt template uses an unknown placeholder {exc}."
+            ) from exc
+
+
+class PromptTemplates:
+    """Loads named templates.
+
+    This is the DI seam for prompts: pass a different directory (or subclass) to
+    swap templates without touching code.
+
+    Args:
+        directory: Directory with the ``.txt`` files. When ``None`` (default),
+            the templates packaged inside ``text2sql.prompts`` are used, so the
+            library works when installed as a package, from any working
+            directory.
+    """
+
+    def __init__(self, directory: str | Path | None = None) -> None:
+        self._dir = Path(directory).expanduser() if directory is not None else None
+
+    def load(self, filename: str) -> PromptTemplate:
+        """Load and parse a template by name.
+
+        Raises:
+            ConfigError: If the template is missing or lacks the two sections.
+        """
+        if self._dir is not None:
+            path = self._dir / filename
+            if not path.is_file():
+                raise ConfigError(f"Prompt template not found: '{path}'.")
+            return self._parse(path.read_text(encoding="utf-8"), str(path))
+
+        # Fall back to the templates packaged with the library.
+        try:
+            resource = resources.files(_PACKAGE).joinpath(filename)
+            text = resource.read_text(encoding="utf-8")
+        except (FileNotFoundError, ModuleNotFoundError) as exc:
+            raise ConfigError(
+                f"Packaged prompt template '{filename}' not found in {_PACKAGE}."
+            ) from exc
+        return self._parse(text, f"{_PACKAGE}/{filename}")
+
+    @staticmethod
+    def _parse(text: str, path: str) -> PromptTemplate:
+        if _SYSTEM_MARKER not in text or _USER_MARKER not in text:
+            raise ConfigError(
+                f"Prompt template '{path}' must have a '{_SYSTEM_MARKER}' and a "
+                f"'{_USER_MARKER}' section."
+            )
+        _, _, remainder = text.partition(_SYSTEM_MARKER)
+        system_part, _, user_part = remainder.partition(_USER_MARKER)
+        return PromptTemplate(system=system_part.strip(), user=user_part.strip())
+
+
+__all__ = ["PromptTemplate", "PromptTemplates"]

text2sql/prompts/generation.txt

@@ -0,0 +1,28 @@
+SYSTEM:
+You are an expert SQL generator. You get a user request and a reduced schema
+(only the relevant tables/columns/joins picked by an upstream linking step).
+Produce one correct SQL query in the target dialect.
+
+Rules:
+- Use only the tables and columns in the reduced schema.
+- Write idiomatic SQL for the target dialect: {dialect}.
+- Prefer explicit JOINs using the given relations.
+- Do not run the query. Only produce it.
+- Return the answer strictly as the requested JSON object with two fields:
+  "sql" (the query) and "explanation" (one or two sentences).
+
+USER:
+Target SQL dialect: {dialect}
+
+Reduced schema (relevant subset):
+--------------------------------
+{linked_schema}
+--------------------------------
+
+Detected entities: {entities}
+Filter hints: {filters}
+
+User request:
+"{request}"
+
+Write the SQL query that answers the request.

text2sql/prompts/linking.txt

@@ -0,0 +1,28 @@
+SYSTEM:
+You are a schema-linking assistant for a Text-to-SQL system. You get a user
+request and the full metadata of a database schema. Your job: find the subset of
+the schema that is relevant to the request.
+
+Rules:
+- Pick only the tables and columns that are actually needed.
+- Propose the join paths (from the foreign keys) needed to connect them.
+- Extract the entities and values named in the request.
+- Note any filter conditions the request implies.
+- Return at most {max_tables} tables.
+- Never invent a table or column that is not in the metadata.
+- Give a short rationale for your choice.
+
+Return the answer strictly as the requested JSON object.
+
+USER:
+Database schema metadata:
+--------------------------------
+{schema}
+--------------------------------
+
+Target SQL dialect: {dialect}
+
+User request:
+"{request}"
+
+Find the relevant tables, columns, joins, entities, and filters.

text2sql/providers/__init__.py

@@ -0,0 +1,19 @@
+"""LLM client layer."""
+
+from __future__ import annotations
+
+from .base import (
+    LLMProvider,
+    get_provider_class,
+    register_provider,
+    registered_providers,
+)
+from .registry import build_provider
+
+__all__ = [
+    "LLMProvider",
+    "register_provider",
+    "registered_providers",
+    "get_provider_class",
+    "build_provider",
+]

text2sql/providers/anthropic_provider.py

@@ -0,0 +1,123 @@
+"""Anthropic provider.
+
+Uses the Messages API. Structured output is forced with tool-calling: one tool
+whose ``input_schema`` is the JSON schema, plus ``tool_choice`` to force it. The
+``anthropic`` SDK is imported lazily.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from ..errors import ProviderError
+from ..logging_utils import get_logger
+from ..types import LLMResponse, TokenUsage
+from .base import LLMProvider, register_provider
+
+logger = get_logger("providers.anthropic")
+
+
+@register_provider("anthropic")
+class AnthropicProvider(LLMProvider):
+    """Provider backed by the Anthropic Messages API."""
+
+    def _client(self) -> Any:
+        try:
+            from anthropic import Anthropic
+        except ImportError as exc:  # pragma: no cover - dependency guard
+            raise ProviderError(
+                "The 'anthropic' package is required for the Anthropic provider. "
+                "Install it with 'pip install anthropic'."
+            ) from exc
+        kwargs: dict[str, Any] = {"api_key": self._api_key, "timeout": self._timeout}
+        if self._base_url:
+            kwargs["base_url"] = self._base_url
+        return Anthropic(**kwargs)
+
+    @staticmethod
+    def _usage(response: Any) -> TokenUsage:
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return TokenUsage()
+        prompt = getattr(usage, "input_tokens", 0) or 0
+        completion = getattr(usage, "output_tokens", 0) or 0
+        return TokenUsage(
+            prompt_tokens=prompt,
+            completion_tokens=completion,
+            total_tokens=prompt + completion,
+        )
+
+    def complete(
+        self,
+        *,
+        system: str,
+        user: str,
+        temperature: float,
+        top_p: float,
+        max_tokens: int,
+    ) -> LLMResponse:
+        client = self._client()
+        logger.debug("Anthropic completion request (model=%s)", self.model)
+        try:
+            response = client.messages.create(
+                model=self.model,
+                system=system,
+                messages=[{"role": "user", "content": user}],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise ProviderError(f"Anthropic request failed: {exc}") from exc
+
+        text = "".join(
+            block.text for block in response.content if getattr(block, "type", "") == "text"
+        )
+        return LLMResponse(text=text, usage=self._usage(response), model=self.model, raw=response)
+
+    def complete_json(
+        self,
+        *,
+        system: str,
+        user: str,
+        json_schema: dict[str, Any],
+        schema_name: str,
+        temperature: float,
+        top_p: float,
+        max_tokens: int,
+    ) -> LLMResponse:
+        client = self._client()
+        tool = {
+            "name": schema_name,
+            "description": "Return the structured result with this tool.",
+            "input_schema": json_schema,
+        }
+        logger.debug("Anthropic structured (tool-use) request (model=%s)", self.model)
+        try:
+            response = client.messages.create(
+                model=self.model,
+                system=system,
+                messages=[{"role": "user", "content": user}],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                tools=[tool],
+                tool_choice={"type": "tool", "name": schema_name},
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise ProviderError(f"Anthropic structured request failed: {exc}") from exc
+
+        # Take the tool-use block's input as the JSON payload.
+        for block in response.content:
+            if getattr(block, "type", "") == "tool_use":
+                text = json.dumps(block.input)
+                return LLMResponse(
+                    text=text, usage=self._usage(response), model=self.model, raw=response
+                )
+
+        # No tool call came back. Return whatever text there is, for the parser.
+        text = "".join(
+            block.text for block in response.content if getattr(block, "type", "") == "text"
+        )
+        return LLMResponse(text=text, usage=self._usage(response), model=self.model, raw=response)