text2sql-engine 0.1.0__py3-none-any.whl

text2sql/__init__.py

@@ -0,0 +1,67 @@
+"""text2sql — schema-agnostic, config-driven Text-to-SQL library.
+
+Entry point::
+
+    from text2sql import Text2SQL
+    engine = Text2SQL.from_config()
+    result = engine.run("get me the top 5 best-selling products last month")
+    print(result.sql)
+    print(result.linked_schema)
+    print(result.metadata)
+
+The library only produces a SQL string. It never connects to or runs against a
+database.
+"""
+
+from __future__ import annotations
+
+from .errors import (
+    ConfigError,
+    EmptyLinkingError,
+    ProviderError,
+    SchemaError,
+    StructuredOutputError,
+    Text2SQLError,
+)
+from .pipeline import Text2SQL
+from .prompts import PromptTemplates
+from .providers import LLMProvider, register_provider
+from .schema import FileSchemaProvider, Schema, SchemaProvider
+from .types import (
+    LinkedColumn,
+    LinkedJoin,
+    LinkedSchema,
+    RunMetadata,
+    StageMetadata,
+    Text2SQLResult,
+    TokenUsage,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Text2SQL",
+    # Schema
+    "Schema",
+    "SchemaProvider",
+    "FileSchemaProvider",
+    # Providers / DI
+    "LLMProvider",
+    "register_provider",
+    "PromptTemplates",
+    # Result types
+    "Text2SQLResult",
+    "LinkedSchema",
+    "LinkedColumn",
+    "LinkedJoin",
+    "RunMetadata",
+    "StageMetadata",
+    "TokenUsage",
+    # Errors
+    "Text2SQLError",
+    "ConfigError",
+    "SchemaError",
+    "ProviderError",
+    "StructuredOutputError",
+    "EmptyLinkingError",
+]

text2sql/config.py

