Text-to-SQL: Help agents turn natural language into SQL queries

DataQuery is the little sister of Google's QueryData product.

Author: amotl

CHANGES.md

@@ -5,6 +5,7 @@
 - Kinesis: Added `ctk kinesis` CLI group with `list-checkpoints` and
   `prune-checkpoints` commands for checkpoint table maintenance
 - Dependencies: Permitted installation of click 8.3
+- QueryData: Help agents turn natural language into SQL queries
 
 ## 2026/03/16 v0.0.46
 - I/O: API improvements: `ctk {load,save} table` became `ctk {load,save}`

cratedb_toolkit/query/cli.py

@@ -1,26 +1,9 @@
-import logging
-
-import click
-from click_aliases import ClickAliasedGroup
-
-from ..util.cli import boot_click
+from ..util.app import make_cli
 from .convert.cli import convert_query
+from .llm.cli import llm_cli
 from .mcp.cli import cli as mcp_cli
 
-logger = logging.getLogger(__name__)
-
-
-@click.group(cls=ClickAliasedGroup)
-@click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
-@click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
-@click.version_option()
-@click.pass_context
-def cli(ctx: click.Context, verbose: bool, debug: bool):
-    """
-    Query utilities.
-    """
-    return boot_click(ctx, verbose, debug)
-
-
+cli = make_cli()
 cli.add_command(convert_query, name="convert")
+cli.add_command(llm_cli, name="llm")
 cli.add_command(mcp_cli, name="mcp")

cratedb_toolkit/query/llm/__init__.py

@@ -0,0 +1,81 @@
+"""
+Use an LLM to query a database in human language via NLSQLTableQueryEngine.
+Example code using LlamaIndex with vanilla Open AI, Azure Open AI, or Ollama.
+"""
+
+import dataclasses
+import logging
+import os
+from typing import Optional
+
+import sqlalchemy as sa
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.base.response.schema import RESPONSE_TYPE
+from llama_index.core.llms import LLM
+from llama_index.core.query_engine import NLSQLTableQueryEngine
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+
+from cratedb_toolkit.query.llm.model import DatabaseInfo, ModelInfo
+from cratedb_toolkit.query.llm.util import configure_llm
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class DataQuery:
+    """
+    DataQuery helps agents turn natural language into SQL queries.
+    It's the little sister of Google's QueryData product. [1]
+
+    We recommend evaluating the Text-to-SQL interface using the Gemma models if you are
+    looking at non-frontier variants that need less resources for inference. However,
+    depending on the complexity of your problem, you may also want to use cutting-edge
+    models with your provider of choice at the cost of higher resource usage.
+
+    Attention: Any natural language SQL table query engine and Text-to-SQL application
+    should be aware that executing arbitrary SQL queries can be a security risk.
+    It is recommended to take precautions as needed, such as using restricted roles,
+    read-only databases, sandboxing, etc.
+
+    [1] https://cloud.google.com/blog/products/databases/introducing-querydata-for-near-100-percent-accurate-data-agents
+    [2] https://github.com/kupp0/multi-db-property-search-data-agents
+    """
+
+    db: DatabaseInfo
+    model: ModelInfo
+    query_engine: Optional[NLSQLTableQueryEngine] = None
+
+    def __post_init__(self):
+        self.setup()
+
+    def setup(self):
+        # Configure database connection and query engine.
+        logger.info("Connecting to CrateDB")
+        engine_crate = sa.create_engine(os.getenv("CRATEDB_SQLALCHEMY_URL", "crate://"))
+        engine_crate.connect()
+
+        # Configure model.
+        logger.info("Configuring LLM model")
+        llm: LLM
+        embed_model: BaseEmbedding
+        llm, embed_model = configure_llm(self.model)
+
+        # Configure query engine.
+        logger.info("Creating query engine")
+        sql_database = SQLDatabase(
+            engine_crate,
+            ignore_tables=self.db.ignore_tables,
+            include_tables=self.db.include_tables,
+        )
+        self.query_engine = NLSQLTableQueryEngine(
+            sql_database=sql_database,
+            llm=llm,
+            embed_model=embed_model,
+        )
+
+    def ask(self, question: str) -> RESPONSE_TYPE:
+        """Invoke an inquiry to the LLM."""
+        if not self.query_engine:
+            raise ValueError("Query engine not configured")
+        logger.debug("Running query: %s", question)
+        return self.query_engine.query(question)

