Text-to-SQL: Add Mistral provider

Author: amotl

cratedb_toolkit/query/nlsql/model.py

@@ -12,6 +12,7 @@ class ModelProvider(Enum):
     OPENAI = "openai"
     ANTHROPIC = "anthropic"
     AZURE = "azure"
+    MISTRAL = "mistral"
     OLLAMA = "ollama"
 
 
@@ -44,6 +45,8 @@ def from_options(
                 llm_name = "gemma3:1b"
             elif provider in [ModelProvider.ANTHROPIC]:
                 llm_name = "claude-sonnet-4-0"
+            elif provider in [ModelProvider.MISTRAL]:
+                llm_name = "mistral-medium-latest"
             else:
                 raise ValueError("LLM completion model not defined")
         if not llm_api_key:
@@ -69,6 +72,12 @@ def from_options(
                         "LLM API key not defined. Use either CLI/API parameter or "
                         "ANTHROPIC_API_KEY environment variable."
                     )
+            elif provider is ModelProvider.MISTRAL:
+                llm_api_key = os.getenv("MISTRAL_API_KEY")
+                if not llm_api_key:
+                    raise ValueError(
+                        "LLM API key not defined. Use either CLI/API parameter or MISTRAL_API_KEY environment variable."
+                    )
         return cls(
             provider=provider,
             endpoint=llm_endpoint,

cratedb_toolkit/query/nlsql/util.py

@@ -8,6 +8,7 @@
 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.mistralai import MistralAI
 from llama_index.llms.ollama import Ollama
 from llama_index.llms.openai import OpenAI
 
@@ -79,6 +80,13 @@ def configure_llm(info: ModelInfo, debug: bool = False) -> LLM:
             base_url=info.endpoint,
             api_key=info.api_key,
         )
+    elif info.provider is ModelProvider.MISTRAL:
+        llm = MistralAI(
+            model=completion_model,
+            temperature=0.0,
+            endpoint=info.endpoint,
+            api_key=info.api_key,
+        )
     else:
         raise ValueError(f"LLM model provider not implemented: {info.provider}")
 

pyproject.toml

@@ -227,8 +227,10 @@ optional-dependencies.mongodb = [
 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-mistralai<0.11; python_version>='3.10'",
   "llama-index-llms-ollama<0.11; python_version>='3.10'",
   "llama-index-llms-openai<0.8; python_version>='3.10'",
+  "mistralai<2",
 ]
 optional-dependencies.pymongo = [
   "jessiql==1.0.0rc1",