text2sql-engine 0.1.0__py3-none-any.whl

text2sql/providers/base.py

@@ -0,0 +1,131 @@
+"""Provider-agnostic LLM client layer.
+
+An :class:`LLMProvider` turns a prompt into text. Two entry points:
+
+* :meth:`complete` — free-form text. Used by Stage 2.
+* :meth:`complete_json` — output forced to a JSON object that matches a given
+  JSON schema. Used by Stage 1. Providers with native structured output
+  (``response_format`` / tool-calling) should use it. The base class gives a
+  plain-text fallback for the rest.
+
+Adding a provider is one subclass plus the ``@register_provider`` decorator.
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from typing import Any, Callable
+
+from ..logging_utils import get_logger
+from ..types import LLMResponse
+
+logger = get_logger("providers")
+
+
+class LLMProvider(ABC):
+    """Base class for all providers.
+
+    Args:
+        model: Model id for this instance.
+        api_key: Secret key. Never logged.
+        base_url: Optional endpoint override, for OpenAI-compatible gateways.
+        timeout: Request timeout in seconds.
+    """
+
+    #: Registry name. Set by ``@register_provider``.
+    name: str = ""
+
+    def __init__(
+        self,
+        model: str,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        timeout: float = 60.0,
+    ) -> None:
+        self.model = model
+        self._api_key = api_key
+        self._base_url = base_url
+        self._timeout = timeout
+
+    @abstractmethod
+    def complete(
+        self,
+        *,
+        system: str,
+        user: str,
+        temperature: float,
+        top_p: float,
+        max_tokens: int,
+    ) -> LLMResponse:
+        """Return a free-form text completion."""
+        raise NotImplementedError
+
+    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:
+        """Return a completion forced to a JSON object.
+
+        Default: append the JSON schema to the prompt and rely on text parsing.
+        Providers with native structured output should override this for a
+        stronger guarantee.
+        """
+        augmented_user = (
+            f"{user}\n\n"
+            f"Respond with a single JSON object that matches this JSON schema. "
+            f"Do not wrap it in markdown fences and do not add any comment.\n"
+            f"JSON schema:\n{json.dumps(json_schema)}"
+        )
+        return self.complete(
+            system=system,
+            user=augmented_user,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+        )
+
+
+# Registry / factory.
+_REGISTRY: dict[str, type[LLMProvider]] = {}
+
+
+def register_provider(name: str) -> Callable[[type[LLMProvider]], type[LLMProvider]]:
+    """Class decorator. Registers a provider under ``name``."""
+
+    def decorator(cls: type[LLMProvider]) -> type[LLMProvider]:
+        cls.name = name
+        _REGISTRY[name.lower()] = cls
+        return cls
+
+    return decorator
+
+
+def registered_providers() -> list[str]:
+    """Return the registered provider names, sorted."""
+    return sorted(_REGISTRY)
+
+
+def get_provider_class(name: str) -> type[LLMProvider]:
+    """Look up a provider class by name.
+
+    Raises:
+        ProviderError: If no provider is registered under ``name``.
+    """
+    from ..errors import ProviderError
+
+    cls = _REGISTRY.get(name.lower())
+    if cls is None:
+        raise ProviderError(
+            f"Unknown provider '{name}'. Registered providers: "
+            f"{', '.join(registered_providers()) or '(none)'}."
+        )
+    return cls

text2sql/providers/openai_provider.py

@@ -0,0 +1,125 @@
+"""OpenAI-compatible provider.
+
+Works with the OpenAI API and any OpenAI-compatible gateway (via ``base_url``).
+Uses ``response_format`` JSON-schema mode for structured output when available.
+The ``openai`` SDK is imported lazily, so the library imports without it.
+"""
+
+from __future__ import annotations
+
+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.openai")
+
+
+@register_provider("openai")
+class OpenAIProvider(LLMProvider):
+    """Provider backed by the OpenAI Chat Completions API."""
+
+    def _client(self) -> Any:
+        try:
+            from openai import OpenAI
+        except ImportError as exc:  # pragma: no cover - dependency guard
+            raise ProviderError(
+                "The 'openai' package is required for the OpenAI provider. "
+                "Install it with 'pip install openai'."
+            ) from exc
+        # base_url=None lets the SDK use its default endpoint.
+        return OpenAI(api_key=self._api_key, base_url=self._base_url, timeout=self._timeout)
+
+    @staticmethod
+    def _usage(response: Any) -> TokenUsage:
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return TokenUsage()
+        return TokenUsage(
+            prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
+            completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
+            total_tokens=getattr(usage, "total_tokens", 0) or 0,
+        )
+
+    def complete(
+        self,
+        *,
+        system: str,
+        user: str,
+        temperature: float,
+        top_p: float,
+        max_tokens: int,
+    ) -> LLMResponse:
+        client = self._client()
+        logger.debug("OpenAI completion request (model=%s)", self.model)
+        try:
+            response = client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system},
+                    {"role": "user", "content": user},
+                ],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+            )
+        except Exception as exc:  # noqa: BLE001 - normalize SDK errors
+            raise ProviderError(f"OpenAI request failed: {exc}") from exc
+
+        text = response.choices[0].message.content or ""
+        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()
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": schema_name,
+                "schema": json_schema,
+                "strict": True,
+            },
+        }
+        logger.debug("OpenAI structured request (model=%s)", self.model)
+        try:
+            response = client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system},
+                    {"role": "user", "content": user},
+                ],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                response_format=response_format,
+            )
+        except Exception as exc:  # noqa: BLE001
+            # Some OpenAI-compatible gateways do not support json_schema mode.
+            # Fall back to schema-in-prompt + text parsing.
+            logger.warning(
+                "OpenAI structured output failed (%s). Falling back to text parsing.",
+                exc,
+            )
+            return super().complete_json(
+                system=system,
+                user=user,
+                json_schema=json_schema,
+                schema_name=schema_name,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+            )
+
+        text = response.choices[0].message.content or ""
+        return LLMResponse(text=text, usage=self._usage(response), model=self.model, raw=response)

text2sql/providers/registry.py

@@ -0,0 +1,38 @@
+"""Factory that builds a provider from a :class:`ProviderSelection`.
+
+Importing this module registers the built-in providers as a side effect.
+"""
+
+from __future__ import annotations
+
+from ..config import ProviderSelection
+from .base import LLMProvider, get_provider_class, registered_providers
+
+# Importing these runs their @register_provider decorators.
+from . import anthropic_provider as _anthropic  # noqa: F401
+from . import openai_provider as _openai  # noqa: F401
+
+
+def build_provider(selection: ProviderSelection, *, timeout: float) -> LLMProvider:
+    """Build the provider described by ``selection``.
+
+    Args:
+        selection: Provider/model/base_url/api_key from ``.env``.
+        timeout: Request timeout in seconds, from the stage config.
+
+    Returns:
+        A ready :class:`LLMProvider`.
+
+    Raises:
+        ProviderError: If the provider name is not registered.
+    """
+    cls = get_provider_class(selection.provider)
+    return cls(
+        model=selection.model,
+        api_key=selection.api_key,
+        base_url=selection.base_url,
+        timeout=timeout,
+    )
+
+
+__all__ = ["build_provider", "registered_providers"]

text2sql/schema/__init__.py

@@ -0,0 +1,18 @@
+"""Schema models and providers."""
+
+from __future__ import annotations
+
+from .file_provider import FileSchemaProvider
+from .models import Column, ForeignKey, Schema, Table
+from .provider import SchemaProvider
+from .serialize import schema_to_prompt
+
+__all__ = [
+    "Column",
+    "ForeignKey",
+    "Schema",
+    "Table",
+    "SchemaProvider",
+    "FileSchemaProvider",
+    "schema_to_prompt",
+]

text2sql/schema/file_provider.py

@@ -0,0 +1,140 @@
+"""File-backed SchemaProvider for JSON and YAML.
+
+The file format (documented in ``examples/schema.yaml``) is::
+
+    name: my_database            # optional
+    tables:
+      - name: customers
+        description: "..."
+        columns:
+          - name: id
+            type: integer
+            description: "..."
+            primary_key: true
+            sample_values: [1, 2, 3]
+        foreign_keys:
+          - column: category_id
+            references_table: categories
+            references_column: id
+
+JSON uses the same structure. Format is picked from the file extension.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from ..errors import SchemaError
+from .models import Column, ForeignKey, Schema, Table
+from .provider import SchemaProvider
+
+
+class FileSchemaProvider(SchemaProvider):
+    """Load schema metadata from a JSON or YAML file.
+
+    Format comes from the extension (``.json``, ``.yaml``, ``.yml``).
+
+    Args:
+        path: Path to the metadata file.
+    """
+
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path).expanduser()
+
+    def load(self) -> Schema:
+        if not self._path.is_file():
+            raise SchemaError(f"Schema file not found: '{self._path}'.")
+
+        raw_text = self._path.read_text(encoding="utf-8")
+        suffix = self._path.suffix.lower()
+
+        try:
+            if suffix in (".yaml", ".yml"):
+                data = self._parse_yaml(raw_text)
+            elif suffix == ".json":
+                data = json.loads(raw_text)
+            else:
+                raise SchemaError(
+                    f"Unsupported schema file extension '{suffix}'. "
+                    f"Use .json, .yaml, or .yml."
+                )
+        except SchemaError:
+            raise
+        except Exception as exc:  # noqa: BLE001 - normalize parser errors
+            raise SchemaError(f"Failed to parse schema file '{self._path}': {exc}") from exc
+
+        return self._build_schema(data)
+
+    @staticmethod
+    def _parse_yaml(text: str) -> Any:
+        try:
+            import yaml
+        except ImportError as exc:  # pragma: no cover - dependency guard
+            raise SchemaError(
+                "PyYAML is required to load YAML schema files. "
+                "Install it with 'pip install PyYAML'."
+            ) from exc
+        return yaml.safe_load(text)
+
+    @staticmethod
+    def _build_schema(data: Any) -> Schema:
+        if not isinstance(data, dict):
+            raise SchemaError("Schema root must be a mapping/object.")
+
+        raw_tables = data.get("tables")
+        if not isinstance(raw_tables, list) or not raw_tables:
+            raise SchemaError("Schema must have a non-empty 'tables' list.")
+
+        tables: list[Table] = []
+        for entry in raw_tables:
+            if not isinstance(entry, dict):
+                raise SchemaError("Each table entry must be a mapping/object.")
+            name = entry.get("name")
+            if not name:
+                raise SchemaError("Every table needs a 'name'.")
+
+            columns: list[Column] = []
+            for col in entry.get("columns", []) or []:
+                if not isinstance(col, dict) or not col.get("name"):
+                    raise SchemaError(
+                        f"Invalid column in table '{name}': every column needs a 'name'."
+                    )
+                columns.append(
+                    Column(
+                        name=str(col["name"]),
+                        type=str(col.get("type", "")),
+                        description=str(col.get("description", "")),
+                        primary_key=bool(col.get("primary_key", False)),
+                        sample_values=list(col.get("sample_values", []) or []),
+                    )
+                )
+
+            foreign_keys: list[ForeignKey] = []
+            for fk in entry.get("foreign_keys", []) or []:
+                if not isinstance(fk, dict):
+                    raise SchemaError(f"Invalid foreign_key in table '{name}'.")
+                try:
+                    foreign_keys.append(
+                        ForeignKey(
+                            column=str(fk["column"]),
+                            references_table=str(fk["references_table"]),
+                            references_column=str(fk["references_column"]),
+                        )
+                    )
+                except KeyError as exc:
+                    raise SchemaError(
+                        f"Foreign key in table '{name}' is missing field {exc}."
+                    ) from exc
+
+            tables.append(
+                Table(
+                    name=str(name),
+                    description=str(entry.get("description", "")),
+                    columns=columns,
+                    foreign_keys=foreign_keys,
+                )
+            )
+
+        return Schema(name=str(data.get("name", "")), tables=tables)

