text2sql-engine 0.1.0__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/LICENSE +1 -1
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/PKG-INFO +12 -5
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/README.md +7 -3
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/pyproject.toml +6 -9
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_config.py +35 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_pipeline.py +42 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_prompts.py +8 -1
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/__init__.py +1 -1
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/config.py +18 -2
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/generation.py +55 -3
- text2sql_engine-0.2.0/text2sql/prompts/generation.txt +42 -0
- text2sql_engine-0.2.0/text2sql/prompts/generation_raw.txt +41 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/PKG-INFO +12 -5
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/SOURCES.txt +1 -0
- text2sql_engine-0.1.0/text2sql/prompts/generation.txt +0 -28
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/setup.cfg +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_schema.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/errors.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/logging_utils.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/__init__.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/json_utils.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/linking.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/pipeline.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/prompts/__init__.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/prompts/linking.txt +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/__init__.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/anthropic_provider.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/base.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/openai_provider.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/registry.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/py.typed +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/__init__.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/file_provider.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/models.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/provider.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/serialize.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/types.py +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/dependency_links.txt +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/requires.txt +0 -0
- {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/top_level.txt +0 -0
|
@@ -1,9 +1,12 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: text2sql-engine
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
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.
|
|
5
|
-
Author:
|
|
5
|
+
Author: Mehmet Kaan Durupunar
|
|
6
6
|
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/defectGI/sqlretrieve
|
|
8
|
+
Project-URL: Repository, https://github.com/defectGI/sqlretrieve
|
|
9
|
+
Project-URL: Issues, https://github.com/defectGI/sqlretrieve/issues
|
|
7
10
|
Keywords: text-to-sql,nl2sql,llm,sql-generation,schema-linking
|
|
8
11
|
Classifier: Development Status :: 4 - Beta
|
|
9
12
|
Classifier: Intended Audience :: Developers
|
|
@@ -67,8 +70,11 @@ user request
|
|
|
67
70
|
reduces the schema to the relevant part and returns structured JSON. Output is
|
|
68
71
|
forced via `response_format` (OpenAI) or tool-calling (Anthropic), with
|
|
69
72
|
strict-JSON parsing and retry as a fallback.
|
|
70
|
-
- **Stage 2 (generation)** takes the request
|
|
71
|
-
|
|
73
|
+
- **Stage 2 (generation)** takes the request, that reduced subset, and Stage 1's
|
|
74
|
+
reasoning (rationale, chosen joins, column reasons). It follows that logic and
|
|
75
|
+
writes SQL in the target dialect. Two output modes (config `output_mode`):
|
|
76
|
+
`structured` forces JSON with `sql` + `explanation`; `raw` asks for plain SQL
|
|
77
|
+
text, for specialized text-to-SQL models.
|
|
72
78
|
- The two stages can use **different providers and models**. For example a cheap
|
|
73
79
|
general model for linking, a stronger SQL model for generation.
|
|
74
80
|
|
|
@@ -165,8 +171,9 @@ No magic numbers in code. Every tunable is here.
|
|
|
165
171
|
| `temperature` | float | `0.0` | Sampling temperature. |
|
|
166
172
|
| `top_p` | float | `1.0` | Nucleus-sampling cutoff. |
|
|
167
173
|
| `max_tokens` | int | `1024` | Max tokens for the SQL response. |
|
|
168
|
-
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. |
|
|
174
|
+
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. In `raw` mode, defaults to `generation_raw.txt` when unset. |
|
|
169
175
|
| `timeout_seconds` | float | `60` | Per-request timeout. |
|
|
176
|
+
| `output_mode` | string | `structured` | `structured`: force JSON with `sql` + `explanation`, reliable parsing, best for general models. `raw`: ask for plain SQL text (no explanation), better for specialized text-to-SQL models trained to emit SQL only. |
|
|
170
177
|
|
|
171
178
|
#### `[retry]` — shared by both stages
|
|
172
179
|
|
|
@@ -33,8 +33,11 @@ user request
|
|
|
33
33
|
reduces the schema to the relevant part and returns structured JSON. Output is
|
|
34
34
|
forced via `response_format` (OpenAI) or tool-calling (Anthropic), with
|
|
35
35
|
strict-JSON parsing and retry as a fallback.
|
|
36
|
-
- **Stage 2 (generation)** takes the request
|
|
37
|
-
|
|
36
|
+
- **Stage 2 (generation)** takes the request, that reduced subset, and Stage 1's
|
|
37
|
+
reasoning (rationale, chosen joins, column reasons). It follows that logic and
|
|
38
|
+
writes SQL in the target dialect. Two output modes (config `output_mode`):
|
|
39
|
+
`structured` forces JSON with `sql` + `explanation`; `raw` asks for plain SQL
|
|
40
|
+
text, for specialized text-to-SQL models.
|
|
38
41
|
- The two stages can use **different providers and models**. For example a cheap
|
|
39
42
|
general model for linking, a stronger SQL model for generation.
|
|
40
43
|
|
|
@@ -131,8 +134,9 @@ No magic numbers in code. Every tunable is here.
|
|
|
131
134
|
| `temperature` | float | `0.0` | Sampling temperature. |
|
|
132
135
|
| `top_p` | float | `1.0` | Nucleus-sampling cutoff. |
|
|
133
136
|
| `max_tokens` | int | `1024` | Max tokens for the SQL response. |
|
|
134
|
-
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. |
|
|
137
|
+
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. In `raw` mode, defaults to `generation_raw.txt` when unset. |
|
|
135
138
|
| `timeout_seconds` | float | `60` | Per-request timeout. |
|
|
139
|
+
| `output_mode` | string | `structured` | `structured`: force JSON with `sql` + `explanation`, reliable parsing, best for general models. `raw`: ask for plain SQL text (no explanation), better for specialized text-to-SQL models trained to emit SQL only. |
|
|
136
140
|
|
|
137
141
|
#### `[retry]` — shared by both stages
|
|
138
142
|
|
|
@@ -6,12 +6,12 @@ build-backend = "setuptools.build_meta"
|
|
|
6
6
|
# Distribution name on PyPI. The IMPORT name stays `text2sql`:
|
|
7
7
|
# pip install text2sql-engine -> import text2sql
|
|
8
8
|
name = "text2sql-engine"
|
|
9
|
-
version = "0.
|
|
9
|
+
version = "0.2.0"
|
|
10
10
|
description = "Schema-agnostic, config-driven Text-to-SQL library (2-stage LLM pipeline). Produces a SQL string; never connects to or executes against a database."
|
|
11
11
|
readme = "README.md"
|
|
12
12
|
requires-python = ">=3.11"
|
|
13
13
|
license = { text = "MIT" }
|
|
14
|
-
authors = [{ name = "
|
|
14
|
+
authors = [{ name = "Mehmet Kaan Durupunar" }]
|
|
15
15
|
keywords = ["text-to-sql", "nl2sql", "llm", "sql-generation", "schema-linking"]
|
|
16
16
|
classifiers = [
|
|
17
17
|
"Development Status :: 4 - Beta",
|
|
@@ -43,13 +43,10 @@ all = ["openai>=1.0", "anthropic>=0.34"]
|
|
|
43
43
|
# Test toolchain.
|
|
44
44
|
dev = ["pytest>=7.4"]
|
|
45
45
|
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
# Homepage = "https://github.com/OWNER/text2sql-engine"
|
|
51
|
-
# Repository = "https://github.com/OWNER/text2sql-engine"
|
|
52
|
-
# Issues = "https://github.com/OWNER/text2sql-engine/issues"
|
|
46
|
+
[project.urls]
|
|
47
|
+
Homepage = "https://github.com/defectGI/sqlretrieve"
|
|
48
|
+
Repository = "https://github.com/defectGI/sqlretrieve"
|
|
49
|
+
Issues = "https://github.com/defectGI/sqlretrieve/issues"
|
|
53
50
|
|
|
54
51
|
[tool.setuptools.packages.find]
|
|
55
52
|
include = ["text2sql*"]
|
|
@@ -78,6 +78,41 @@ def test_missing_required_env_raises(tmp_path, monkeypatch):
|
|
|
78
78
|
assert "SCHEMA_PATH" in str(exc.value)
|
|
79
79
|
|
|
80
80
|
|
|
81
|
+
def _set_min_env(monkeypatch):
|
|
82
|
+
monkeypatch.setenv("SCHEMA_PATH", "examples/schema.yaml")
|
|
83
|
+
monkeypatch.setenv("PROMPT_DIR", "text2sql/prompts")
|
|
84
|
+
monkeypatch.setenv("LINKING_PROVIDER", "openai")
|
|
85
|
+
monkeypatch.setenv("LINKING_MODEL", "gpt-4o-mini")
|
|
86
|
+
monkeypatch.setenv("SQL_PROVIDER", "openai")
|
|
87
|
+
monkeypatch.setenv("SQL_MODEL", "gpt-4o")
|
|
88
|
+
monkeypatch.setenv("OPENAI_API_KEY", "sk-openai")
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def test_output_mode_defaults_to_structured(tmp_path, monkeypatch):
|
|
92
|
+
_set_min_env(monkeypatch)
|
|
93
|
+
cfg = _write_config(tmp_path) # no output_mode set
|
|
94
|
+
settings = load_settings(config_path=cfg)
|
|
95
|
+
assert settings.generation.output_mode == "structured"
|
|
96
|
+
assert settings.generation.prompt_template == "generation.txt"
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_raw_mode_defaults_to_raw_template(tmp_path, monkeypatch):
|
|
100
|
+
_set_min_env(monkeypatch)
|
|
101
|
+
cfg = tmp_path / "config.toml"
|
|
102
|
+
cfg.write_text('[generation]\noutput_mode = "raw"\n', encoding="utf-8")
|
|
103
|
+
settings = load_settings(config_path=cfg)
|
|
104
|
+
assert settings.generation.output_mode == "raw"
|
|
105
|
+
assert settings.generation.prompt_template == "generation_raw.txt"
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def test_invalid_output_mode_raises(tmp_path, monkeypatch):
|
|
109
|
+
_set_min_env(monkeypatch)
|
|
110
|
+
cfg = tmp_path / "config.toml"
|
|
111
|
+
cfg.write_text('[generation]\noutput_mode = "banana"\n', encoding="utf-8")
|
|
112
|
+
with pytest.raises(ConfigError):
|
|
113
|
+
load_settings(config_path=cfg)
|
|
114
|
+
|
|
115
|
+
|
|
81
116
|
def test_missing_config_toml_raises(tmp_path, monkeypatch):
|
|
82
117
|
monkeypatch.setenv("SCHEMA_PATH", "examples/schema.yaml")
|
|
83
118
|
monkeypatch.setenv("PROMPT_DIR", "text2sql/prompts")
|
|
@@ -59,6 +59,21 @@ def test_stage2_receives_only_linked_tables(
|
|
|
59
59
|
assert "Table: customers" not in stage2_user_prompt
|
|
60
60
|
|
|
61
61
|
|
|
62
|
+
def test_stage2_prompt_carries_linking_reasoning(
|
|
63
|
+
settings, schema_provider, prompts, linked_payload, generation_payload
|
|
64
|
+
):
|
|
65
|
+
# Stage 1's rationale and joins must reach the Stage 2 prompt so the SQL
|
|
66
|
+
# model follows the linking logic instead of guessing.
|
|
67
|
+
linking = FakeProvider([linked_payload])
|
|
68
|
+
generation = FakeProvider([generation_payload])
|
|
69
|
+
engine = _engine(settings, schema_provider, prompts, linking, generation)
|
|
70
|
+
engine.run("top products")
|
|
71
|
+
|
|
72
|
+
stage2_user_prompt = generation.calls[0]["user"]
|
|
73
|
+
assert "rank by units sold" in stage2_user_prompt # from the rationale
|
|
74
|
+
assert "order_items.product_id = products.id" in stage2_user_prompt # a join
|
|
75
|
+
|
|
76
|
+
|
|
62
77
|
def test_hallucinated_names_are_dropped(
|
|
63
78
|
settings, schema_provider, prompts, generation_payload
|
|
64
79
|
):
|
|
@@ -140,3 +155,30 @@ def test_empty_request_rejected(
|
|
|
140
155
|
)
|
|
141
156
|
with pytest.raises(ValueError):
|
|
142
157
|
engine.run(" ")
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
def test_raw_mode_returns_plain_sql(settings, schema_provider, prompts, linked_payload):
|
|
161
|
+
# Raw mode: Stage 2 returns plain SQL, no explanation.
|
|
162
|
+
settings.generation.output_mode = "raw"
|
|
163
|
+
settings.generation.prompt_template = "generation_raw.txt"
|
|
164
|
+
linking = FakeProvider([linked_payload])
|
|
165
|
+
generation = FakeProvider(["SELECT p.name FROM products p;"])
|
|
166
|
+
engine = _engine(settings, schema_provider, prompts, linking, generation)
|
|
167
|
+
|
|
168
|
+
result = engine.run("list products")
|
|
169
|
+
assert result.sql == "SELECT p.name FROM products p;"
|
|
170
|
+
assert result.explanation == ""
|
|
171
|
+
# complete() was used, not complete_json() -> no schema_name recorded.
|
|
172
|
+
assert "schema_name" not in generation.calls[0]
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def test_raw_mode_strips_markdown_fences(settings, schema_provider, prompts, linked_payload):
|
|
176
|
+
settings.generation.output_mode = "raw"
|
|
177
|
+
settings.generation.prompt_template = "generation_raw.txt"
|
|
178
|
+
fenced = "```sql\nSELECT 1;\n```"
|
|
179
|
+
linking = FakeProvider([linked_payload])
|
|
180
|
+
generation = FakeProvider([fenced])
|
|
181
|
+
engine = _engine(settings, schema_provider, prompts, linking, generation)
|
|
182
|
+
|
|
183
|
+
result = engine.run("anything")
|
|
184
|
+
assert result.sql == "SELECT 1;"
|
|
@@ -29,7 +29,14 @@ def test_packaged_fallback_when_no_directory():
|
|
|
29
29
|
def test_render_substitutes_placeholders():
|
|
30
30
|
tmpl = PromptTemplates().load("generation.txt")
|
|
31
31
|
system, user = tmpl.render(
|
|
32
|
-
linked_schema="X",
|
|
32
|
+
linked_schema="X",
|
|
33
|
+
request="Y",
|
|
34
|
+
dialect="postgres",
|
|
35
|
+
entities="a",
|
|
36
|
+
filters="b",
|
|
37
|
+
columns="c",
|
|
38
|
+
joins="d",
|
|
39
|
+
rationale="e",
|
|
33
40
|
)
|
|
34
41
|
assert "postgres" in system
|
|
35
42
|
assert "Y" in user
|
|
@@ -52,6 +52,9 @@ class StageSettings:
|
|
|
52
52
|
timeout_seconds: float = 60.0
|
|
53
53
|
# Used by linking only. Generation ignores it.
|
|
54
54
|
max_tables: int = 8
|
|
55
|
+
# Used by generation only. "structured" forces JSON (sql + explanation).
|
|
56
|
+
# "raw" asks for plain SQL text. Linking ignores it.
|
|
57
|
+
output_mode: str = "structured"
|
|
55
58
|
|
|
56
59
|
|
|
57
60
|
@dataclass
|
|
@@ -149,6 +152,7 @@ def _stage_from_toml(
|
|
|
149
152
|
if include_max_tables
|
|
150
153
|
else defaults.max_tables
|
|
151
154
|
),
|
|
155
|
+
output_mode=str(table.get("output_mode", defaults.output_mode)).lower(),
|
|
152
156
|
)
|
|
153
157
|
|
|
154
158
|
|
|
@@ -226,9 +230,21 @@ def load_settings(
|
|
|
226
230
|
default_template="linking.txt",
|
|
227
231
|
include_max_tables=True,
|
|
228
232
|
)
|
|
233
|
+
|
|
234
|
+
gen_tbl = toml_data.get("generation", {})
|
|
235
|
+
gen_output_mode = str(gen_tbl.get("output_mode", "structured")).lower()
|
|
236
|
+
if gen_output_mode not in ("structured", "raw"):
|
|
237
|
+
raise ConfigError(
|
|
238
|
+
f"[generation] output_mode must be 'structured' or 'raw', "
|
|
239
|
+
f"got '{gen_output_mode}'."
|
|
240
|
+
)
|
|
241
|
+
# In raw mode, default to the raw template unless the user set one.
|
|
242
|
+
gen_default_template = "generation.txt"
|
|
243
|
+
if gen_output_mode == "raw" and "prompt_template" not in gen_tbl:
|
|
244
|
+
gen_default_template = "generation_raw.txt"
|
|
229
245
|
generation = _stage_from_toml(
|
|
230
|
-
|
|
231
|
-
default_template=
|
|
246
|
+
gen_tbl,
|
|
247
|
+
default_template=gen_default_template,
|
|
232
248
|
include_max_tables=False,
|
|
233
249
|
)
|
|
234
250
|
|
|
@@ -2,13 +2,22 @@
|
|
|
2
2
|
|
|
3
3
|
Takes the reduced schema from Stage 1 and produces a SQL string, plus an
|
|
4
4
|
optional short explanation, in the configured dialect.
|
|
5
|
+
|
|
6
|
+
Two output modes (config: [generation] output_mode):
|
|
7
|
+
|
|
8
|
+
* "structured" (default) — force JSON with "sql" and "explanation" fields.
|
|
9
|
+
Reliable parsing. Best for general chat/instruct models.
|
|
10
|
+
* "raw" — ask for plain SQL text. No explanation. Better for specialized
|
|
11
|
+
text-to-SQL models that are trained to emit SQL only.
|
|
5
12
|
"""
|
|
6
13
|
|
|
7
14
|
from __future__ import annotations
|
|
8
15
|
|
|
16
|
+
import re
|
|
9
17
|
from typing import Any
|
|
10
18
|
|
|
11
19
|
from ..config import Settings
|
|
20
|
+
from ..errors import StructuredOutputError
|
|
12
21
|
from ..logging_utils import get_logger
|
|
13
22
|
from ..prompts import PromptTemplates
|
|
14
23
|
from ..providers.base import LLMProvider
|
|
@@ -16,6 +25,18 @@ from ..schema import Schema, schema_to_prompt
|
|
|
16
25
|
from ..types import LinkedSchema, LLMResponse, StageMetadata, _Timer
|
|
17
26
|
from .json_utils import extract_json, with_retry
|
|
18
27
|
|
|
28
|
+
# Matches a ```sql ... ``` or ``` ... ``` fenced block.
|
|
29
|
+
_FENCE_RE = re.compile(r"```(?:sql)?\s*(.*?)```", re.DOTALL | re.IGNORECASE)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _strip_sql(text: str) -> str:
|
|
33
|
+
"""Return plain SQL from a raw response. Drops markdown fences if present."""
|
|
34
|
+
stripped = text.strip()
|
|
35
|
+
fence = _FENCE_RE.search(stripped)
|
|
36
|
+
if fence:
|
|
37
|
+
return fence.group(1).strip()
|
|
38
|
+
return stripped
|
|
39
|
+
|
|
19
40
|
logger = get_logger("pipeline.generation")
|
|
20
41
|
|
|
21
42
|
GENERATION_SCHEMA_NAME = "sql_generation"
|
|
@@ -61,6 +82,19 @@ class GenerationStage:
|
|
|
61
82
|
reduced_schema = full_schema.subset(linked.tables)
|
|
62
83
|
reduced_text = schema_to_prompt(reduced_schema) if reduced_schema.tables else "(none)"
|
|
63
84
|
|
|
85
|
+
# Forward Stage 1's reasoning to Stage 2. The linking model already
|
|
86
|
+
# worked out the logic (ranges, joins, which column means what); the SQL
|
|
87
|
+
# model should follow it instead of guessing.
|
|
88
|
+
columns_text = "\n".join(
|
|
89
|
+
f"- {c.table}.{c.column}" + (f" — {c.reason}" if c.reason else "")
|
|
90
|
+
for c in linked.columns
|
|
91
|
+
) or "(none)"
|
|
92
|
+
joins_text = "\n".join(
|
|
93
|
+
f"- {j.left_table}.{j.left_column} = {j.right_table}.{j.right_column}"
|
|
94
|
+
for j in linked.joins
|
|
95
|
+
) or "(none)"
|
|
96
|
+
rationale_text = linked.rationale.strip() or "(none)"
|
|
97
|
+
|
|
64
98
|
template = self._prompts.load(stage_cfg.prompt_template)
|
|
65
99
|
system, user = template.render(
|
|
66
100
|
linked_schema=reduced_text,
|
|
@@ -68,11 +102,14 @@ class GenerationStage:
|
|
|
68
102
|
dialect=self._settings.general.sql_dialect,
|
|
69
103
|
entities=", ".join(linked.entities) or "(none)",
|
|
70
104
|
filters="; ".join(linked.filters) or "(none)",
|
|
105
|
+
columns=columns_text,
|
|
106
|
+
joins=joins_text,
|
|
107
|
+
rationale=rationale_text,
|
|
71
108
|
)
|
|
72
109
|
|
|
73
110
|
attempts = {"n": 0}
|
|
74
111
|
|
|
75
|
-
def
|
|
112
|
+
def attempt_structured() -> tuple[str, str, LLMResponse]:
|
|
76
113
|
attempts["n"] += 1
|
|
77
114
|
response = self._provider.complete_json(
|
|
78
115
|
system=system,
|
|
@@ -87,11 +124,26 @@ class GenerationStage:
|
|
|
87
124
|
sql = str(data.get("sql", "")).strip()
|
|
88
125
|
explanation = str(data.get("explanation", "")).strip()
|
|
89
126
|
if not sql:
|
|
90
|
-
from ..errors import StructuredOutputError
|
|
91
|
-
|
|
92
127
|
raise StructuredOutputError("Generation returned an empty 'sql' field.")
|
|
93
128
|
return sql, explanation, response
|
|
94
129
|
|
|
130
|
+
def attempt_raw() -> tuple[str, str, LLMResponse]:
|
|
131
|
+
attempts["n"] += 1
|
|
132
|
+
response = self._provider.complete(
|
|
133
|
+
system=system,
|
|
134
|
+
user=user,
|
|
135
|
+
temperature=stage_cfg.temperature,
|
|
136
|
+
top_p=stage_cfg.top_p,
|
|
137
|
+
max_tokens=stage_cfg.max_tokens,
|
|
138
|
+
)
|
|
139
|
+
sql = _strip_sql(response.text)
|
|
140
|
+
if not sql:
|
|
141
|
+
raise StructuredOutputError("Generation returned empty SQL.")
|
|
142
|
+
# Raw mode does not produce an explanation.
|
|
143
|
+
return sql, "", response
|
|
144
|
+
|
|
145
|
+
attempt = attempt_raw if stage_cfg.output_mode == "raw" else attempt_structured
|
|
146
|
+
|
|
95
147
|
with _Timer() as timer:
|
|
96
148
|
sql, explanation, response = with_retry(
|
|
97
149
|
attempt,
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
SYSTEM:
|
|
2
|
+
You are an expert SQL generator. You get a user request, a reduced schema (only
|
|
3
|
+
the relevant tables/columns/joins picked by an upstream linking step), and that
|
|
4
|
+
step's analysis. Produce one correct SQL query in the target dialect.
|
|
5
|
+
|
|
6
|
+
Rules:
|
|
7
|
+
- Use only the tables and columns in the reduced schema.
|
|
8
|
+
- Follow the linking analysis. It already worked out the logic; do not ignore it.
|
|
9
|
+
- Respect the intended comparison semantics from the analysis and filter hints.
|
|
10
|
+
If a value should be matched by a numeric range or an exact compare, implement
|
|
11
|
+
it that way. Do not fall back to substring/text (LIKE) matching unless that is
|
|
12
|
+
clearly what is meant.
|
|
13
|
+
- Use the suggested joins to connect tables. Prefer explicit JOINs.
|
|
14
|
+
- Write idiomatic SQL for the target dialect: {dialect}.
|
|
15
|
+
- Do not run the query. Only produce it.
|
|
16
|
+
- Return the answer strictly as the requested JSON object with two fields:
|
|
17
|
+
"sql" (the query) and "explanation" (one or two sentences).
|
|
18
|
+
|
|
19
|
+
USER:
|
|
20
|
+
Target SQL dialect: {dialect}
|
|
21
|
+
|
|
22
|
+
Reduced schema (relevant subset):
|
|
23
|
+
--------------------------------
|
|
24
|
+
{linked_schema}
|
|
25
|
+
--------------------------------
|
|
26
|
+
|
|
27
|
+
Linking analysis (reasoning from the previous step — follow this logic):
|
|
28
|
+
{rationale}
|
|
29
|
+
|
|
30
|
+
Relevant columns:
|
|
31
|
+
{columns}
|
|
32
|
+
|
|
33
|
+
Suggested joins:
|
|
34
|
+
{joins}
|
|
35
|
+
|
|
36
|
+
Detected entities: {entities}
|
|
37
|
+
Filter hints: {filters}
|
|
38
|
+
|
|
39
|
+
User request:
|
|
40
|
+
"{request}"
|
|
41
|
+
|
|
42
|
+
Write the SQL query that answers the request.
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
SYSTEM:
|
|
2
|
+
You are an expert SQL generator. You get a user request, a reduced schema (only
|
|
3
|
+
the relevant tables/columns/joins picked by an upstream linking step), and that
|
|
4
|
+
step's analysis. Produce one correct SQL query in the target dialect.
|
|
5
|
+
|
|
6
|
+
Rules:
|
|
7
|
+
- Use only the tables and columns in the reduced schema.
|
|
8
|
+
- Follow the linking analysis. It already worked out the logic; do not ignore it.
|
|
9
|
+
- Respect the intended comparison semantics from the analysis and filter hints.
|
|
10
|
+
If a value should be matched by a numeric range or an exact compare, implement
|
|
11
|
+
it that way. Do not fall back to substring/text (LIKE) matching unless that is
|
|
12
|
+
clearly what is meant.
|
|
13
|
+
- Use the suggested joins to connect tables. Prefer explicit JOINs.
|
|
14
|
+
- Write idiomatic SQL for the target dialect: {dialect}.
|
|
15
|
+
- Do not run the query. Only produce it.
|
|
16
|
+
- Output ONLY the SQL query. No explanation, no comment, no markdown fences.
|
|
17
|
+
|
|
18
|
+
USER:
|
|
19
|
+
Target SQL dialect: {dialect}
|
|
20
|
+
|
|
21
|
+
Reduced schema (relevant subset):
|
|
22
|
+
--------------------------------
|
|
23
|
+
{linked_schema}
|
|
24
|
+
--------------------------------
|
|
25
|
+
|
|
26
|
+
Linking analysis (reasoning from the previous step — follow this logic):
|
|
27
|
+
{rationale}
|
|
28
|
+
|
|
29
|
+
Relevant columns:
|
|
30
|
+
{columns}
|
|
31
|
+
|
|
32
|
+
Suggested joins:
|
|
33
|
+
{joins}
|
|
34
|
+
|
|
35
|
+
Detected entities: {entities}
|
|
36
|
+
Filter hints: {filters}
|
|
37
|
+
|
|
38
|
+
User request:
|
|
39
|
+
"{request}"
|
|
40
|
+
|
|
41
|
+
Write the SQL query that answers the request.
|
|
@@ -1,9 +1,12 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: text2sql-engine
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
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.
|
|
5
|
-
Author:
|
|
5
|
+
Author: Mehmet Kaan Durupunar
|
|
6
6
|
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/defectGI/sqlretrieve
|
|
8
|
+
Project-URL: Repository, https://github.com/defectGI/sqlretrieve
|
|
9
|
+
Project-URL: Issues, https://github.com/defectGI/sqlretrieve/issues
|
|
7
10
|
Keywords: text-to-sql,nl2sql,llm,sql-generation,schema-linking
|
|
8
11
|
Classifier: Development Status :: 4 - Beta
|
|
9
12
|
Classifier: Intended Audience :: Developers
|
|
@@ -67,8 +70,11 @@ user request
|
|
|
67
70
|
reduces the schema to the relevant part and returns structured JSON. Output is
|
|
68
71
|
forced via `response_format` (OpenAI) or tool-calling (Anthropic), with
|
|
69
72
|
strict-JSON parsing and retry as a fallback.
|
|
70
|
-
- **Stage 2 (generation)** takes the request
|
|
71
|
-
|
|
73
|
+
- **Stage 2 (generation)** takes the request, that reduced subset, and Stage 1's
|
|
74
|
+
reasoning (rationale, chosen joins, column reasons). It follows that logic and
|
|
75
|
+
writes SQL in the target dialect. Two output modes (config `output_mode`):
|
|
76
|
+
`structured` forces JSON with `sql` + `explanation`; `raw` asks for plain SQL
|
|
77
|
+
text, for specialized text-to-SQL models.
|
|
72
78
|
- The two stages can use **different providers and models**. For example a cheap
|
|
73
79
|
general model for linking, a stronger SQL model for generation.
|
|
74
80
|
|
|
@@ -165,8 +171,9 @@ No magic numbers in code. Every tunable is here.
|
|
|
165
171
|
| `temperature` | float | `0.0` | Sampling temperature. |
|
|
166
172
|
| `top_p` | float | `1.0` | Nucleus-sampling cutoff. |
|
|
167
173
|
| `max_tokens` | int | `1024` | Max tokens for the SQL response. |
|
|
168
|
-
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. |
|
|
174
|
+
| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. In `raw` mode, defaults to `generation_raw.txt` when unset. |
|
|
169
175
|
| `timeout_seconds` | float | `60` | Per-request timeout. |
|
|
176
|
+
| `output_mode` | string | `structured` | `structured`: force JSON with `sql` + `explanation`, reliable parsing, best for general models. `raw`: ask for plain SQL text (no explanation), better for specialized text-to-SQL models trained to emit SQL only. |
|
|
170
177
|
|
|
171
178
|
#### `[retry]` — shared by both stages
|
|
172
179
|
|
|
@@ -18,6 +18,7 @@ text2sql/pipeline/linking.py
|
|
|
18
18
|
text2sql/pipeline/pipeline.py
|
|
19
19
|
text2sql/prompts/__init__.py
|
|
20
20
|
text2sql/prompts/generation.txt
|
|
21
|
+
text2sql/prompts/generation_raw.txt
|
|
21
22
|
text2sql/prompts/linking.txt
|
|
22
23
|
text2sql/providers/__init__.py
|
|
23
24
|
text2sql/providers/anthropic_provider.py
|
|
@@ -1,28 +0,0 @@
|
|
|
1
|
-
SYSTEM:
|
|
2
|
-
You are an expert SQL generator. You get a user request and a reduced schema
|
|
3
|
-
(only the relevant tables/columns/joins picked by an upstream linking step).
|
|
4
|
-
Produce one correct SQL query in the target dialect.
|
|
5
|
-
|
|
6
|
-
Rules:
|
|
7
|
-
- Use only the tables and columns in the reduced schema.
|
|
8
|
-
- Write idiomatic SQL for the target dialect: {dialect}.
|
|
9
|
-
- Prefer explicit JOINs using the given relations.
|
|
10
|
-
- Do not run the query. Only produce it.
|
|
11
|
-
- Return the answer strictly as the requested JSON object with two fields:
|
|
12
|
-
"sql" (the query) and "explanation" (one or two sentences).
|
|
13
|
-
|
|
14
|
-
USER:
|
|
15
|
-
Target SQL dialect: {dialect}
|
|
16
|
-
|
|
17
|
-
Reduced schema (relevant subset):
|
|
18
|
-
--------------------------------
|
|
19
|
-
{linked_schema}
|
|
20
|
-
--------------------------------
|
|
21
|
-
|
|
22
|
-
Detected entities: {entities}
|
|
23
|
-
Filter hints: {filters}
|
|
24
|
-
|
|
25
|
-
User request:
|
|
26
|
-
"{request}"
|
|
27
|
-
|
|
28
|
-
Write the SQL query that answers the request.
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/dependency_links.txt
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|