Text-to-SQL: Copy editing. Suggestions by CodeRabbit.

Author: amotl

.github/workflows/main.yml

@@ -35,8 +35,6 @@ jobs:
     env:
       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
@@ -76,6 +74,8 @@ jobs:
       env:
         KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
         KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
+        ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
       run: |
         poe check
 

cratedb_toolkit/query/nlsql/api.py

@@ -11,14 +11,15 @@
 
 logger = logging.getLogger(__name__)
 
+llama_index_import_error: Optional[ImportError] = None
 
 try:
     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
-except ImportError:
-    pass
+except ImportError as exc:
+    llama_index_import_error = exc
 
 
 @dataclasses.dataclass
@@ -53,6 +54,11 @@ def setup(self):
         """Configure database connection and query engine."""
         from cratedb_toolkit.query.nlsql.util import configure_llm
 
+        if llama_index_import_error:
+            raise ImportError(
+                "NLSQL support requires installing `cratedb-toolkit[nlsql]`"
+            ) from llama_index_import_error
+
         # Configure model.
         logger.info("Configuring LLM model")
         llm: LLM = configure_llm(self.model)

cratedb_toolkit/query/nlsql/model.py

@@ -47,17 +47,27 @@ def from_options(
             else:
                 raise ValueError("LLM completion model not defined")
         if not llm_api_key:
-            if provider in [ModelProvider.OPENAI, ModelProvider.AZURE]:
+            if provider is ModelProvider.OPENAI:
                 llm_api_key = os.getenv("OPENAI_API_KEY")
                 if not llm_api_key:
                     raise ValueError(
-                        "LLM API key not defined. Use either API option or OPENAI_API_KEY environment variable."
+                        "LLM API key not defined. Use either CLI/API parameter or OPENAI_API_KEY environment variable."
                     )
-            elif provider in [ModelProvider.ANTHROPIC]:
+            elif provider is ModelProvider.AZURE:
+                llm_endpoint = llm_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
+                llm_api_key = os.getenv("AZURE_OPENAI_API_KEY")
+                llm_api_version = llm_api_version or os.getenv("OPENAI_API_VERSION")
+                if not llm_api_key:
+                    raise ValueError(
+                        "LLM API key not defined. Use either CLI/API parameter or "
+                        "AZURE_OPENAI_API_KEY environment variable."
+                    )
+            elif provider is 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."
+                        "LLM API key not defined. Use either CLI/API parameter or "
+                        "ANTHROPIC_API_KEY environment variable."
                     )
         return cls(
             provider=provider,

tests/query/test_nlsql.py

@@ -8,13 +8,13 @@
 from cratedb_toolkit.query.cli import cli
 
 if sys.version_info < (3, 10):
-    pytest.skip("Only available for Python 3.10+", allow_module_level=True)
+    pytest.skip("Only available for Python 3.10+", allow_module_level=True)  # ty: ignore[invalid-argument-type,too-many-positional-arguments]
 
 
 @pytest.fixture
 def provision_db(cratedb):
     sql_ddl = """
-CREATE TABLE IF NOT EXISTS testdrive.time_series_data (
+CREATE TABLE testdrive.time_series_data (
     timestamp TIMESTAMP,
     value DOUBLE,
     location STRING,