text2sql/schema/models.py

@@ -0,0 +1,84 @@
+"""Schema metadata models.
+
+Schema-agnostic containers. They describe any relational schema, loaded from a
+file or from live introspection. Nothing here is tied to one database product.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class Column:
+    """A column in a table."""
+
+    name: str
+    type: str
+    description: str = ""
+    primary_key: bool = False
+    sample_values: list[Any] = field(default_factory=list)
+
+
+@dataclass
+class ForeignKey:
+    """A foreign key: a column pointing to another table's column."""
+
+    column: str
+    references_table: str
+    references_column: str
+
+
+@dataclass
+class Table:
+    """A table with its columns, primary keys, and foreign keys."""
+
+    name: str
+    description: str = ""
+    columns: list[Column] = field(default_factory=list)
+    foreign_keys: list[ForeignKey] = field(default_factory=list)
+
+    @property
+    def primary_keys(self) -> list[str]:
+        """Names of the primary-key columns."""
+        return [c.name for c in self.columns if c.primary_key]
+
+    def column(self, name: str) -> Column | None:
+        """Return the named column, or ``None``. Case-insensitive."""
+        lowered = name.lower()
+        for col in self.columns:
+            if col.name.lower() == lowered:
+                return col
+        return None
+
+
+@dataclass
+class Schema:
+    """A full schema: a list of tables."""
+
+    tables: list[Table] = field(default_factory=list)
+    name: str = ""
+
+    def table(self, name: str) -> Table | None:
+        """Return the named table, or ``None``. Case-insensitive."""
+        lowered = name.lower()
+        for tbl in self.tables:
+            if tbl.name.lower() == lowered:
+                return tbl
+        return None
+
+    @property
+    def table_names(self) -> list[str]:
+        return [t.name for t in self.tables]
+
+    def subset(self, table_names: list[str]) -> "Schema":
+        """Return a new Schema with only the named tables.
+
+        Matches names case-insensitively. Unknown names are skipped.
+        """
+        wanted = {n.lower() for n in table_names}
+        return Schema(
+            name=self.name,
+            tables=[t for t in self.tables if t.name.lower() in wanted],
+        )