@@ -0,0 +1,251 @@
+"""Typed config.
+
+Two separate sources merge into one typed :class:`Settings`:
+
+* ``.env``        -> secrets, paths, provider/model selection.
+* ``config.toml`` -> behavioral parameters. No magic numbers in code.
+
+Loading fails fast with a clear :class:`~text2sql.errors.ConfigError` when a
+required ``.env`` variable is missing or ``config.toml`` is broken.
+"""
+
+from __future__ import annotations
+
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+
+from .errors import ConfigError
+
+# .env keys, declared once so they are easy to audit and document.
+ENV_SCHEMA_PATH = "SCHEMA_PATH"
+ENV_PROMPT_DIR = "PROMPT_DIR"
+ENV_CONFIG_PATH = "CONFIG_PATH"
+
+ENV_LINKING_PROVIDER = "LINKING_PROVIDER"
+ENV_LINKING_MODEL = "LINKING_MODEL"
+ENV_LINKING_BASE_URL = "LINKING_BASE_URL"
+
+ENV_SQL_PROVIDER = "SQL_PROVIDER"
+ENV_SQL_MODEL = "SQL_MODEL"
+ENV_SQL_BASE_URL = "SQL_BASE_URL"
+
+# API-key variable per provider. A provider needs its key only if a stage uses it.
+PROVIDER_KEY_ENV = {
+    "openai": "OPENAI_API_KEY",
+    "anthropic": "ANTHROPIC_API_KEY",
+}
+
+
+@dataclass
+class StageSettings:
+    """Sampling and behavior for one stage."""
+
+    temperature: float = 0.0
+    top_p: float = 1.0
+    max_tokens: int = 1024
+    prompt_template: str = ""
+    timeout_seconds: float = 60.0
+    # Used by linking only. Generation ignores it.
+    max_tables: int = 8
+
+
+@dataclass
+class RetrySettings:
+    """Retry/backoff, shared by both stages."""
+
+    max_attempts: int = 3
+    backoff_seconds: float = 1.0
+    backoff_factor: float = 2.0
+
+
+@dataclass
+class GeneralSettings:
+    """Top-level behavior."""
+
+    sql_dialect: str = "postgres"
+    log_level: str = "INFO"
+    strict_json: bool = True
+
+
+@dataclass
+class ProviderSelection:
+    """Which provider/model/endpoint a stage uses. Comes from ``.env``."""
+
+    provider: str
+    model: str
+    base_url: str | None = None
+    api_key: str | None = None
+
+
+@dataclass
+class Settings:
+    """Full config for a run."""
+
+    # Paths (.env)
+    schema_path: Path
+    # None -> use the prompt templates packaged inside text2sql.prompts.
+    prompt_dir: Path | None
+
+    # Provider selection (.env)
+    linking_selection: ProviderSelection
+    sql_selection: ProviderSelection
+
+    # Behavior (config.toml)
+    general: GeneralSettings = field(default_factory=GeneralSettings)
+    linking: StageSettings = field(default_factory=StageSettings)
+    generation: StageSettings = field(default_factory=StageSettings)
+    retry: RetrySettings = field(default_factory=RetrySettings)
+
+
+def _require_env(key: str) -> str:
+    """Return an env var, or raise a clear ConfigError if it is missing."""
+    value = os.environ.get(key)
+    if value is None or value.strip() == "":
+        raise ConfigError(
+            f"Required environment variable '{key}' is missing or empty. "
+            f"See .env.example for the full list."
+        )
+    return value
+
+
+def _optional_env(key: str) -> str | None:
+    value = os.environ.get(key)
+    if value is None or value.strip() == "":
+        return None
+    return value
+
+
+def _load_toml(config_path: Path) -> dict[str, Any]:
+    if not config_path.is_file():
+        raise ConfigError(f"config.toml not found at '{config_path}'.")
+    try:
+        with config_path.open("rb") as fh:
+            return tomllib.load(fh)
+    except tomllib.TOMLDecodeError as exc:
+        raise ConfigError(f"Failed to parse config.toml at '{config_path}': {exc}") from exc
+
+
+def _stage_from_toml(
+    table: dict[str, Any],
+    *,
+    default_template: str,
+    include_max_tables: bool,
+) -> StageSettings:
+    """Build a StageSettings. Falls back to dataclass defaults per key."""
+    defaults = StageSettings()
+    return StageSettings(
+        temperature=float(table.get("temperature", defaults.temperature)),
+        top_p=float(table.get("top_p", defaults.top_p)),
+        max_tokens=int(table.get("max_tokens", defaults.max_tokens)),
+        prompt_template=str(table.get("prompt_template", default_template)),
+        timeout_seconds=float(table.get("timeout_seconds", defaults.timeout_seconds)),
+        max_tables=(
+            int(table.get("max_tables", defaults.max_tables))
+            if include_max_tables
+            else defaults.max_tables
+        ),
+    )
+
+
+def _selection_for_stage(
+    provider_key: str,
+    model_key: str,
+    base_url_key: str,
+) -> ProviderSelection:
+    provider = _require_env(provider_key).strip().lower()
+    model = _require_env(model_key)
+    base_url = _optional_env(base_url_key)
+
+    # Pick the API key for the selected provider, if we know one for it.
+    key_env = PROVIDER_KEY_ENV.get(provider)
+    api_key = _optional_env(key_env) if key_env else None
+
+    return ProviderSelection(
+        provider=provider,
+        model=model,
+        base_url=base_url,
+        api_key=api_key,
+    )
+
+
+def load_settings(
+    env_file: str | os.PathLike[str] | None = None,
+    config_path: str | os.PathLike[str] | None = None,
+) -> Settings:
+    """Load and merge ``.env`` and ``config.toml`` into a :class:`Settings`.
+
+    Args:
+        env_file: Optional path to a ``.env`` file. If ``None``, dotenv's default
+            lookup is used (nearest ``.env`` walking up from CWD). Real process
+            env vars always win over the file.
+        config_path: Optional path to ``config.toml``. If ``None``, uses the
+            ``CONFIG_PATH`` env var, else ``./config.toml``.
+
+    Returns:
+        A full, typed settings object.
+
+    Raises:
+        ConfigError: If a required variable is missing or a file is invalid.
+    """
+    # override=False -> real environment wins over the file. Safer for secrets.
+    load_dotenv(dotenv_path=env_file, override=False)
+
+    schema_path = Path(_require_env(ENV_SCHEMA_PATH)).expanduser()
+    # PROMPT_DIR is optional. If unset, packaged prompt templates are used.
+    prompt_dir_raw = _optional_env(ENV_PROMPT_DIR)
+    prompt_dir = Path(prompt_dir_raw).expanduser() if prompt_dir_raw else None
+
+    linking_selection = _selection_for_stage(
+        ENV_LINKING_PROVIDER, ENV_LINKING_MODEL, ENV_LINKING_BASE_URL
+    )
+    sql_selection = _selection_for_stage(
+        ENV_SQL_PROVIDER, ENV_SQL_MODEL, ENV_SQL_BASE_URL
+    )
+
+    resolved_config_path = Path(
+        config_path
+        or _optional_env(ENV_CONFIG_PATH)
+        or "config.toml"
+    ).expanduser()
+    toml_data = _load_toml(resolved_config_path)
+
+    general_tbl = toml_data.get("general", {})
+    general = GeneralSettings(
+        sql_dialect=str(general_tbl.get("sql_dialect", GeneralSettings.sql_dialect)),
+        log_level=str(general_tbl.get("log_level", GeneralSettings.log_level)),
+        strict_json=bool(general_tbl.get("strict_json", GeneralSettings.strict_json)),
+    )
+
+    linking = _stage_from_toml(
+        toml_data.get("linking", {}),
+        default_template="linking.txt",
+        include_max_tables=True,
+    )
+    generation = _stage_from_toml(
+        toml_data.get("generation", {}),
+        default_template="generation.txt",
+        include_max_tables=False,
+    )
+
+    retry_tbl = toml_data.get("retry", {})
+    retry = RetrySettings(
+        max_attempts=int(retry_tbl.get("max_attempts", RetrySettings.max_attempts)),
+        backoff_seconds=float(retry_tbl.get("backoff_seconds", RetrySettings.backoff_seconds)),
+        backoff_factor=float(retry_tbl.get("backoff_factor", RetrySettings.backoff_factor)),
+    )
+
+    return Settings(
+        schema_path=schema_path,
+        prompt_dir=prompt_dir,
+        linking_selection=linking_selection,
+        sql_selection=sql_selection,
+        general=general,
+        linking=linking,
+        generation=generation,
+        retry=retry,
+    )

