text2sql-engine 0.1.0__py3-none-any.whl

text2sql/types.py

@@ -0,0 +1,138 @@
+"""Dataclasses returned by the pipeline.
+
+Plain dataclasses, no external deps. Easy to inspect and serialize.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class TokenUsage:
+    """Token counts for one LLM call. Zero when the provider reports nothing."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+    def __add__(self, other: "TokenUsage") -> "TokenUsage":
+        return TokenUsage(
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+            total_tokens=self.total_tokens + other.total_tokens,
+        )
+
+
+@dataclass
+class LLMResponse:
+    """One provider call, normalized.
+
+    Same shape for every provider, so the pipeline never touches an SDK type.
+    ``raw`` keeps the native response for advanced callers.
+    """
+
+    text: str
+    usage: TokenUsage = field(default_factory=TokenUsage)
+    model: str = ""
+    raw: Any = None
+
+
+@dataclass
+class LinkedColumn:
+    """A column picked by Stage 1."""
+
+    table: str
+    column: str
+    reason: str = ""
+
+
+@dataclass
+class LinkedJoin:
+    """A candidate join between two tables, found during linking."""
+
+    left_table: str
+    left_column: str
+    right_table: str
+    right_column: str
+
+
+@dataclass
+class LinkedSchema:
+    """Stage 1 output: the reduced schema.
+
+    This is what Stage 2 receives. It is also returned on the final result so
+    callers can see what the pipeline picked as relevant.
+    """
+
+    tables: list[str] = field(default_factory=list)
+    columns: list[LinkedColumn] = field(default_factory=list)
+    joins: list[LinkedJoin] = field(default_factory=list)
+    entities: list[str] = field(default_factory=list)
+    filters: list[str] = field(default_factory=list)
+    rationale: str = ""
+
+    def is_empty(self) -> bool:
+        """True when linking found no usable table or column."""
+        return not self.tables and not self.columns
+
+
+@dataclass
+class StageMetadata:
+    """Metadata for one stage: which model, how many tokens, how long."""
+
+    provider: str
+    model: str
+    usage: TokenUsage = field(default_factory=TokenUsage)
+    latency_seconds: float = 0.0
+    attempts: int = 1
+
+
+@dataclass
+class RunMetadata:
+    """Metadata for a full run (both stages)."""
+
+    linking: StageMetadata
+    generation: StageMetadata
+    total_latency_seconds: float = 0.0
+
+    @property
+    def total_usage(self) -> TokenUsage:
+        """Token usage of both stages combined."""
+        return self.linking.usage + self.generation.usage
+
+
+@dataclass
+class Text2SQLResult:
+    """Final result of :meth:`text2sql.Text2SQL.run`.
+
+    Attributes:
+        sql: The generated SQL string. This is the deliverable.
+        explanation: Optional short explanation from Stage 2.
+        linked_schema: Stage 1 output (the reduced schema).
+        metadata: Models, tokens, and latency for the run.
+        request: The original request.
+    """
+
+    sql: str
+    linked_schema: LinkedSchema
+    metadata: RunMetadata
+    request: str
+    explanation: str = ""
+
+
+class _Timer:
+    """Small monotonic timer for stage latency. Internal only."""
+
+    def __init__(self) -> None:
+        self._start = 0.0
+        self.elapsed = 0.0
+
+    def __enter__(self) -> "_Timer":
+        self._start = time.perf_counter()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.elapsed = time.perf_counter() - self._start

text2sql_engine-0.1.0.dist-info/METADATA