cratedb_toolkit/query/llm/api.py

@@ -0,0 +1,36 @@
+import dataclasses
+from enum import Enum
+from typing import List, Optional
+
+import sqlalchemy as sa
+
+
+class ModelProvider(Enum):
+    """Model provider choices."""
+
+    OPENAI = "openai"
+    AZURE = "azure"
+    OLLAMA = "ollama"
+
+
+@dataclasses.dataclass
+class ModelInfo:
+    """Information about the model."""
+
+    provider: ModelProvider
+    completion: str
+    embedding: str
+    endpoint: Optional[str] = None
+    instance: Optional[str] = None
+    api_key: Optional[str] = None
+    api_version: Optional[str] = None
+
+
+@dataclasses.dataclass
+class DatabaseInfo:
+    """Information about the database."""
+
+    engine: sa.engine.Engine
+    schema: Optional[str] = None
+    ignore_tables: Optional[List[str]] = None
+    include_tables: Optional[List[str]] = None

cratedb_toolkit/query/llm/model.py

@@ -0,0 +1,84 @@
+import os
+from typing import Tuple
+
+import llama_index.core
+from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.llms import LLM
+from llama_index.embeddings.langchain import LangchainEmbedding
+from llama_index.llms.azure_openai import AzureOpenAI
+from llama_index.llms.ollama import Ollama
+from llama_index.llms.openai import OpenAI
+
+from cratedb_toolkit.query.llm.model import ModelInfo, ModelProvider
+
+
+def configure_llm(info: ModelInfo, debug: bool = False) -> Tuple[LLM, BaseEmbedding]:
+    """
+    Configure LLM access and model types. Use either vanilla Open AI, Azure Open AI, or Ollama.
+
+    TODO: What about Hugging Face, Runpod, vLLM, and others?
+
+    Notes about text embedding models:
+
+    > The new model, `text-embedding-ada-002`, replaces five separate models for text search,
+    > text similarity, and code search, and outperforms our previous most capable model,
+    > Davinci, at most tasks, while being priced 99.8% lower.
+
+    - https://openai.com/index/new-and-improved-embedding-model/
+    - https://community.openai.com/t/models-embedding-vs-similarity-vs-search-models/291265
+    """
+
+    completion_model = info.completion
+    embedding_model = info.embedding or "text-embedding-3-large"
+
+    if not info.provider:
+        raise ValueError("LLM model type not defined")
+    if not completion_model:
+        raise ValueError("LLM model name not defined")
+
+    # https://docs.llamaindex.ai/en/stable/understanding/tracing_and_debugging/tracing_and_debugging/
+    if debug:
+        llama_index.core.set_global_handler("simple")
+
+    if info.provider is ModelProvider.OPENAI:
+        llm = OpenAI(
+            model=completion_model,
+            temperature=0.0,
+            api_key=info.api_key,
+            api_version=info.api_version,
+        )
+    elif info.provider is ModelProvider.AZURE:
+        llm = AzureOpenAI(
+            model=completion_model,
+            temperature=0.0,
+            engine=info.instance,
+            azure_endpoint=info.endpoint,
+            api_key=info.api_key,
+            api_version=info.api_version,
+        )
+    elif info.provider is ModelProvider.OLLAMA:
+        # https://docs.llamaindex.ai/en/stable/api_reference/llms/ollama/
+        llm = Ollama(
+            base_url=info.endpoint or "http://localhost:11434",
+            model=completion_model,
+            temperature=0.0,
+            request_timeout=120.0,
+            keep_alive=-1,
+        )
+    else:
+        raise ValueError("LLM model type invalid: %s", info.provider)
+
+    if info.provider is ModelProvider.OPENAI:
+        embed_model = LangchainEmbedding(OpenAIEmbeddings(model=embedding_model))
+    elif info.provider is ModelProvider.AZURE:
+        embed_model = LangchainEmbedding(
+            AzureOpenAIEmbeddings(
+                azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"),
+                model=embedding_model,
+            )
+        )
+    else:
+        embed_model = "local"
+
+    return llm, embed_model  # ty: ignore[invalid-return-type]

cratedb_toolkit/query/llm/util.py

