text2sql-engine 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

text2sql/__init__.py

@@ -37,7 +37,7 @@ from .types import (
     TokenUsage,
 )
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     "Text2SQL",

text2sql/config.py

@@ -52,6 +52,9 @@ class StageSettings:
     timeout_seconds: float = 60.0
     # Used by linking only. Generation ignores it.
     max_tables: int = 8
+    # Used by generation only. "structured" forces JSON (sql + explanation).
+    # "raw" asks for plain SQL text. Linking ignores it.
+    output_mode: str = "structured"
 
 
 @dataclass
@@ -149,6 +152,7 @@ def _stage_from_toml(
             if include_max_tables
             else defaults.max_tables
         ),
+        output_mode=str(table.get("output_mode", defaults.output_mode)).lower(),
     )
 
 
@@ -226,9 +230,21 @@ def load_settings(
         default_template="linking.txt",
         include_max_tables=True,
     )
+
+    gen_tbl = toml_data.get("generation", {})
+    gen_output_mode = str(gen_tbl.get("output_mode", "structured")).lower()
+    if gen_output_mode not in ("structured", "raw"):
+        raise ConfigError(
+            f"[generation] output_mode must be 'structured' or 'raw', "
+            f"got '{gen_output_mode}'."
+        )
+    # In raw mode, default to the raw template unless the user set one.
+    gen_default_template = "generation.txt"
+    if gen_output_mode == "raw" and "prompt_template" not in gen_tbl:
+        gen_default_template = "generation_raw.txt"
     generation = _stage_from_toml(
-        toml_data.get("generation", {}),
-        default_template="generation.txt",
+        gen_tbl,
+        default_template=gen_default_template,
         include_max_tables=False,
     )
 

text2sql/pipeline/generation.py

@@ -2,13 +2,22 @@
 
 Takes the reduced schema from Stage 1 and produces a SQL string, plus an
 optional short explanation, in the configured dialect.