@@ -0,0 +1,304 @@
+Metadata-Version: 2.4
+Name: text2sql-engine
+Version: 0.1.0
+Summary: Schema-agnostic, config-driven Text-to-SQL library (2-stage LLM pipeline). Produces a SQL string; never connects to or executes against a database.
+Author: stajyer14
+License: MIT
+Keywords: text-to-sql,nl2sql,llm,sql-generation,schema-linking
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Database
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: python-dotenv>=1.0
+Requires-Dist: PyYAML>=6.0
+Provides-Extra: openai
+Requires-Dist: openai>=1.0; extra == "openai"
+Provides-Extra: anthropic
+Requires-Dist: anthropic>=0.34; extra == "anthropic"
+Provides-Extra: all
+Requires-Dist: openai>=1.0; extra == "all"
+Requires-Dist: anthropic>=0.34; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Dynamic: license-file
+
+# text2sql
+
+A **schema-agnostic, config-driven Text-to-SQL library**. Give it a
+natural-language request and metadata for any relational schema, and it returns
+a SQL string.
+
+It is a **library**, not an app. No UI, no web server, no CLI. It **never
+connects to or runs against a database**. Generating the SQL is the last step.
+Running it is the caller's job.
+
+## How it works — a 2-stage LLM pipeline
+
+```
+user request
+      │
+      ▼
+┌─────────────────────────────┐   full schema metadata
+│ Stage 1: Schema Linking LLM │◄──────────────────────
+│  reduce schema → subset     │
+└─────────────────────────────┘
+      │  linked schema (tables, columns, joins, entities, filters, rationale)
+      ▼
+┌─────────────────────────────┐
+│ Stage 2: SQL Generation LLM │
+│  produce SQL in dialect     │
+└─────────────────────────────┘
+      │
+      ▼
+  Text2SQLResult(sql, linked_schema, metadata, explanation)
+```
+
+- **Stage 1 (linking)** takes the request and the *full* schema metadata. It
+  reduces the schema to the relevant part and returns structured JSON. Output is
+  forced via `response_format` (OpenAI) or tool-calling (Anthropic), with
+  strict-JSON parsing and retry as a fallback.
+- **Stage 2 (generation)** takes the request and that reduced subset. It writes
+  SQL in the target dialect, plus a short explanation.
+- The two stages can use **different providers and models**. For example a cheap
+  general model for linking, a stronger SQL model for generation.
+
+## Installation
+
+```bash
+pip install text2sql-engine            # from PyPI (import name: text2sql)
+pip install "text2sql-engine[all]"     # + openai and anthropic SDKs
+```
+
+The PyPI name is `text2sql-engine`. The import name is `text2sql`
+(`import text2sql`). Different names on purpose, since `text2sql` was taken.
+
+From a checkout of this repo:
+
+```bash
+pip install -e .              # core library
+pip install -e ".[all]"       # + openai and anthropic SDKs
+pip install -e ".[dev]"       # + pytest
+```
+
+The provider SDKs (`openai`, `anthropic`) are imported lazily. So the library
+imports, and the tests run, without them installed.
+
+## Quick start
+
+```python
+from text2sql import Text2SQL
+
+engine = Text2SQL.from_config()          # reads .env + config.toml
+result = engine.run("get me the top 5 best-selling products last month")
+
+print(result.sql)              # the generated SQL string
+print(result.linked_schema)    # Stage 1 output (reduced schema)
+print(result.metadata)         # models, tokens, latency, attempts
+```
+
+1. Copy `.env.example` to `.env`. Fill in provider selection and API keys.
+2. Adjust `config.toml` for behavior (dialect, sampling, retries).
+3. Point `SCHEMA_PATH` at your metadata file. See `examples/schema.yaml`.
+
+## Configuration — two separate layers
+
+Config is split by concern. **Secrets, paths, and provider/model *selection* go
+in `.env`. Everything tunable about *behavior* goes in `config.toml`.** One typed
+layer (`text2sql.config.Settings`) loads both. It fails fast with a clear error
+if a required `.env` variable is missing.
+
+### `.env` — secrets, paths, provider & model selection
+
+| Variable | Required | Description |
+|----------|----------|-------------|
+| `SCHEMA_PATH` | yes | Path to the schema file (`.json`, `.yaml`, `.yml`). |
+| `PROMPT_DIR` | no | Directory with prompt templates (`linking.txt`, `generation.txt`). If unset, the templates **packaged inside the library** are used. So it works after pip install, from any working directory. |
+| `CONFIG_PATH` | no | Override the `config.toml` location. Defaults to `./config.toml`. |
+| `LINKING_PROVIDER` | yes | Provider for Stage 1. One of the registered names: `openai`, `anthropic`. |
+| `LINKING_MODEL` | yes | Model id for Stage 1 (schema linking). |
+| `LINKING_BASE_URL` | no | Endpoint override for Stage 1 (for OpenAI-compatible gateways). |
+| `SQL_PROVIDER` | yes | Provider for Stage 2. |
+| `SQL_MODEL` | yes | Model id for Stage 2 (SQL generation). |
+| `SQL_BASE_URL` | no | Endpoint override for Stage 2. |
+| `OPENAI_API_KEY` | if used | API key. Needed only if a stage uses provider `openai`. |
+| `ANTHROPIC_API_KEY` | if used | API key. Needed only if a stage uses provider `anthropic`. |
+
+Provider, model, and `base_url` are picked **per stage**, independently.
+
+### `config.toml` — behavior
+
+No magic numbers in code. Every tunable is here.
+
+#### `[general]`
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `sql_dialect` | string | `postgres` | Target dialect (e.g. `postgres`, `mysql`, `sqlite`). Passed to the prompts. |
+| `log_level` | string | `INFO` | `DEBUG` / `INFO` / `WARNING` / `ERROR`. |
+| `strict_json` | bool | `true` | `true`: parse LLM JSON strictly. `false`: tolerate fences / extra prose. |
+
+#### `[linking]` — Stage 1
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `temperature` | float | `0.0` | Sampling temperature. |
+| `top_p` | float | `1.0` | Nucleus-sampling cutoff. |
+| `max_tokens` | int | `1024` | Max tokens for the linking response. |
+| `prompt_template` | string | `linking.txt` | Template file, looked up in `PROMPT_DIR`. |
+| `max_tables` | int | `8` | Max tables linking may return. |
+| `timeout_seconds` | float | `60` | Per-request timeout. |
+
+#### `[generation]` — Stage 2
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `temperature` | float | `0.0` | Sampling temperature. |
+| `top_p` | float | `1.0` | Nucleus-sampling cutoff. |
+| `max_tokens` | int | `1024` | Max tokens for the SQL response. |
+| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. |
+| `timeout_seconds` | float | `60` | Per-request timeout. |
+
+#### `[retry]` — shared by both stages
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `max_attempts` | int | `3` | Attempts per LLM call before failing. Retries on JSON-parse and transient provider errors. |
+| `backoff_seconds` | float | `1.0` | First delay between retries. |
+| `backoff_factor` | float | `2.0` | Delay is multiplied by this each retry. |
+
+## Schema metadata format
+
+The schema is **never hardcoded**. It is loaded from an external file through a
+`SchemaProvider`. `examples/schema.yaml` (and the same `examples/schema.json`)
+show the full format: a small e-commerce database with `customers`,
+`categories`, `products`, `orders`, and `order_items`. Each table has a
+description. Each column has a name, type, description, primary-key flag, and
+sample values. Foreign keys are listed. Both JSON and YAML work, picked by
+extension.
+
+The example runs out of the box. The tests exercise the pipeline against it with
+mocked LLM calls.
+
+## Extending
+
+Everything is swappable via **config + dependency injection**:
+
+```python
+from text2sql import Text2SQL, SchemaProvider, PromptTemplates
+from text2sql.schema import Schema
+
+# 1. Custom schema source, e.g. live DB introspection instead of a file.
+class MyIntrospectionProvider(SchemaProvider):
+    def load(self) -> Schema:
+        ...   # build and return a Schema
+
+# 2. Inject any part. Anything left None is built from config.
+engine = Text2SQL.from_config(
+    schema_provider=MyIntrospectionProvider(),
+    prompts=PromptTemplates("/path/to/my/templates"),
+    # linking_provider=..., generation_provider=...,
+)
+```
+
+### Adding a new LLM provider
+
+Subclass `LLMProvider`, implement `complete` (and optionally `complete_json` for
+native structured output), and register it. One class, one decorator:
+
+```python
+from text2sql import LLMProvider, register_provider
+from text2sql.types import LLMResponse
+
+@register_provider("myprovider")
+class MyProvider(LLMProvider):
+    def complete(self, *, system, user, temperature, top_p, max_tokens) -> LLMResponse:
+        ...
+```
+
+Then set `LINKING_PROVIDER=myprovider` (or `SQL_PROVIDER`) in `.env`.
+
+### Prompts
+
+Prompt templates are plain, editable text files in `PROMPT_DIR`. Each has a
+`SYSTEM:` and a `USER:` section with `{placeholder}` substitution. Edit them
+without touching code.
+
+## Result object
+
+`engine.run(...)` returns a `Text2SQLResult`:
+
+| Field | Description |
+|-------|-------------|
+| `sql` | The generated SQL string. The deliverable. |
+| `explanation` | Short explanation from Stage 2. |
+| `linked_schema` | Stage 1 output: `tables`, `columns`, `joins`, `entities`, `filters`, `rationale`. |
+| `metadata` | `RunMetadata`: per-stage provider/model, token usage, latency, attempts. Plus `total_usage` and `total_latency_seconds`. |
+| `request` | The original request. |
+
+## Error handling
+
+Every error subclasses `Text2SQLError`:
+
+- `ConfigError` — missing/invalid `.env` var or `config.toml`.
+- `SchemaError` — schema file missing or invalid.
+- `ProviderError` — provider build or API call failed.
+- `StructuredOutputError` — output not parseable as JSON after all retries.
+- `EmptyLinkingError` — Stage 1 matched no table/column.
+
+Uses stdlib `logging` throughout, never `print`. Secrets are never logged.
+
+## Try it end to end
+
+With a real provider and key set in `.env`:
+
+```bash
+python examples/run_example.py "top 5 best-selling products last month"
+```
+
+Makes real LLM calls for both stages. Prints the linked schema, the SQL, and the
+metadata. Never connects to a database.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest
+```
+
+Tests mock all LLM calls (no network) and run on `examples/schema.yaml`.
+
+## Project layout
+
+```
+text2sql/
+  config.py              .env + config.toml loading, typed settings
+  types.py               result dataclasses
+  errors.py              exceptions
+  logging_utils.py       logging setup
+  schema/                SchemaProvider, file loader, models, serializer
+  providers/             base + registry + openai + anthropic
+  prompts/               editable linking.txt / generation.txt templates
+  pipeline/              linking (Stage 1), generation (Stage 2), orchestrator
+examples/
+  schema.yaml / schema.json    sample metadata "spec"
+  run_example.py               runnable end-to-end example (real LLM call)
+config.toml              behavioral parameters
+.env.example             secrets / paths / model selection template
+tests/                   pipeline, schema, config, and prompt tests
+.github/workflows/ci.yml GitHub Actions: pytest on 3.11 and 3.12
+LICENSE                  MIT
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE).

