prazy1208/text2sql

1# Few-Shot Example Storage — Implementation Pack2 3This document contains **ready-to-add** artifacts for the Few-Shot Agent pipeline.  4**Plan mode** prevented writing `.sql` / `.py` files directly; copy sections into new files as named, or switch to **Agent mode** and ask to apply this pack.5 6**Storage (v1):** Postgres is the source of truth; a **build script** writes **`metadata_store/few_shot_examples_metadata.json`** (same idea as [`metadata_store/relationships_*_metadata.json`](metadata_store/) — list of objects for agents). Objects are **`{ id, question_text, sql_query, query_type }`** with **no** `embedding` field in v1 (text-only catalog).7 8**Retrieval:** **`fewshot_retrieval.py`** loads the catalog from **that JSON file** (optional fallback: read from DB if the file is missing). **`few_shot_agent.py`** uses an **LLM** to pick the **top `k`** examples (default 2). No FAISS.9 10---11 12## 1. SQL — `scripts/create_system_schema_few_shot_examples.sql`13 14```sql15-- Few-shot SQL pattern examples (generic) for the Few-Shot Agent.16-- Run against text2sql_db (or: python scripts/run_create_few_shot_examples_schema.py)17 18CREATE SCHEMA IF NOT EXISTS system_schema;19 20CREATE TABLE IF NOT EXISTS system_schema.few_shot_examples (21    id              SERIAL PRIMARY KEY,22    question_text   TEXT NOT NULL,23    sql_query       TEXT NOT NULL,24    query_type      TEXT NOT NULL25);26 27CREATE INDEX IF NOT EXISTS idx_few_shot_examples_query_type28    ON system_schema.few_shot_examples (query_type);29 30COMMENT ON SCHEMA system_schema IS 'Cross-cutting metadata (few-shot examples, etc.)';31COMMENT ON TABLE system_schema.few_shot_examples IS 'Curated question/SQL pairs for few-shot retrieval';32COMMENT ON COLUMN system_schema.few_shot_examples.query_type IS 'Pattern label (e.g. aggregation_groupby, join)';33```34 35**Migration (idempotent):** `scripts/migration_add_system_schema_few_shot_examples.sql` — same `CREATE TABLE IF NOT EXISTS` block without comments.36 37---38 39## 2. Schema runner — `scripts/run_create_few_shot_examples_schema.py`40 41```python42"""Apply scripts/create_system_schema_few_shot_examples.sql. Run from project root."""43 44import os45import sys46from pathlib import Path47 48PROJECT_ROOT = Path(__file__).resolve().parent.parent49if str(PROJECT_ROOT) not in sys.path:50    sys.path.insert(0, str(PROJECT_ROOT))51 52os.chdir(PROJECT_ROOT)53 54from dotenv import load_dotenv55from sqlalchemy import create_engine56 57load_dotenv()58 59 60def get_engine():61    database_url = os.getenv("DATABASE_URL")62    if database_url:63        return create_engine(database_url)64    host = os.getenv("DB_HOST", "localhost")65    port = os.getenv("DB_PORT", "5432")66    user = os.getenv("DB_USER", "postgres")67    password = os.getenv("DB_PASSWORD", "")68    dbname = os.getenv("DB_NAME", "text2sql_db")69    url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"70    return create_engine(url)71 72 73def main():74    sql_file = PROJECT_ROOT / "scripts" / "create_system_schema_few_shot_examples.sql"75    if not sql_file.exists():76        print(f"SQL file not found: {sql_file}")77        sys.exit(1)78    sql = sql_file.read_text(encoding="utf-8")79    engine = get_engine()80    raw_conn = engine.raw_connection()81    try:82        cur = raw_conn.cursor()83        cur.execute(sql)84        raw_conn.commit()85    finally:86        raw_conn.close()87    print("Done. system_schema.few_shot_examples created.")88 89 90if __name__ == "__main__":91    main()92```93 94---95 96## 3. Insert script — `scripts/insert_few_shot_examples.py`97 98- Deletes existing rows (optional) then inserts 20 curated rows.99- Run after the table exists: `python scripts/insert_few_shot_examples.py`100 101```python102"""Insert 20 generic few-shot examples into system_schema.few_shot_examples."""103 104import os105import sys106from pathlib import Path107 108PROJECT_ROOT = Path(__file__).resolve().parent.parent109if str(PROJECT_ROOT) not in sys.path:110    sys.path.insert(0, str(PROJECT_ROOT))111 112os.chdir(PROJECT_ROOT)113 114from dotenv import load_dotenv115from sqlalchemy import create_engine, text116 117load_dotenv()118 119EXAMPLES = [120    {"question": "Retrieve specific columns from a dataset", "sql": "SELECT column1, column2 FROM table_name;", "query_type": "select"},121    {"question": "Filter records based on a specific condition", "sql": "SELECT * FROM table_name WHERE column = value;", "query_type": "filter"},122    {"question": "Filter records using multiple conditions", "sql": "SELECT * FROM table_name WHERE column1 = value1 AND column2 = value2;", "query_type": "filter_multiple"},123    {"question": "Retrieve records where a column value falls within a range", "sql": "SELECT * FROM table_name WHERE column BETWEEN value1 AND value2;", "query_type": "filter_range"},124    {"question": "Calculate the total value of a numeric column", "sql": "SELECT SUM(numeric_column) FROM table_name;", "query_type": "aggregation_sum"},125    {"question": "Count the number of records", "sql": "SELECT COUNT(*) FROM table_name;", "query_type": "aggregation_count"},126    {"question": "Calculate average value of a numeric column", "sql": "SELECT AVG(numeric_column) FROM table_name;", "query_type": "aggregation_avg"},127    {"question": "Group data by a category and calculate total values", "sql": "SELECT category, SUM(numeric_column) FROM table_name GROUP BY category;", "query_type": "aggregation_groupby"},128    {"question": "Filter grouped results based on aggregated values", "sql": "SELECT category, SUM(numeric_column) FROM table_name GROUP BY category HAVING SUM(numeric_column) > threshold;", "query_type": "having"},129    {"question": "Sort records in descending order based on a metric", "sql": "SELECT * FROM table_name ORDER BY metric DESC;", "query_type": "ordering"},130    {"question": "Retrieve top N records based on a metric", "sql": "SELECT column_name FROM table_name ORDER BY metric DESC LIMIT 5;", "query_type": "ranking_limit"},131    {"question": "Retrieve distinct values from a column", "sql": "SELECT DISTINCT column_name FROM table_name;", "query_type": "distinct"},132    {"question": "Join two related datasets using a common identifier", "sql": "SELECT a.column1, b.column2 FROM table_a a JOIN table_b b ON a.id = b.id;", "query_type": "join"},133    {"question": "Combine multiple related datasets through chained joins", "sql": "SELECT a.column1, b.column2, c.column3 FROM table_a a JOIN table_b b ON a.id = b.id JOIN table_c c ON b.id = c.id;", "query_type": "multi_join"},134    {"question": "Filter records based on a date range", "sql": "SELECT * FROM table_name WHERE date_column >= 'start_date' AND date_column <= 'end_date';", "query_type": "time_filter"},135    {"question": "Compare aggregated values across two time periods", "sql": "SELECT SUM(CASE WHEN date_column >= period1_start AND date_column <= period1_end THEN value END) AS period1_total, SUM(CASE WHEN date_column >= period2_start AND date_column <= period2_end THEN value END) AS period2_total FROM table_name;", "query_type": "time_comparison"},136    {"question": "Create conditional labels based on column values", "sql": "SELECT column_name, CASE WHEN condition THEN 'Category A' ELSE 'Category B' END FROM table_name;", "query_type": "case_when"},137    {"question": "Use a subquery to filter results", "sql": "SELECT * FROM table_name WHERE column IN (SELECT column FROM another_table);", "query_type": "subquery"},138    {"question": "Rank records within groups based on a metric", "sql": "SELECT column_name, ROW_NUMBER() OVER (PARTITION BY group_column ORDER BY metric DESC) FROM table_name;", "query_type": "window_function"},139    {"question": "Handle missing values in a column", "sql": "SELECT COALESCE(column_name, default_value) FROM table_name;", "query_type": "null_handling"},140]141 142 143def get_engine():144    database_url = os.getenv("DATABASE_URL")145    if database_url:146        return create_engine(database_url)147    host = os.getenv("DB_HOST", "localhost")148    port = os.getenv("DB_PORT", "5432")149    user = os.getenv("DB_USER", "postgres")150    password = os.getenv("DB_PASSWORD", "")151    dbname = os.getenv("DB_NAME", "text2sql_db")152    url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"153    return create_engine(url)154 155 156def main():157    engine = get_engine()158    insert_sql = text("""159        INSERT INTO system_schema.few_shot_examples (question_text, sql_query, query_type)160        VALUES (:question_text, :sql_query, :query_type)161    """)162    with engine.begin() as conn:163        conn.execute(text("DELETE FROM system_schema.few_shot_examples"))164        for ex in EXAMPLES:165            conn.execute(166                insert_sql,167                {168                    "question_text": ex["question"],169                    "sql_query": ex["sql"],170                    "query_type": ex["query_type"],171                },172            )173    print(f"Inserted {len(EXAMPLES)} row(s) into system_schema.few_shot_examples.")174 175 176if __name__ == "__main__":177    main()178```179 180---181 182## 4. Export to metadata store — `build_few_shot_metadata_store.py` (project root)183 184- Read all rows: `SELECT id, question_text, sql_query, query_type FROM system_schema.few_shot_examples ORDER BY id`.185- Write **`metadata_store/few_shot_examples_metadata.json`** as a JSON array (one object per row).186- Register **`FEWSHOT_METADATA_NAME`** in `backend/config.py` (same pattern as `RELATIONSHIP_METADATA_NAMES` / `METADATA_STORE_DIR`).187- **Run after** insert whenever the table changes.188 189---190 191## 5. Catalog load — `backend/services/fewshot_retrieval.py`192 193**No vector search.** Single function:194 195- `list_all_few_shot_examples() -> list[dict]` — **`json.load`** from `METADATA_STORE_DIR / FEWSHOT_METADATA_NAME`. Optional: if file missing, `SELECT ...` from Postgres and log once.196- Returns `[]` if no data.197 198---199 200## 6. Few-Shot Agent — `backend/agents/few_shot_agent.py`201 202- Call `list_all_few_shot_examples()`; if empty, return `{ "few_shot_examples": [] }`.203- Build a **prompt** (`FEW_SHOT_AGENT_PROMPT`) listing candidates with **`question`** and **`query_type` only** (no SQL in the LLM prompt; Gen-SQL receives SQL from resolved rows).204- Include **rephrased_question** and **keywords** in the user message.205- Ask the LLM (`chat_completion`) to return **JSON** `{"selected_examples": [{ "question", "query_type" }, ...]}` — **all** examples it considers relevant (no fixed cap).206- Map each pair back to catalog rows by matching `question_text` and `query_type`; return full `{ "few_shot_examples": [ {...}, ... ] }` with `id`, `question_text`, `sql_query`, `query_type` for downstream Gen-SQL.207 208---209 210## 7. Run order211 2121. `python scripts/run_create_few_shot_examples_schema.py` (or run SQL in pgAdmin)2132. `python scripts/insert_few_shot_examples.py`2143. `python build_few_shot_metadata_store.py` — generates `metadata_store/few_shot_examples_metadata.json`2154. From API or Gen-SQL: call `run_few_shot_agent(rephrased, keywords, business_insights)` (retrieval reads JSON)216 217---218 219## 8. Switch to Agent mode220 221To have these files **created automatically** in the repo, switch to **Agent mode** and ask: *“Implement the Few-Shot pipeline from docs/FEWSHOT_PIPELINE_IMPLEMENTATION.md”*.222