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.
Files changed (40) hide show
  1. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/LICENSE +1 -1
  2. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/PKG-INFO +12 -5
  3. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/README.md +7 -3
  4. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/pyproject.toml +6 -9
  5. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_config.py +35 -0
  6. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_pipeline.py +42 -0
  7. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_prompts.py +8 -1
  8. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/__init__.py +1 -1
  9. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/config.py +18 -2
  10. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/generation.py +55 -3
  11. text2sql_engine-0.2.0/text2sql/prompts/generation.txt +42 -0
  12. text2sql_engine-0.2.0/text2sql/prompts/generation_raw.txt +41 -0
  13. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/PKG-INFO +12 -5
  14. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/SOURCES.txt +1 -0
  15. text2sql_engine-0.1.0/text2sql/prompts/generation.txt +0 -28
  16. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/setup.cfg +0 -0
  17. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/tests/test_schema.py +0 -0
  18. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/errors.py +0 -0
  19. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/logging_utils.py +0 -0
  20. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/__init__.py +0 -0
  21. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/json_utils.py +0 -0
  22. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/linking.py +0 -0
  23. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/pipeline/pipeline.py +0 -0
  24. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/prompts/__init__.py +0 -0
  25. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/prompts/linking.txt +0 -0
  26. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/__init__.py +0 -0
  27. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/anthropic_provider.py +0 -0
  28. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/base.py +0 -0
  29. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/openai_provider.py +0 -0
  30. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/providers/registry.py +0 -0
  31. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/py.typed +0 -0
  32. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/__init__.py +0 -0
  33. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/file_provider.py +0 -0
  34. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/models.py +0 -0
  35. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/provider.py +0 -0
  36. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/schema/serialize.py +0 -0
  37. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql/types.py +0 -0
  38. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/dependency_links.txt +0 -0
  39. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/requires.txt +0 -0
  40. {text2sql_engine-0.1.0 → text2sql_engine-0.2.0}/text2sql_engine.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  MIT License
2
2
 
3
- Copyright (c) 2026 text2sql contributors
3
+ Copyright (c) 2026 Mehmet Kaan Durupunar
4
4
 
5
5
  Permission is hereby granted, free of charge, to any person obtaining a copy
6
6
  of this software and associated documentation files (the "Software"), to deal
@@ -1,9 +1,12 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: text2sql-engine
3
- Version: 0.1.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: stajyer14
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 and that reduced subset. It writes
71
- SQL in the target dialect, plus a short explanation.
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 and that reduced subset. It writes
37
- SQL in the target dialect, plus a short explanation.
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.1.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 = "stajyer14" }]
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
- # Replace OWNER with your GitHub username/org once the repo is pushed, then
47
- # uncomment. Dead links look worse than no links, so this stays commented out
48
- # until the repo exists.
49
- # [project.urls]
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", request="Y", dialect="postgres", entities="a", filters="b"
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
@@ -37,7 +37,7 @@ from .types import (
37
37
  TokenUsage,
38
38
  )
39
39
 
40
- __version__ = "0.1.0"
40
+ __version__ = "0.2.0"
41
41
 
42
42
  __all__ = [
43
43
  "Text2SQL",
@@ -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
- toml_data.get("generation", {}),
231
- default_template="generation.txt",
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 attempt() -> tuple[str, str, LLMResponse]:
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.1.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: stajyer14
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 and that reduced subset. It writes
71
- SQL in the target dialect, plus a short explanation.
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.