Conversation
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 45 minutes and 18 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
WalkthroughA new LLM-to-database query module is introduced with domain models for LLM and database configuration, a utility function to initialize LLM backends (OpenAI, Azure, Ollama), and a Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant QueryData
participant LLMConfig as LLM Config
participant LlamaIndex
participant CrateDB
User->>QueryData: Create QueryData(db, model)
activate QueryData
QueryData->>QueryData: __post_init__()
QueryData->>QueryData: setup()
QueryData->>LLMConfig: configure_llm(model)
activate LLMConfig
LLMConfig->>LLMConfig: Instantiate LLM Backend<br/>(OpenAI/Azure/Ollama)
LLMConfig->>LLMConfig: Create Embedding Model
LLMConfig-->>QueryData: (llm, embedding)
deactivate LLMConfig
QueryData->>LlamaIndex: Create NLSQLTableQueryEngine<br/>(db, llm, embedding)
LlamaIndex->>CrateDB: Connect & Inspect Schema
deactivate QueryData
User->>QueryData: ask(question)
activate QueryData
QueryData->>QueryData: Validate query_engine
QueryData->>LlamaIndex: query(question)
activate LlamaIndex
LlamaIndex->>CrateDB: Generate & Execute SQL
CrateDB-->>LlamaIndex: Results
LlamaIndex-->>QueryData: Response with metadata
deactivate LlamaIndex
QueryData->>QueryData: Log response & sources
QueryData-->>User: Return LlamaIndex Response
deactivate QueryData
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cratedb_toolkit/query/llm/api.py`:
- Around line 84-88: Change the four logger.info calls that print full
query/answer/sources/metadata to DEBUG-level logging (or redact sensitive
fields) to avoid leaking query results: replace logger.info("Query was: %s",
question), logger.info("Answer was: %s", response),
logger.info(response.get_formatted_sources()), and
logger.info(response.metadata) with logger.debug equivalents (or sanitize
response/metadata before logging) so only DEBUG mode exposes full content while
normal INFO logs remain safe.
- Around line 54-59: The setup method creates a new SQLAlchemy engine and calls
an unscoped connect(), which ignores the caller-provided DatabaseInfo.engine and
leaks a connection; change setup() to use the existing engine on self.db
(self.db.engine) instead of creating engine_crate and remove the standalone
engine_crate.connect() call (or if a scoped connection is required, acquire it
via a context manager from self.db.engine). Ensure all references to
engine_crate are replaced with self.db.engine so the caller's configuration is
honored and no permanent open connection is left.
- Around line 67-70: SQLDatabase is being instantiated without passing through
DatabaseInfo filters, so schema/include_tables/ignore_tables set on DatabaseInfo
are ignored; update the setup() method where SQLDatabase(engine_crate, ...) is
created to forward DatabaseInfo.schema, DatabaseInfo.include_tables, and
DatabaseInfo.ignore_tables into the SQLDatabase constructor (or compute the
include/ignore lists from DatabaseInfo if names differ), ensuring the
SQLDatabase instance receives the schema and table filters defined on the
DatabaseInfo object.
In `@cratedb_toolkit/query/llm/util.py`:
- Around line 47-60: The code currently ignores ModelInfo fields and always
reads Azure/OpenAI config from environment variables; update the
OpenAI/AzureOpenAI construction to prefer values from the passed ModelInfo
instance (info.api_key, info.endpoint, info.instance, info.api_version) and only
fall back to os.getenv(...) when those ModelInfo fields are empty/None;
specifically modify the branches that create OpenAI and AzureOpenAI (referencing
llm, completion_model, info, ModelType.AZURE) so they pass api_key=info.api_key
or os.getenv(...), azure_endpoint=info.endpoint or os.getenv(...),
engine=info.instance or os.getenv(...), and api_version=info.api_version or
os.getenv(...).
- Around line 73-85: The embed_model variable currently returns the string
"local" in the Ollama branch which violates the declared BaseEmbedding return
type; change that branch to return None (since NLSQLTableQueryEngine.embed_model
accepts Optional[BaseEmbedding]) and update the function's return type
annotation to reflect Optional[BaseEmbedding] (or a tuple containing
Optional[BaseEmbedding] if the function returns (llm, embed_model)); also remove
the ty: ignore[invalid-return-type] comment and ensure references to ModelType
and the return tuple (llm, embed_model) in this function are updated
accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: e342e93e-bc9c-4a49-9446-79f272a76782
📒 Files selected for processing (5)
CHANGES.mdcratedb_toolkit/query/llm/__init__.pycratedb_toolkit/query/llm/api.pycratedb_toolkit/query/llm/model.pycratedb_toolkit/query/llm/util.py
| llm = OpenAI( | ||
| model=completion_model, | ||
| temperature=0.0, | ||
| api_key=os.getenv("OPENAI_API_KEY"), | ||
| ) | ||
| elif info.type is ModelType.AZURE: | ||
| llm = AzureOpenAI( | ||
| model=completion_model, | ||
| temperature=0.0, | ||
| engine=os.getenv("LLM_INSTANCE"), | ||
| azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"), | ||
| api_key=os.getenv("OPENAI_API_KEY"), | ||
| api_version=os.getenv("OPENAI_AZURE_API_VERSION"), | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n cratedb_toolkit/query/llm/util.py | head -100Repository: crate/cratedb-toolkit
Length of output: 3759
🏁 Script executed:
cat -n cratedb_toolkit/query/llm/model.pyRepository: crate/cratedb-toolkit
Length of output: 906
Honor ModelInfo fields before falling back to env vars.
ModelInfo.endpoint, instance, api_key, and api_version are defined but ignored here, which makes programmatic configuration unreliable. Prefer ModelInfo values first, then env fallbacks.
Proposed fix
if info.type is ModelType.OPENAI:
llm = OpenAI(
model=completion_model,
temperature=0.0,
- api_key=os.getenv("OPENAI_API_KEY"),
+ api_key=info.api_key or os.getenv("OPENAI_API_KEY"),
)
@@
llm = AzureOpenAI(
model=completion_model,
temperature=0.0,
- engine=os.getenv("LLM_INSTANCE"),
- azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"),
- api_key=os.getenv("OPENAI_API_KEY"),
- api_version=os.getenv("OPENAI_AZURE_API_VERSION"),
+ engine=info.instance or os.getenv("LLM_INSTANCE"),
+ azure_endpoint=info.endpoint or os.getenv("OPENAI_AZURE_ENDPOINT"),
+ api_key=info.api_key or os.getenv("OPENAI_API_KEY"),
+ api_version=info.api_version or os.getenv("OPENAI_AZURE_API_VERSION"),
)
@@
AzureOpenAIEmbeddings(
- azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"),
+ azure_endpoint=info.endpoint or os.getenv("OPENAI_AZURE_ENDPOINT"),
model=embedding_model,
)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cratedb_toolkit/query/llm/util.py` around lines 47 - 60, The code currently
ignores ModelInfo fields and always reads Azure/OpenAI config from environment
variables; update the OpenAI/AzureOpenAI construction to prefer values from the
passed ModelInfo instance (info.api_key, info.endpoint, info.instance,
info.api_version) and only fall back to os.getenv(...) when those ModelInfo
fields are empty/None; specifically modify the branches that create OpenAI and
AzureOpenAI (referencing llm, completion_model, info, ModelType.AZURE) so they
pass api_key=info.api_key or os.getenv(...), azure_endpoint=info.endpoint or
os.getenv(...), engine=info.instance or os.getenv(...), and
api_version=info.api_version or os.getenv(...).
| if info.type is ModelType.OPENAI: | ||
| embed_model = LangchainEmbedding(OpenAIEmbeddings(model=embedding_model)) | ||
| elif info.type is ModelType.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] |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
In current LlamaIndex docs for NLSQLTableQueryEngine, what is the accepted type of embed_model? Is a plain string like "local" supported?
💡 Result:
The accepted type for embed_model in the current LlamaIndex docs for NLSQLTableQueryEngine is Optional[BaseEmbedding]. A plain string like "local" is not directly supported as the parameter value; it expects an instance of an embedding model class implementing BaseEmbedding, such as HuggingFaceEmbedding(model_name="BAAI/bge-small-en"). However, if not provided (None), it falls back to Settings.embed_model, which can be set globally to such an instance.
Citations:
- 1: https://developers.llamaindex.ai/python/framework-api-reference/query_engine/NL_SQL_table/
- 2: https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/indices/struct_store/sql_query.py
- 3: https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/indices/struct_store/sql_retriever.py
- 4: https://docs.llamaindex.ai/en/stable/api_reference/query_engine/retriever/
🏁 Script executed:
# First, verify the file exists and examine the function
find . -name "util.py" -path "*/query/llm/*" -type fRepository: crate/cratedb-toolkit
Length of output: 100
🏁 Script executed:
# Once found, read the file to verify code state
cat -n cratedb_toolkit/query/llm/util.py | head -100Repository: crate/cratedb-toolkit
Length of output: 3759
🏁 Script executed:
# Search for usages of configure_llm
rg "configure_llm" -A 3 -B 3Repository: crate/cratedb-toolkit
Length of output: 1541
🏁 Script executed:
# Search for usages of NLSQLTableQueryEngine to understand how embed_model is passed
rg "NLSQLTableQueryEngine" -A 5 -B 2Repository: crate/cratedb-toolkit
Length of output: 2228
Fix type inconsistency in embed_model for Ollama branch.
Line 83 returns a string ("local") while the function signature declares it returns BaseEmbedding. This breaks the contract and will cause a runtime error when NLSQLTableQueryEngine receives the value. The type mismatch is currently masked by ty: ignore[invalid-return-type] on line 85.
Per LlamaIndex API documentation, NLSQLTableQueryEngine.embed_model accepts Optional[BaseEmbedding], not plain strings. Return None instead of "local" and update the return type annotation accordingly.
Proposed fix
-from typing import Tuple
+from typing import Optional, Tuple
@@
-def configure_llm(info: ModelInfo, debug: bool = False) -> Tuple[LLM, BaseEmbedding]:
+def configure_llm(info: ModelInfo, debug: bool = False) -> Tuple[LLM, Optional[BaseEmbedding]]:
@@
- else:
- embed_model = "local"
+ else:
+ embed_model = None
- return llm, embed_model # ty: ignore[invalid-return-type]
+ return llm, embed_model🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cratedb_toolkit/query/llm/util.py` around lines 73 - 85, The embed_model
variable currently returns the string "local" in the Ollama branch which
violates the declared BaseEmbedding return type; change that branch to return
None (since NLSQLTableQueryEngine.embed_model accepts Optional[BaseEmbedding])
and update the function's return type annotation to reflect
Optional[BaseEmbedding] (or a tuple containing Optional[BaseEmbedding] if the
function returns (llm, embed_model)); also remove the ty:
ignore[invalid-return-type] comment and ensure references to ModelType and the
return tuple (llm, embed_model) in this function are updated accordingly.
b231b23 to
51706ec
Compare
DataQuery is the little sister of Google's QueryData product.
About
Exploring Text-to-SQL with non-frontier models.
What's inside
DataQuery is the little sister of Google's QueryData product.
References