@@ -6,6 +6,7 @@ expressions: Adapters, converters, migration support tasks, etc.
 ```{toctree}
 :maxdepth: 2
 
+llm/index
 mcp/index
 convert
 ```

doc/query/index.md

@@ -0,0 +1,75 @@
+# Text-to-SQL query adapter
+
+## Install
+
+```shell
+uv pip install 'cratedb-toolkit[llm]'
+```
+
+## Usage
+
+### CLI
+
+```shell
+export CRATEDB_CLUSTER_URL=crate://localhost/
+export LLM_PROVIDER=openai
+export OPENAI_API_KEY=<YOUR_OPENAI_API_KEY>
+```
+
+### API
+
+```python
+import sqlalchemy as sa
+from cratedb_toolkit.query.llm.api import DataQuery
+from cratedb_toolkit.query.llm.model import DatabaseInfo, ModelInfo, ModelProvider
+
+engine = sa.create_engine("crate://")
+schema = "doc"
+
+# Use Open AI GPT-4.
+dq = DataQuery(
+    db=DatabaseInfo(engine=engine, schema=schema),
+    model=ModelInfo(provider=ModelProvider.OPENAI, completion="gpt-4.1", embedding="text-embedding-3-large"),
+)
+
+# Use Gemma3 via Ollama.
+dq = DataQuery(
+    db=DatabaseInfo(engine=engine, schema=schema),
+    model=ModelInfo(provider=ModelProvider.OLLAMA, completion="gemma3:1b", embedding="local"),
+)
+```
+
+## Example
+
+```sql
+CREATE TABLE IF NOT EXISTS time_series_data (
+    timestamp TIMESTAMP,
+    value DOUBLE,
+    location STRING,
+    sensor_id INT
+);
+
+INSERT INTO time_series_data (timestamp, value, location, sensor_id)
+VALUES
+    ('2023-09-14T00:00:00', 10.5, 'Sensor A', 1),
+    ('2023-09-14T01:00:00', 15.2, 'Sensor A', 1),
+    ('2023-09-14T02:00:00', 18.9, 'Sensor A', 1),
+    ('2023-09-14T03:00:00', 12.7, 'Sensor B', 2),
+    ('2023-09-14T04:00:00', 17.3, 'Sensor B', 2),
+    ('2023-09-14T05:00:00', 20.1, 'Sensor B', 2),
+    ('2023-09-14T06:00:00', 22.5, 'Sensor A', 1),
+    ('2023-09-14T07:00:00', 18.3, 'Sensor A', 1),
+    ('2023-09-14T08:00:00', 16.8, 'Sensor A', 1),
+    ('2023-09-14T09:00:00', 14.6, 'Sensor B', 2),
+    ('2023-09-14T10:00:00', 13.2, 'Sensor B', 2),
+    ('2023-09-14T11:00:00', 11.7, 'Sensor B', 2);
+
+REFRESH TABLE time_series_data;
+```
+
+```shell
+ctk query llm "What is the average value for sensor 1?"
+```
+```text
+Answer: The average value for sensor 1 is approximately 17.03.
+```

doc/query/llm/index.md

@@ -165,7 +165,7 @@ optional-dependencies.dynamodb = [
   "cratedb-toolkit[io-recipe]",
 ]
 optional-dependencies.full = [
-  "cratedb-toolkit[cfr,datasets,docs-api,io-curated,mcp,service]",
+  "cratedb-toolkit[cfr,datasets,docs-api,io-curated,llm,mcp,service]",
 ]
 optional-dependencies.iceberg = [
   "cratedb-toolkit[io]",
@@ -212,6 +212,15 @@ optional-dependencies.kinesis = [
   "cratedb-toolkit[io-recipe]",
   "lorrystream[carabas]>=0.0.6",
 ]
+optional-dependencies.llm = [
+  "langchain-openai<1.2",
+  "llama-index-embeddings-langchain<0.6",
+  "llama-index-embeddings-openai<0.7",
+  "llama-index-llms-azure-openai<0.6",
+  "llama-index-llms-ollama<0.11",
+  "llama-index-llms-openai<0.8",
+  "llama-index-tools-mcp<0.5",
+]
 optional-dependencies.mcp = [
   "mcp<1.28; python_version>='3.10'",
 ]