+
+Two output modes (config: [generation] output_mode):
+
+* "structured" (default) — force JSON with "sql" and "explanation" fields.
+  Reliable parsing. Best for general chat/instruct models.
+* "raw" — ask for plain SQL text. No explanation. Better for specialized
+  text-to-SQL models that are trained to emit SQL only.
 """
 
 from __future__ import annotations
 
+import re
 from typing import Any
 
 from ..config import Settings
+from ..errors import StructuredOutputError
 from ..logging_utils import get_logger
 from ..prompts import PromptTemplates
 from ..providers.base import LLMProvider
@@ -16,6 +25,18 @@ from ..schema import Schema, schema_to_prompt
 from ..types import LinkedSchema, LLMResponse, StageMetadata, _Timer
 from .json_utils import extract_json, with_retry
 
+# Matches a ```sql ... ``` or ``` ... ``` fenced block.
+_FENCE_RE = re.compile(r"```(?:sql)?\s*(.*?)```", re.DOTALL | re.IGNORECASE)
+
+
+def _strip_sql(text: str) -> str:
+    """Return plain SQL from a raw response. Drops markdown fences if present."""
+    stripped = text.strip()
+    fence = _FENCE_RE.search(stripped)
+    if fence:
+        return fence.group(1).strip()
+    return stripped
+
 logger = get_logger("pipeline.generation")
 
 GENERATION_SCHEMA_NAME = "sql_generation"
@@ -61,6 +82,19 @@ class GenerationStage:
         reduced_schema = full_schema.subset(linked.tables)
         reduced_text = schema_to_prompt(reduced_schema) if reduced_schema.tables else "(none)"
 
+        # Forward Stage 1's reasoning to Stage 2. The linking model already
+        # worked out the logic (ranges, joins, which column means what); the SQL
+        # model should follow it instead of guessing.
+        columns_text = "\n".join(
+            f"- {c.table}.{c.column}" + (f" — {c.reason}" if c.reason else "")
+            for c in linked.columns
+        ) or "(none)"
+        joins_text = "\n".join(
+            f"- {j.left_table}.{j.left_column} = {j.right_table}.{j.right_column}"
+            for j in linked.joins
+        ) or "(none)"
+        rationale_text = linked.rationale.strip() or "(none)"
+
         template = self._prompts.load(stage_cfg.prompt_template)
         system, user = template.render(
             linked_schema=reduced_text,
@@ -68,11 +102,14 @@ class GenerationStage:
             dialect=self._settings.general.sql_dialect,
             entities=", ".join(linked.entities) or "(none)",
             filters="; ".join(linked.filters) or "(none)",
+            columns=columns_text,
+            joins=joins_text,
+            rationale=rationale_text,
         )
 
         attempts = {"n": 0}
 
-        def attempt() -> tuple[str, str, LLMResponse]:
+        def attempt_structured() -> tuple[str, str, LLMResponse]:
             attempts["n"] += 1
             response = self._provider.complete_json(
                 system=system,
@@ -87,11 +124,26 @@ class GenerationStage:
             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
 
+        def attempt_raw() -> tuple[str, str, LLMResponse]:
+            attempts["n"] += 1
+            response = self._provider.complete(
+                system=system,
+                user=user,
+                temperature=stage_cfg.temperature,
+                top_p=stage_cfg.top_p,
+                max_tokens=stage_cfg.max_tokens,
+            )
+            sql = _strip_sql(response.text)
+            if not sql:
+                raise StructuredOutputError("Generation returned empty SQL.")
+            # Raw mode does not produce an explanation.
+            return sql, "", response
+
+        attempt = attempt_raw if stage_cfg.output_mode == "raw" else attempt_structured
+
         with _Timer() as timer:
             sql, explanation, response = with_retry(
                 attempt,

text2sql/prompts/generation.txt

@@ -1,12 +1,17 @@
 SYSTEM:
-You are an expert SQL generator. You get a user request and a reduced schema
-(only the relevant tables/columns/joins picked by an upstream linking step).
-Produce one correct SQL query in the target dialect.
+You are an expert SQL generator. You get a user request, a reduced schema (only
+the relevant tables/columns/joins picked by an upstream linking step), and that
+step's analysis. Produce one correct SQL query in the target dialect.
 
 Rules:
 - Use only the tables and columns in the reduced schema.
+- Follow the linking analysis. It already worked out the logic; do not ignore it.
+- Respect the intended comparison semantics from the analysis and filter hints.
+  If a value should be matched by a numeric range or an exact compare, implement
+  it that way. Do not fall back to substring/text (LIKE) matching unless that is
+  clearly what is meant.
+- Use the suggested joins to connect tables. Prefer explicit JOINs.
 - Write idiomatic SQL for the target dialect: {dialect}.
-- Prefer explicit JOINs using the given relations.
 - Do not run the query. Only produce it.
 - Return the answer strictly as the requested JSON object with two fields:
   "sql" (the query) and "explanation" (one or two sentences).
@@ -19,6 +24,15 @@ Reduced schema (relevant subset):
 {linked_schema}
 --------------------------------
 
+Linking analysis (reasoning from the previous step — follow this logic):
+{rationale}
+
+Relevant columns:
+{columns}
+
+Suggested joins:
+{joins}
+
 Detected entities: {entities}
 Filter hints: {filters}
 

text2sql/prompts/generation_raw.txt

@@ -0,0 +1,41 @@
+SYSTEM:
+You are an expert SQL generator. You get a user request, a reduced schema (only
+the relevant tables/columns/joins picked by an upstream linking step), and that
+step's analysis. Produce one correct SQL query in the target dialect.
+
+Rules:
+- Use only the tables and columns in the reduced schema.
+- Follow the linking analysis. It already worked out the logic; do not ignore it.
+- Respect the intended comparison semantics from the analysis and filter hints.
+  If a value should be matched by a numeric range or an exact compare, implement
+  it that way. Do not fall back to substring/text (LIKE) matching unless that is
+  clearly what is meant.
+- Use the suggested joins to connect tables. Prefer explicit JOINs.
+- Write idiomatic SQL for the target dialect: {dialect}.
+- Do not run the query. Only produce it.
+- Output ONLY the SQL query. No explanation, no comment, no markdown fences.
+
+USER:
+Target SQL dialect: {dialect}
+
+Reduced schema (relevant subset):
+--------------------------------
+{linked_schema}
+--------------------------------
+
+Linking analysis (reasoning from the previous step — follow this logic):
+{rationale}
+
+Relevant columns:
+{columns}
+
+Suggested joins:
+{joins}
+
+Detected entities: {entities}
+Filter hints: {filters}
+
+User request:
+"{request}"
+
+Write the SQL query that answers the request.

{text2sql_engine-0.1.0.dist-info → text2sql_engine-0.2.0.dist-info}/METADATA

@@ -1,9 +1,12 @@
 Metadata-Version: 2.4
 Name: text2sql-engine
-Version: 0.1.0
+Version: 0.2.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
+Author: Mehmet Kaan Durupunar
 License: MIT
+Project-URL: Homepage, https://github.com/defectGI/sqlretrieve
+Project-URL: Repository, https://github.com/defectGI/sqlretrieve
+Project-URL: Issues, https://github.com/defectGI/sqlretrieve/issues
 Keywords: text-to-sql,nl2sql,llm,sql-generation,schema-linking
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -67,8 +70,11 @@ user request
   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.
+- **Stage 2 (generation)** takes the request, that reduced subset, and Stage 1's
+  reasoning (rationale, chosen joins, column reasons). It follows that logic and
+  writes SQL in the target dialect. Two output modes (config `output_mode`):
+  `structured` forces JSON with `sql` + `explanation`; `raw` asks for plain SQL
+  text, for specialized text-to-SQL models.
 - The two stages can use **different providers and models**. For example a cheap
   general model for linking, a stronger SQL model for generation.
 
@@ -165,8 +171,9 @@ No magic numbers in code. Every tunable is here.
 | `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`. |
+| `prompt_template` | string | `generation.txt` | Template file, looked up in `PROMPT_DIR`. In `raw` mode, defaults to `generation_raw.txt` when unset. |
 | `timeout_seconds` | float | `60` | Per-request timeout. |
+| `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. |
 
 #### `[retry]` — shared by both stages
 

{text2sql_engine-0.1.0.dist-info → text2sql_engine-0.2.0.dist-info}/RECORD

@@ -1,16 +1,17 @@
-text2sql/__init__.py,sha256=EUP3cVW8vLecps3wqmPqqr6Apaht2tD9QsfdjP5-CQM,1442
-text2sql/config.py,sha256=jCpdUogpzRmkL5kN6CPqYXhUNhRA7EdaCFlCY65QnLA,7709
+text2sql/__init__.py,sha256=dwQ98mfnAHRbUhKmOEnC0BcGK-8Slxosi8hqDgpkheM,1442
+text2sql/config.py,sha256=6snv99rf8jf8NtD3xetSrvIaltbLyq2eoISoMoPadSw,8503
 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/generation.py,sha256=M9wehTcMAq1QoK1nnQ4SJd1mszSfXs9E06l26IXpKNw,5750
 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/generation.txt,sha256=th97BZsHFlRtH30OqHj4DY_BooCsa7Zad9EoXC9M_cM,1381
+text2sql/prompts/generation_raw.txt,sha256=EsKLl7WrNWQzcoxAFW9l6yucUtBa7sIimCsvIjRx7ZI,1321
 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
@@ -22,8 +23,8 @@ text2sql/schema/file_provider.py,sha256=Nv9x2-Eh9mXISk5XvUQaLfagth2s3LbX5F7jWp-o
 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.2.0.dist-info/licenses/LICENSE,sha256=qfdfVnpb5NOM3-SK13PnmWagjCLy5sE9ORlgOGXc8hU,1078
+text2sql_engine-0.2.0.dist-info/METADATA,sha256=Xn9bMUbUlTrirl1JWpZ99bOrfZRZzK2IdDx8bQVmRFw,13061
+text2sql_engine-0.2.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+text2sql_engine-0.2.0.dist-info/top_level.txt,sha256=ZONWO1r8wsWpyHhF5qPzAjnsNoCEhwZ2Z-M_V8nKALc,9
+text2sql_engine-0.2.0.dist-info/RECORD,,

{text2sql_engine-0.1.0.dist-info → text2sql_engine-0.2.0.dist-info}/licenses/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2026 text2sql contributors
+Copyright (c) 2026 Mehmet Kaan Durupunar
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal