Text-to-SQL: Add Anthropic provider

Author: amotl

.github/workflows/main.yml

@@ -36,6 +36,7 @@ jobs:
       OS: ${{ matrix.os }}
       PYTHON: ${{ matrix.python-version }}
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
       # Do not tear down Testcontainers
       TC_KEEPALIVE: true
     # https://docs.github.com/en/actions/using-containerized-services/about-service-containers

cratedb_toolkit/query/nlsql/api.py

@@ -51,14 +51,11 @@ def __post_init__(self):
 
     def setup(self):
         """Configure database connection and query engine."""
-        logger.info("Connecting to CrateDB")
+        from cratedb_toolkit.query.nlsql.util import configure_llm
 
         # Configure model.
         logger.info("Configuring LLM model")
-        llm: LLM
-        from cratedb_toolkit.query.nlsql.util import configure_llm
-
-        llm = configure_llm(self.model)
+        llm: LLM = configure_llm(self.model)
 
         # Configure query engine.
         logger.info("Creating query engine")

cratedb_toolkit/query/nlsql/model.py

@@ -10,6 +10,7 @@ class ModelProvider(Enum):
     """Model provider choices."""
 
     OPENAI = "openai"
+    ANTHROPIC = "anthropic"
     AZURE = "azure"
     OLLAMA = "ollama"
 
@@ -41,6 +42,8 @@ def from_options(
                 llm_name = "gpt-4.1"
             elif provider in [ModelProvider.OLLAMA]:
                 llm_name = "gemma3:1b"
+            elif provider in [ModelProvider.ANTHROPIC]:
+                llm_name = "claude-sonnet-4-0"
             else:
                 raise ValueError("LLM completion model not defined")
         if not llm_api_key:
@@ -50,6 +53,12 @@ def from_options(
                     raise ValueError(
                         "LLM API key not defined. Use either API option or OPENAI_API_KEY environment variable."
                     )
+            elif provider in [ModelProvider.ANTHROPIC]:
+                llm_api_key = os.getenv("ANTHROPIC_API_KEY")
+                if not llm_api_key:
+                    raise ValueError(
+                        "LLM API key not defined. Use either API option or ANTHROPIC_API_KEY environment variable."
+                    )
         return cls(
             provider=provider,
             endpoint=llm_endpoint,

cratedb_toolkit/query/nlsql/util.py

@@ -1,12 +1,27 @@
-import llama_index.core
+from typing import Optional
+
+from llama_index.core import MockEmbedding, set_global_handler, settings
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.embeddings import utils
+from llama_index.core.embeddings.utils import EmbedType
 from llama_index.core.llms import LLM
+from llama_index.llms.anthropic import Anthropic
 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.nlsql.model import ModelInfo, ModelProvider
 
 
+def resolve_embed_model(
+    embed_model: Optional[EmbedType] = None,
+    callback_manager: Optional[CallbackManager] = None,
+) -> BaseEmbedding:
+    """Stub function for disabling embeddings without the `print` and other side effects."""
+    return MockEmbedding(embed_dim=1)
+
+
 def configure_llm(info: ModelInfo, debug: bool = False) -> LLM:
     """
     Configure LLM access and model types. Use either vanilla Open AI, Azure Open AI, or Ollama.
@@ -16,14 +31,18 @@ def configure_llm(info: ModelInfo, debug: bool = False) -> LLM:
 
     completion_model = info.name
 
+    # Disable embeddings.
+    utils.resolve_embed_model = resolve_embed_model  # ty: ignore[invalid-assignment]
+    settings.resolve_embed_model = resolve_embed_model  # ty: ignore[invalid-assignment]
+
     if not info.provider:
         raise ValueError("LLM model provider 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")
+        set_global_handler("simple")
 
     # Select completions model.
     if info.provider is ModelProvider.OPENAI:
@@ -53,7 +72,14 @@ def configure_llm(info: ModelInfo, debug: bool = False) -> LLM:
             request_timeout=120.0,
             keep_alive=-1,
         )
+    elif info.provider is ModelProvider.ANTHROPIC:
+        llm = Anthropic(
+            model=completion_model,
+            temperature=0.0,
+            base_url=info.endpoint,
+            api_key=info.api_key,
+        )
     else:
-        raise ValueError(f"LLM model provider not found: {info.provider}")
+        raise ValueError(f"LLM model provider not implemented: {info.provider}")
 
     return llm

doc/query/nlsql/index.md

@@ -18,6 +18,12 @@ export LLM_PROVIDER=openai
 export OPENAI_API_KEY=<YOUR_OPENAI_API_KEY>
 ```
 
+```shell
+export CRATEDB_CLUSTER_URL=crate://localhost/
+export LLM_PROVIDER=anthropic
+export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>
+```
+
 ```shell
 export CRATEDB_CLUSTER_URL=crate://localhost/
 export LLM_PROVIDER=ollama
@@ -38,16 +44,25 @@ engine = sa.create_engine("crate://")
 schema = "doc"
 
 # Use Open AI GPT-4.
-dq = DataQuery(
+dataquery = DataQuery(
     db=DatabaseInfo(engine=engine, schema=schema),
     model=ModelInfo(provider=ModelProvider.OPENAI, name="gpt-4.1"),
 )
 
-# Use Gemma3 via Ollama.
-dq = DataQuery(
+# Use Anthropic Claude Sonnet.
+dataquery = DataQuery(
+    db=DatabaseInfo(engine=engine, schema=schema),
+    model=ModelInfo(provider=ModelProvider.ANTHROPIC, name="claude-sonnet-4-0"),
+)
+
+# Use Google Gemma3 via Ollama.
+dataquery = DataQuery(
     db=DatabaseInfo(engine=engine, schema=schema),
     model=ModelInfo(provider=ModelProvider.OLLAMA, name="gemma3:1b"),
 )
+
+response = dataquery.ask("What is the average value for sensor 1?")
+print(response)
 ```
 
 ## Example

pyproject.toml

@@ -225,6 +225,7 @@ optional-dependencies.mongodb = [
   "undatum<1.2",
 ]
 optional-dependencies.nlsql = [
+  "llama-index-llms-anthropic<0.12; python_version>='3.10'",
   "llama-index-llms-azure-openai<0.6; python_version>='3.10'",
   "llama-index-llms-ollama<0.11; python_version>='3.10'",
   "llama-index-llms-openai<0.8; python_version>='3.10'",

tests/query/test_nlsql.py

@@ -14,15 +14,15 @@
 @pytest.fixture
 def provision_db(cratedb):
     sql_ddl = """
-CREATE TABLE IF NOT EXISTS time_series_data (
+CREATE TABLE IF NOT EXISTS testdrive.time_series_data (
     timestamp TIMESTAMP,
     value DOUBLE,
     location STRING,
     sensor_id INT
 );
 """
     sql_dml = """
-INSERT INTO time_series_data (timestamp, value, location, sensor_id)
+INSERT INTO testdrive.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),
@@ -42,14 +42,15 @@ def provision_db(cratedb):
 
 
 @pytest.mark.skipif("OPENAI_API_KEY" not in os.environ, reason="OPENAI_API_KEY not set")
-def test_query_llm(cratedb, provision_db):
+def test_query_nlsql_openai(cratedb, provision_db):
     """
-    Verify `ctk query nlsql ...`.
+    Verify `ctk query nlsql ...` with Open AI.
     """
 
     runner = CliRunner(
         env={
             "CRATEDB_CLUSTER_URL": cratedb.get_connection_url(),
+            "CRATEDB_SCHEMA": "testdrive",
             "LLM_PROVIDER": "openai",
         }
     )
@@ -64,3 +65,29 @@ def test_query_llm(cratedb, provision_db):
     output = json.loads(result.output)
     assert output["answer"] == "The average value for sensor 1 is approximately 17.03."
     assert output["sql_query"] == "SELECT AVG(value) FROM time_series_data WHERE sensor_id = 1"
+
+
+@pytest.mark.skipif("ANTHROPIC_API_KEY" not in os.environ, reason="ANTHROPIC_API_KEY not set")
+def test_query_nlsql_anthropic(cratedb, provision_db):
+    """
+    Verify `ctk query nlsql ...` with Anthropic.
+    """
+
+    runner = CliRunner(
+        env={
+            "CRATEDB_CLUSTER_URL": cratedb.get_connection_url(),
+            "CRATEDB_SCHEMA": "testdrive",
+            "LLM_PROVIDER": "anthropic",
+        }
+    )
+
+    result = runner.invoke(
+        cli,
+        input="What is the average value for sensor 1?",
+        args="nlsql -",
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0, result.output
+    output = json.loads(result.output)
+    assert "the average value for sensor 1 is approximately **17.03**" in output["answer"]
+    assert output["sql_query"] == "SELECT AVG(value) as average_value FROM time_series_data WHERE sensor_id = 1;"