text2sql_engine-0.1.0.dist-info/RECORD

@@ -0,0 +1,29 @@
+text2sql/__init__.py,sha256=EUP3cVW8vLecps3wqmPqqr6Apaht2tD9QsfdjP5-CQM,1442
+text2sql/config.py,sha256=jCpdUogpzRmkL5kN6CPqYXhUNhRA7EdaCFlCY65QnLA,7709
+text2sql/errors.py,sha256=Y5ORy7g_UsFvDJ2KorB2JFvTyOBqmrPv9MWc8Cl9MI8,1052
+text2sql/logging_utils.py,sha256=hiEH1SDqP16kctUj4uPP0tbXdrCO9AlRcdS-mHYDg5s,1633
+text2sql/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+text2sql/types.py,sha256=MxcdHP7Gnftk1yKdcBQKaVVyeqRnBu5kxbhSZ2xJJ80,3532
+text2sql/pipeline/__init__.py,sha256=rjEOi4U9C-SKP15cyKPdvG_EZBOW0VW01LouS3PJVdY,245
+text2sql/pipeline/generation.py,sha256=LpnsKmSgu4Gfg-nHCCk6akg7ToI-G9Vk0LaVRn9WK5Y,3635
+text2sql/pipeline/json_utils.py,sha256=92viA3xyw_GMF_nKdjSPgAyQ8iBUZgnyNRAoyWC5ZlE,3145
+text2sql/pipeline/linking.py,sha256=U28LOcWJn3cItwAsv2fG4fYv9KcJ80XmP1Fi8nNx4qc,6966
+text2sql/pipeline/pipeline.py,sha256=TE0cHSEJyeY43P89HsQDbnPGSfgRQmtw7KZAq-Yjj1U,5675
+text2sql/prompts/__init__.py,sha256=HfeflfAGsZl3xHnNrmklTP6K_8EKxX6vWJrayq6QuEQ,3315
+text2sql/prompts/generation.txt,sha256=r4125WCr7gsJIzTa5wZTAmbTyn5WvP3QrfGf4CWNP_w,849
+text2sql/prompts/linking.txt,sha256=nLPSkaq3PiWiJ7wXLeD0BvK_-rR3hhc_tLhZ-daPnGY,889
+text2sql/providers/__init__.py,sha256=R5a-YKtyaUFpkVJGWrmhDGyRHeh84gA3F7nEH57dSGo,345
+text2sql/providers/anthropic_provider.py,sha256=Zhh8zM6EfHqkjTgNmhYKdZLHQ69b0Pu2-5WbCW2rD6s,4342
+text2sql/providers/base.py,sha256=lqy_-Z923EfuKvOY247lyifdE5HjN0o52tyi9M2Ufns,3655
+text2sql/providers/openai_provider.py,sha256=EoIkCB-RnFFHzyBMAY2eglzu7WLrVF5i9q8MhcqEyG0,4370
+text2sql/providers/registry.py,sha256=epe7OLm07_jhsspItpauupTy3UEyuUheZ20VYQcq1Ec,1147
+text2sql/schema/__init__.py,sha256=AvXHTwCDiOolg-_btcLReFDWupztDswEsFkz_1EGOmo,395
+text2sql/schema/file_provider.py,sha256=Nv9x2-Eh9mXISk5XvUQaLfagth2s3LbX5F7jWp-oZqw,4918
+text2sql/schema/models.py,sha256=yX5zdKbTuy-g-nkwT0v-0HUTK81jBRoabV-OC-kD-JI,2263
+text2sql/schema/provider.py,sha256=JPMRsjRgEB4RdQR_pFwbP2LwLIXZhQVWSQPoYzaXC7s,798
+text2sql/schema/serialize.py,sha256=_2MeA1IwDJ51-69PXyqXBSFZSyzovbqJkRIMpekkfJA,1833
+text2sql_engine-0.1.0.dist-info/licenses/LICENSE,sha256=dzrVwSmgUgCDYwDn21Job7Yb4huyHjC74evyzsHXas8,1078
+text2sql_engine-0.1.0.dist-info/METADATA,sha256=Cni4NycxDnoS3HgP5eYTRiUXgOmnLhSYBI27LTBm08U,12302
+text2sql_engine-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+text2sql_engine-0.1.0.dist-info/top_level.txt,sha256=ZONWO1r8wsWpyHhF5qPzAjnsNoCEhwZ2Z-M_V8nKALc,9
+text2sql_engine-0.1.0.dist-info/RECORD,,

text2sql_engine-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

text2sql_engine-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 text2sql contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

text2sql_engine-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+text2sql