text2sql/schema/provider.py

@@ -0,0 +1,30 @@
+"""SchemaProvider.
+
+This is the seam that keeps the pipeline schema-agnostic. The default loads
+from a file (see :mod:`.file_provider`). A caller can inject any subclass
+instead, for example one that reads a live database, as long as it returns a
+:class:`~text2sql.schema.models.Schema`.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from .models import Schema
+
+
+class SchemaProvider(ABC):
+    """Gives schema metadata to the pipeline.
+
+    Subclasses return a full :class:`Schema`. The pipeline does not care where
+    it comes from.
+    """
+
+    @abstractmethod
+    def load(self) -> Schema:
+        """Load and return the schema.
+
+        Raises:
+            SchemaError: If the schema cannot be built or is invalid.
+        """
+        raise NotImplementedError

text2sql/schema/serialize.py

@@ -0,0 +1,55 @@
+"""Render a :class:`Schema` to text for LLM prompts.
+
+Kept apart from the models so prompt formatting can change without touching the
+data classes.
+"""
+
+from __future__ import annotations
+
+from .models import Schema
+
+
+def schema_to_prompt(schema: Schema, *, max_sample_values: int = 5) -> str:
+    """Render a schema as text for humans and LLMs.
+
+    Includes table descriptions, columns (type + description), primary keys,
+    foreign keys, and a capped list of sample values.
+
+    Args:
+        schema: The schema to render.
+        max_sample_values: Max sample values shown per column.
+    """
+    lines: list[str] = []
+    if schema.name:
+        lines.append(f"Database: {schema.name}")
+        lines.append("")
+
+    for table in schema.tables:
+        lines.append(f"Table: {table.name}")
+        if table.description:
+            lines.append(f"  Description: {table.description}")
+        pks = table.primary_keys
+        if pks:
+            lines.append(f"  Primary key: {', '.join(pks)}")
+
+        lines.append("  Columns:")
+        for col in table.columns:
+            parts = [f"    - {col.name} ({col.type})"]
+            if col.description:
+                parts.append(f": {col.description}")
+            lines.append("".join(parts))
+            if col.sample_values:
+                sample = col.sample_values[:max_sample_values]
+                rendered = ", ".join(repr(v) for v in sample)
+                lines.append(f"        samples: {rendered}")
+
+        if table.foreign_keys:
+            lines.append("  Foreign keys:")
+            for fk in table.foreign_keys:
+                lines.append(
+                    f"    - {table.name}.{fk.column} -> "
+                    f"{fk.references_table}.{fk.references_column}"
+                )
+        lines.append("")
+
+    return "\n".join(lines).rstrip() + "\n"