text2sql/errors.py

@@ -0,0 +1,43 @@
+"""Exceptions used across the library.
+
+Everything subclasses ``Text2SQLError``. Catch that to catch all of them.
+"""
+
+from __future__ import annotations
+
+
+class Text2SQLError(Exception):
+    """Base class for every error this library raises."""
+
+
+class ConfigError(Text2SQLError):
+    """Config is missing or invalid.
+
+    Example: a required ``.env`` variable is not set, or ``config.toml`` is
+    broken, or a provider/model selection is incomplete.
+    """
+
+
+class SchemaError(Text2SQLError):
+    """Schema metadata cannot be loaded or is not valid."""
+
+
+class ProviderError(Text2SQLError):
+    """An LLM provider failed to build or a call failed.
+
+    Wraps SDK/transport errors so callers do not deal with SDK-specific types.
+    """
+
+
+class StructuredOutputError(Text2SQLError):
+    """LLM output could not be parsed as the expected JSON.
+
+    Raised only after all retries are used.
+    """
+
+
+class EmptyLinkingError(Text2SQLError):
+    """Stage 1 linked no tables or columns.
+
+    A valid but empty match means we cannot generate SQL.
+    """

text2sql/logging_utils.py

@@ -0,0 +1,54 @@
+"""Logging helpers.
+
+The library uses stdlib :mod:`logging` only, never ``print``, and never logs
+secrets. Call :func:`configure_logging` once with the level from config. Use
+:func:`get_logger` everywhere else.
+"""
+
+from __future__ import annotations
+
+import logging
+
+_LIBRARY_LOGGER_NAME = "text2sql"
+_configured = False
+
+
+def configure_logging(level: str = "INFO") -> None:
+    """Set up the library logger.
+
+    Attaches one stream handler to the ``text2sql`` logger, not the root logger,
+    so the host application's logging is left alone. Idempotent: later calls
+    only change the level.
+
+    Args:
+        level: A standard level name, e.g. ``"DEBUG"``, ``"INFO"``.
+    """
+    global _configured
+    logger = logging.getLogger(_LIBRARY_LOGGER_NAME)
+    numeric_level = logging.getLevelName(level.upper())
+    if not isinstance(numeric_level, int):
+        numeric_level = logging.INFO
+    logger.setLevel(numeric_level)
+
+    if not _configured:
+        handler = logging.StreamHandler()
+        handler.setFormatter(
+            logging.Formatter(
+                "%(asctime)s %(levelname)-7s [%(name)s] %(message)s",
+                datefmt="%Y-%m-%dT%H:%M:%S",
+            )
+        )
+        logger.addHandler(handler)
+        logger.propagate = False
+        _configured = True
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Return a child logger under the ``text2sql`` namespace.
+
+    Args:
+        name: Optional dotted suffix, e.g. ``"pipeline.linking"``.
+    """
+    if name:
+        return logging.getLogger(f"{_LIBRARY_LOGGER_NAME}.{name}")
+    return logging.getLogger(_LIBRARY_LOGGER_NAME)

text2sql/pipeline/__init__.py

@@ -0,0 +1,9 @@
+"""Pipeline stages and the orchestrator."""
+
+from __future__ import annotations
+
+from .generation import GenerationStage
+from .linking import LinkingStage
+from .pipeline import Text2SQL
+
+__all__ = ["Text2SQL", "LinkingStage", "GenerationStage"]

text2sql/pipeline/generation.py

@@ -0,0 +1,110 @@
+"""Stage 2 — SQL generation.
+
+Takes the reduced schema from Stage 1 and produces a SQL string, plus an
+optional short explanation, in the configured dialect.
+"""
+
+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 LinkedSchema, LLMResponse, StageMetadata, _Timer
+from .json_utils import extract_json, with_retry
+
+logger = get_logger("pipeline.generation")
+
+GENERATION_SCHEMA_NAME = "sql_generation"
+
+GENERATION_JSON_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "additionalProperties": False,
+    "properties": {
+        "sql": {"type": "string"},
+        "explanation": {"type": "string"},
+    },
+    "required": ["sql", "explanation"],
+}
+
+
+class GenerationStage:
+    """Runs SQL generation (Stage 2)."""
+
+    def __init__(
+        self,
+        provider: LLMProvider,
+        settings: Settings,
+        prompts: PromptTemplates,
+    ) -> None:
+        self._provider = provider
+        self._settings = settings
+        self._prompts = prompts
+
+    def run(
+        self,
+        request: str,
+        linked: LinkedSchema,
+        full_schema: Schema,
+    ) -> tuple[str, str, StageMetadata]:
+        """Generate SQL for ``request`` using the linked subset.
+
+        Returns ``(sql, explanation, metadata)``.
+        """
+        stage_cfg = self._settings.generation
+
+        # Render the reduced schema from the real metadata (types, keys, FKs),
+        # limited to the linked tables. Richer than just table names.
+        reduced_schema = full_schema.subset(linked.tables)
+        reduced_text = schema_to_prompt(reduced_schema) if reduced_schema.tables else "(none)"
+
+        template = self._prompts.load(stage_cfg.prompt_template)
+        system, user = template.render(
+            linked_schema=reduced_text,
+            request=request,
+            dialect=self._settings.general.sql_dialect,
+            entities=", ".join(linked.entities) or "(none)",
+            filters="; ".join(linked.filters) or "(none)",
+        )
+
+        attempts = {"n": 0}
+
+        def attempt() -> tuple[str, str, LLMResponse]:
+            attempts["n"] += 1
+            response = self._provider.complete_json(
+                system=system,
+                user=user,
+                json_schema=GENERATION_JSON_SCHEMA,
+                schema_name=GENERATION_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)
+            sql = str(data.get("sql", "")).strip()
+            explanation = str(data.get("explanation", "")).strip()
+            if not sql:
+                from ..errors import StructuredOutputError
+
+                raise StructuredOutputError("Generation returned an empty 'sql' field.")
+            return sql, explanation, response
+
+        with _Timer() as timer:
+            sql, explanation, response = with_retry(
+                attempt,
+                retry=self._settings.retry,
+                description="SQL generation",
+            )
+
+        metadata = StageMetadata(
+            provider=self._provider.name,
+            model=self._provider.model,
+            usage=response.usage,
+            latency_seconds=timer.elapsed,
+            attempts=attempts["n"],
+        )
+        logger.info("Generated SQL (%d chars).", len(sql))
+        return sql, explanation, metadata

text2sql/pipeline/json_utils.py

@@ -0,0 +1,108 @@
+"""JSON extraction and retry helpers. Used by both stages."""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+from typing import Any, Callable, TypeVar
+
+from ..config import RetrySettings
+from ..errors import ProviderError, StructuredOutputError
+from ..logging_utils import get_logger
+
+logger = get_logger("pipeline.json")
+
+_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL)
+
+T = TypeVar("T")
+
+
+def extract_json(text: str, *, strict: bool) -> dict[str, Any]:
+    """Parse a JSON object from raw model text.
+
+    Args:
+        text: The raw response body.
+        strict: If ``True``, only try a direct ``json.loads``. If ``False``,
+            also strip markdown fences and grab the first ``{...}`` block.
+
+    Raises:
+        StructuredOutputError: If no JSON object can be parsed.
+    """
+    candidate = text.strip()
+    try:
+        parsed = json.loads(candidate)
+        if isinstance(parsed, dict):
+            return parsed
+        raise StructuredOutputError("Expected a JSON object, got something else.")
+    except json.JSONDecodeError:
+        if strict:
+            raise StructuredOutputError(
+                "Response was not valid JSON (strict_json is on)."
+            )
+
+    # Lenient path: strip fences, then take the first {...} block.
+    fence = _FENCE_RE.search(candidate)
+    if fence:
+        candidate = fence.group(1).strip()
+        try:
+            parsed = json.loads(candidate)
+            if isinstance(parsed, dict):
+                return parsed
+        except json.JSONDecodeError:
+            pass
+
+    start = candidate.find("{")
+    end = candidate.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        try:
+            parsed = json.loads(candidate[start : end + 1])
+            if isinstance(parsed, dict):
+                return parsed
+        except json.JSONDecodeError:
+            pass
+
+    raise StructuredOutputError("Could not find a JSON object in the response.")
+
+
+def with_retry(
+    func: Callable[[], T],
+    *,
+    retry: RetrySettings,
+    description: str,
+) -> T:
+    """Run ``func`` with retry and backoff.
+
+    Retries on :class:`StructuredOutputError` and transient
+    :class:`ProviderError`. Re-raises the last error when attempts run out.
+
+    Args:
+        func: Zero-arg callable. One attempt.
+        retry: Retry config.
+        description: Label used in log messages.
+    """
+    attempts = max(1, retry.max_attempts)
+    delay = retry.backoff_seconds
+    last_exc: Exception | None = None
+
+    for attempt in range(1, attempts + 1):
+        try:
+            return func()
+        except (StructuredOutputError, ProviderError) as exc:
+            last_exc = exc
+            if attempt >= attempts:
+                break
+            logger.warning(
+                "%s failed on attempt %d/%d (%s). Retrying in %.1fs.",
+                description,
+                attempt,
+                attempts,
+                type(exc).__name__,
+                delay,
+            )
+            if delay > 0:
+                time.sleep(delay)
+            delay *= retry.backoff_factor
+
+    assert last_exc is not None
+    raise last_exc