feat(wren-ai-service): get knowledge from ibis (ai-env-changed) #2043
Conversation
…e SQL generation system prompt to a new location and updated its format. Added TEXT_TO_SQL_RULES for better clarity in SQL query generation. Cleaned up the construct_instructions function for improved readability.
… to 0.21.3, 0.21.3, and 0.31.3 respectively.
Implemented the `get_sql_knowledge` method in the WrenIbis class to fetch SQL knowledge from the specified data source. This method handles API requests, manages timeouts, and logs errors appropriately.
…tency. Moved the prompt definition to a centralized location and removed redundant code. This update enhances the structure of SQL query generation in the AI service.
…d structure. The prompt has been redefined to improve the step-by-step instructions for data analysts, ensuring consistency in SQL query generation.
Enhanced the service container by integrating the `SqlKnowledges` retrieval component. This addition allows for improved handling of SQL knowledge within the AI service, ensuring better data processing capabilities.
…structions and prompts. This change enhances modularity and maintainability by centralizing the retrieval of SQL-related instructions and prompts across various generation classes.
…tainer This update introduces the `allow_sql_knowledge_retrieval` setting to the configuration, enabling the integration of SQL knowledge retrieval within the service container. This enhancement improves the overall functionality and data processing capabilities of the AI service.
This update introduces the `allow_sql_knowledge_retrieval` parameter across various services, including AskFeedbackService, AskService, QuestionRecommendation, and SqlCorrectionService. This enhancement allows for conditional execution of SQL knowledge retrieval, improving the flexibility and functionality of the AI service's data processing capabilities.
…n class This update introduces the `SqlKnowledge` class and the `SqlKnowledges` pipeline for enhanced SQL knowledge retrieval. Additionally, a print statement in the `SQLGeneration` class has been removed to streamline the code. These changes improve the modularity and functionality of the AI service's data processing capabilities.
…line component setup. This change enhances modularity by focusing on the model name and prompt builder, improving maintainability across the SQL generation pipeline.
WalkthroughAdds a toggleable SQL knowledge retrieval pipeline (SqlKnowledge/SqlKnowledges) with TTL caching and engine integration, threads sql_knowledge into multiple generation/correction pipelines, refactors generation prompts to use dynamic getters and SqlKnowledge, and defers LLM generator creation to runtime; new flag Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant AskService
participant SqlKnowledgeRetrieval as sql_knowledge_retrieval
participant Cache as TTLCache
participant Engine as WrenIbis
participant Generation as sql_generation
User->>AskService: ask(query, project_id, ...)
alt allow_sql_knowledge_retrieval == true
AskService->>sql_knowledge_retrieval: run(project_id)
sql_knowledge_retrieval->>Cache: check(data_source)
alt cache miss
sql_knowledge_retrieval->>Engine: get_sql_knowledge(session, data_source)
Engine-->>sql_knowledge_retrieval: knowledge JSON / None
sql_knowledge_retrieval->>Cache: store(knowledge)
end
sql_knowledge_retrieval-->>AskService: SqlKnowledge (or None)
else allow_sql_knowledge_retrieval == false
AskService-->>AskService: sql_knowledge = None
end
AskService->>Generation: run(query, docs, sql_knowledge, ...)
Generation->>Generation: get_*_instructions(sql_knowledge)
Generation-->>AskService: generated SQL / responses
AskService-->>User: result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wren-ai-service/src/globals.py (1)
296-333: Makecreate_service_metadatamore robust to missing/irregularpyproject.tomland component fields.Two potential robustness issues:
_get_version_from_pyprojectassumes thatpyproject_pathexists and that["tool"]["poetry"]["version"]is present. If the service is run from a different working directory or thepyproject.tomlstructure changes, this will raise and can prevent startup, even though service metadata is only used for logging.
_convert_pipe_metadataassumes everyPipelineComponentdataclass has bothllm_providerandembedder_providerfields. If any component lacks one of these,asdict(component)will not contain the key and the**expansion will fail with aTypeError.You can mitigate both by:
- Wrapping
_get_version_from_pyprojectintry/exceptand falling back to"unknown"while logging a warning when the file or keys are missing.- Making
embedder_provider(and possiblyllm_provider) optional in_convert_pipe_metadata’s signature, e.g.embedder_provider: EmbedderProvider | None = None, and using.get(...)on the dict passed fromasdict(component).That way, metadata collection won’t break service startup if the environment or component shape evolves.
🧹 Nitpick comments (10)
wren-ai-service/tools/dev/.env (1)
14-17: Minor: Reorder environment variables to match linter conventions.The dotenv-linter reports key ordering violations. While this doesn't affect functionality, reordering the version keys alphabetically would align with linting standards and improve consistency.
Apply this diff to reorder the version keys:
-WREN_PRODUCT_VERSION=development -WREN_ENGINE_VERSION=0.21.3 WREN_AI_SERVICE_VERSION=0.27.14 +WREN_ENGINE_VERSION=0.21.3 IBIS_SERVER_VERSION=0.21.3 +WREN_PRODUCT_VERSION=development WREN_UI_VERSION=0.31.3wren-ai-service/src/web/v1/services/sql_corrections.py (1)
97-100: Consider adding error handling for the knowledge retrieval pipeline.The
sql_knowledge_retrievalpipeline call is not wrapped in error handling. If this call fails, it could propagate an exception and fail the entire correction flow, even though knowledge retrieval is arguably an optional enhancement.Consider wrapping the call in a try-except block:
if self._allow_sql_knowledge_retrieval: - await self._pipelines["sql_knowledge_retrieval"].run( - project_id=project_id, - ) + try: + await self._pipelines["sql_knowledge_retrieval"].run( + project_id=project_id, + ) + except Exception as e: + logger.warning(f"SQL knowledge retrieval failed: {e}")wren-ai-service/src/web/v1/services/ask_feedback.py (1)
166-169: Consider adding error handling for the knowledge retrieval pipeline.Similar to the SQL corrections service, the
sql_knowledge_retrievalpipeline call lacks error handling. If this optional feature fails, it could unnecessarily fail the entire ask feedback flow.Consider wrapping the call in error handling:
if self._allow_sql_knowledge_retrieval: - await self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_feedback_request.project_id, - ) + try: + await self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_feedback_request.project_id, + ) + except Exception as e: + logger.warning(f"SQL knowledge retrieval failed: {e}")wren-ai-service/src/web/v1/services/ask.py (1)
344-347: Consider adding error handling for the knowledge retrieval pipeline.Like the other services, this
sql_knowledge_retrievalcall lacks error handling. If the optional knowledge retrieval fails, it could unnecessarily terminate the entire ask flow.Consider adding error handling:
if self._allow_sql_knowledge_retrieval: - await self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) + try: + await self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + except Exception as e: + logger.warning(f"SQL knowledge retrieval failed: {e}")wren-ai-service/src/web/v1/services/question_recommendation.py (1)
36-44: Knowledge retrieval flag wiring is correct; consider avoiding per-question calls.The new
allow_sql_knowledge_retrievalflag and gated call into thesql_knowledge_retrievalpipeline look consistent with the rest of the service wiring.However, placing the call inside
_validate_questionmeanssql_knowledge_retrievalwill run once per candidate question, even though the knowledge is keyed by data source and internally cached. That still incurs repeated metadata lookups andset_sql_knowledgecalls for the sameproject_id.If you want to trim overhead, consider moving the knowledge retrieval to
_recommend(once perevent_id/ request) and relying on the global state it sets, or running it as a separate task in the initialasyncio.gatherrather than inside each validation call.Also applies to: 117-120
wren-ai-service/src/globals.py (1)
153-156: Share a singleSqlKnowledgespipeline instance across services to avoid duplicate caches.
retrieval.SqlKnowledgesis instantiated separately forask_service,ask_feedback_service,question_recommendation, andsql_correction_service, each with its ownTTLCache. Since SQL knowledge is keyed by data source and shared across these flows, this leads to duplicated caches and potentially duplicated engine calls for the sameproject_id/data_source.Consider constructing a single
_sql_knowledge_retrieval_pipeline = retrieval.SqlKnowledges(**pipe_components["sql_knowledge_retrieval"])alongside the other shared pipelines (_db_schema_retrieval_pipeline,_sql_pair_retrieval_pipeline, etc.) and reusing that instance in all servicepipelinesdicts. That keeps caching andset_sql_knowledgebehavior centralized and should slightly reduce load on the engine.Also applies to: 161-162, 178-185, 236-242, 271-276
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (2)
21-33: ClarifySqlKnowledge.emptyintent (returns a boolean, not an instance).
SqlKnowledge.emptyis a classmethod that returns a boolean based on the rawsql_knowledgedict. The nameemptyreads more like a constructor or mutator than a predicate.To improve readability, consider renaming it to something like
is_empty_dictoris_emptyand optionally making it a@staticmethodsince it does not usecls.
96-111: AlignSqlKnowledgesengine typing with Wren-specific usage.
SqlKnowledges.__init__currently typesengineasEngine, butget_knowledgeexpects aWrenIbisand callsengine.get_sql_knowledge(...), which is not part of theEngineinterface.At runtime this works as long as the concrete engine is a
WrenIbis, but the mismatch can confuse readers and static analyzers.Consider either:
- Typing
engineinSqlKnowledges.__init__asWrenIbisto reflect the actual requirement, or- Extending the
Engineabstraction (or introducing a protocol) to includeget_sql_knowledge, and updatingget_knowledge’s annotation to use that abstraction instead of the concrete Wren engine.This will make the dependency on Wren’s knowledge API explicit and easier to reason about.
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
165-165: Static default instruction constants look good (optional structuring only).Centralizing the text‑to‑SQL, calculated field, metric, and JSON instructions into
_DEFAULT_*constants makes the prompts easier to manage and override; content-wise they’re fine. If these blocks keep growing, consider moving them into a separate module or configuration resource to keep this utility file lighter, but that’s purely optional.Also applies to: 225-225, 272-272, 363-363
575-584: Filter out missing/empty instructions to avoidNoneentries.
construct_instructionscurrently appendsinstruction.get("instruction")for every dict, which can introduceNonevalues (or empty strings) into the returned list if the key is missing or empty. If downstream code expectslist[str], this can lead to subtle issues.You can defensively filter here:
-def construct_instructions( - instructions: list[dict] | None = None, -): - _instructions = [] - if instructions: - _instructions += [ - instruction.get("instruction") for instruction in instructions - ] - - return _instructions +def construct_instructions( + instructions: list[dict] | None = None, +) -> list[str]: + _instructions: list[str] = [] + if instructions: + _instructions += [ + instr_val + for instr in instructions + if (instr_val := instr.get("instruction")) + ] + return _instructions(You can drop the walrus operator if you prefer a more conservative style.)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (15)
wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(4 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(5 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(5 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(6 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py(1 hunks)wren-ai-service/src/providers/engine/wren.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/web/v1/services/ask_feedback.py(3 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(2 hunks)wren-ai-service/src/web/v1/services/sql_corrections.py(2 hunks)wren-ai-service/tools/dev/.env(1 hunks)
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: paopa
Repo: Canner/WrenAI PR: 976
File: wren-ai-service/tools/dev/.env:17-17
Timestamp: 2025-02-03T06:25:14.605Z
Learning: The IBIS server version sha-ce21e44 refers to a commit in the internal wren-engine repository (https://github.com/Canner/wren-engine), not the upstream ibis-project/ibis repository. This version includes fixes for server startup issues.
Learnt from: wwwy3y3
Repo: Canner/WrenAI PR: 1585
File: wren-ui/src/pages/api/v1/generate_sql.ts:98-106
Timestamp: 2025-04-24T16:10:43.308Z
Learning: In the generate_sql API, allow users to specify language codes not predefined in the WrenAILanguage enum, passing them through directly rather than strictly validating against the enum.
📚 Learning: 2025-02-03T06:25:14.605Z
Learnt from: paopa
Repo: Canner/WrenAI PR: 976
File: wren-ai-service/tools/dev/.env:17-17
Timestamp: 2025-02-03T06:25:14.605Z
Learning: The IBIS server version sha-ce21e44 refers to a commit in the internal wren-engine repository (https://github.com/Canner/wren-engine), not the upstream ibis-project/ibis repository. This version includes fixes for server startup issues.
Applied to files:
wren-ai-service/tools/dev/.env
📚 Learning: 2025-02-12T22:05:37.109Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1293
File: wren-ai-service/pyproject.toml:38-38
Timestamp: 2025-02-12T22:05:37.109Z
Learning: The qdrant-client package version in wren-ai-service must match the Qdrant server version (1.11.0) to maintain compatibility.
Applied to files:
wren-ai-service/tools/dev/.env
📚 Learning: 2025-06-20T02:37:21.292Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.
Applied to files:
wren-ai-service/src/config.py
📚 Learning: 2025-01-27T08:55:44.919Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1236
File: wren-ai-service/src/providers/llm/litellm.py:49-49
Timestamp: 2025-01-27T08:55:44.919Z
Learning: In the LitellmLLMProvider class, model parameters (`self._model_kwargs`) are intentionally designed to take precedence over runtime parameters (`generation_kwargs`), which differs from other LLM providers. This is the intended behavior specific to this provider.
Applied to files:
wren-ai-service/src/pipelines/generation/followup_sql_generation.pywren-ai-service/src/pipelines/generation/sql_generation.pywren-ai-service/src/pipelines/generation/sql_regeneration.py
🧬 Code graph analysis (12)
wren-ai-service/src/web/v1/services/ask_feedback.py (3)
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(177-224)wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
run(181-216)wren-ai-service/src/pipelines/retrieval/sql_functions.py (1)
run(104-125)
wren-ai-service/src/pipelines/retrieval/__init__.py (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
SqlKnowledges(96-140)
wren-ai-service/src/web/v1/services/ask.py (2)
wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py (2)
run(23-34)run(151-161)wren-ai-service/tests/pytest/services/mocks.py (6)
run(12-13)run(20-21)run(28-29)run(36-37)run(44-53)run(61-74)
wren-ai-service/src/web/v1/services/sql_corrections.py (4)
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(158-193)wren-ai-service/src/pipelines/retrieval/sql_functions.py (1)
run(104-125)wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)
run(499-519)wren-ai-service/tests/pytest/services/mocks.py (6)
run(12-13)run(20-21)run(28-29)run(36-37)run(44-53)run(61-74)
wren-ai-service/src/web/v1/services/question_recommendation.py (2)
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(177-224)wren-ai-service/src/pipelines/retrieval/instructions.py (3)
run(23-34)run(42-54)run(215-228)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (2)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(509-514)get_json_field_instructions(525-530)get_metric_instructions(517-522)get_sql_generation_system_prompt(533-557)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)
wren-ai-service/src/pipelines/generation/sql_correction.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
get_text_to_sql_rules(501-506)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)wren-ai-service/src/core/provider.py (1)
get_generator(8-9)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (6)
wren-ai-service/src/core/engine.py (1)
Engine(17-26)wren-ai-service/src/core/pipeline.py (1)
BasicPipeline(14-20)wren-ai-service/src/core/provider.py (1)
DocumentStoreProvider(34-41)wren-ai-service/src/pipelines/common.py (1)
retrieve_metadata(67-85)wren-ai-service/src/pipelines/generation/utils/sql.py (1)
set_sql_knowledge(459-498)wren-ai-service/src/providers/engine/wren.py (1)
get_sql_knowledge(266-286)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (6)
SqlKnowledge(21-59)text_to_sql_rule(36-37)calculated_field_instructions(44-45)metric_instructions(48-49)json_field_instructions(52-53)instructions(40-41)
wren-ai-service/src/pipelines/generation/sql_generation.py (4)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(509-514)get_json_field_instructions(525-530)get_metric_instructions(517-522)get_sql_generation_system_prompt(533-557)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (2)
metric_instructions(48-49)json_field_instructions(52-53)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)wren-ai-service/src/core/provider.py (1)
get_generator(8-9)
wren-ai-service/src/pipelines/generation/sql_regeneration.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(509-514)get_json_field_instructions(525-530)get_metric_instructions(517-522)get_text_to_sql_rules(501-506)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)wren-ai-service/src/core/provider.py (1)
get_generator(8-9)
wren-ai-service/src/globals.py (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
SqlKnowledges(96-140)
🪛 dotenv-linter (4.0.0)
wren-ai-service/tools/dev/.env
[warning] 14-14: [UnorderedKey] The WREN_ENGINE_VERSION key should go before the WREN_PRODUCT_VERSION key
(UnorderedKey)
[warning] 15-15: [UnorderedKey] The WREN_AI_SERVICE_VERSION key should go before the WREN_ENGINE_VERSION key
(UnorderedKey)
[warning] 16-16: [UnorderedKey] The IBIS_SERVER_VERSION key should go before the WREN_AI_SERVICE_VERSION key
(UnorderedKey)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (17)
wren-ai-service/tools/dev/.env (1)
14-14: All specified versions are available and valid.Verification confirms:
- WREN_ENGINE_VERSION 0.21.3 exists in Canner/wren-engine
- IBIS_SERVER_VERSION 0.21.3 exists in Canner/wren-engine
- WREN_UI_VERSION 0.31.3 exists as a container image (ghcr.io/canner/wren-ui)
The development environment configuration is ready.
wren-ai-service/src/config.py (1)
44-44: LGTM! Configuration flag follows established patterns.The new
allow_sql_knowledge_retrievalfield is consistently placed and named, following the same pattern as other feature flags in the generation config section.wren-ai-service/src/web/v1/services/sql_corrections.py (1)
29-38: LGTM! Constructor parameter follows service patterns.The
allow_sql_knowledge_retrievalparameter is properly added to the constructor and stored as an instance variable, consistent with how other feature flags are handled in this service.wren-ai-service/src/pipelines/retrieval/__init__.py (1)
7-7: LGTM! Module exposure follows established patterns.The
SqlKnowledgesimport and export are correctly added to the retrieval package's public API, maintaining consistency with other retrieval components.Also applies to: 18-18
wren-ai-service/src/web/v1/services/ask_feedback.py (1)
59-74: LGTM! Constructor parameter is properly integrated.The
allow_sql_knowledge_retrievalparameter and instance variable follow the same pattern as other feature flags in this service.wren-ai-service/src/pipelines/generation/sql_generation.py (3)
18-21: LGTM! Dynamic prompt getters enable runtime configuration.The transition from static imports to getter functions allows the system prompt and instructions to be determined at runtime, incorporating retrieved SQL knowledge when available.
104-110: LGTM! Conditional instruction retrieval is properly implemented.The prompt builder correctly calls the instruction getter functions only when the corresponding feature flags are enabled, providing empty strings otherwise.
162-162: Architectural shift: Generator initialization moved to runtime.The generator is now created on each
run()call rather than once at initialization. This enables the system prompt to incorporate dynamically retrieved SQL knowledge, but introduces per-request generator instantiation overhead.This trade-off is reasonable given the feature requirements, and the pattern is consistently applied across generation pipelines. However, if performance becomes a concern, consider caching generators per knowledge state.
Also applies to: 200-203
wren-ai-service/src/web/v1/services/ask.py (1)
98-123: LGTM! Constructor parameter is consistently integrated.The
allow_sql_knowledge_retrievalparameter follows the same pattern used for other feature flags in the AskService.wren-ai-service/src/pipelines/generation/followup_sql_generation.py (3)
19-22: LGTM! Import changes align with sql_generation.py.The dynamic prompt getter imports are consistently applied across generation pipelines.
108-113: LGTM! Prompt construction matches the established pattern.The conditional instruction retrieval using getters is consistently implemented across generation pipelines.
165-165: LGTM! Runtime generator initialization is consistently applied.The follow-up SQL generation pipeline adopts the same runtime generator initialization pattern as the base SQL generation pipeline, ensuring consistent behavior across the codebase.
Also applies to: 203-206
wren-ai-service/src/providers/engine/wren.py (1)
266-286: LGTM! New knowledge retrieval method follows established patterns.The
get_sql_knowledge()method is well-implemented:
- Consistent error handling with
get_func_list()- Appropriate use of
Optional[Dict[str, Any]]return type for graceful degradation- Proper timeout and exception handling
- Clear logging for debugging
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
27-51: Dynamic correction system prompt and lazy generator initialization look good.Building the SQL-correction system prompt via
get_text_to_sql_rules()and instantiating the generator insiderun()is consistent with the other generation pipelines and ensures any updated SQL rules are reflected at execution time. No issues spotted here.Also applies to: 170-173
wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
29-47: Regeneration prompt and generator initialization are consistent with dynamic SQL knowledge design.The new
get_sql_regeneration_system_prompt()plus the dynamic{calculated,metric,json}_field_instructionsblocks in the user template nicely align regeneration with the shared text-to-SQL rules and connector-specific instructions.Creating the generator in
run()withSQL_GENERATION_MODEL_KWARGSmirrorsSQLGeneration/SQLCorrectionand ensures it sees the latest rules set viasql_knowledge_retrieval. Looks solid.Also applies to: 122-127, 166-167, 196-199
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
426-450: Reasoning system prompt is well-structured and consistent.The updated
sql_generation_reasoning_system_promptclearly constrains the model behavior (no SQL, numbered bold steps, table/column formatting) and aligns with the rest of the pipeline. I don’t see functional issues here.
533-557: Dynamic SQL generation system prompt wiring looks correct.
get_sql_generation_system_prompt()now pulls rules viaget_text_to_sql_rules(), which correctly falls back to_DEFAULT_TEXT_TO_SQL_RULESwhen no override is set. The f‑string escaping for the JSON output template ({{ ... }}) is also correct. Aside from the global-state concern called out separately, this function itself is sound.
…egration This update modifies the SQL generation, correction, and regeneration classes to utilize the `SqlKnowledge` class for improved instruction retrieval. The changes include passing `sql_knowledge` to various functions and methods, ensuring that SQL-related instructions are dynamically fetched based on the available knowledge. This enhancement improves the flexibility and accuracy of SQL query generation and correction processes across the AI service.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
79-88: Verify global SQL knowledge synchronization.The
cache()function now returnsget_knowledgeconsistently, which addresses the return value concern from the previous review. However, the prior review also flagged that global SQL rules should be reset when retrieval fails to prevent cross-request leakage.I don't see any
set_sql_knowledge()call in this file, and looking atsql.py, the global variablesTEXT_TO_SQL_RULES,CALCULATED_FIELD_INSTRUCTIONS, etc. still exist (lines 454-457). If these globals are still being mutated elsewhere in the codebase, the race condition concern remains valid.Run the following script to verify if
set_sql_knowledgeexists and where it's called:#!/bin/bash # Check if set_sql_knowledge function exists and where it's used rg -n "def set_sql_knowledge" --type py -A 5 echo "---" # Check all calls to set_sql_knowledge rg -n "set_sql_knowledge\(" --type py -B 2 -A 2wren-ai-service/src/pipelines/generation/utils/sql.py (1)
454-457: Global mutable SQL knowledge still poses concurrency risk.These module-level globals remain from the previous review where a major race condition issue was identified. The getters (lines 470-521) now accept an optional
sql_knowledgeparameter, which helps when it's provided, but they still fall back to these globals whensql_knowledgeisNone.If any code path mutates these globals (via a
set_sql_knowledge()call or direct assignment), concurrent requests in FastAPI will experience cross-request leakage: Project A's SQL knowledge will overwrite and affect Project B's SQL generation.The architectural recommendation from the previous review remains valid: thread
SqlKnowledgethrough the request context and avoid mutating module-level state per-request.Based on learnings from the previous review, this remains a critical concurrency concern.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py(7 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(5 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(7 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(7 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(7 hunks)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(7 hunks)wren-ai-service/src/web/v1/services/ask_feedback.py(6 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(3 hunks)wren-ai-service/src/web/v1/services/sql_corrections.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- wren-ai-service/src/web/v1/services/ask_feedback.py
- wren-ai-service/src/web/v1/services/question_recommendation.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-01-27T08:55:44.919Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1236
File: wren-ai-service/src/providers/llm/litellm.py:49-49
Timestamp: 2025-01-27T08:55:44.919Z
Learning: In the LitellmLLMProvider class, model parameters (`self._model_kwargs`) are intentionally designed to take precedence over runtime parameters (`generation_kwargs`), which differs from other LLM providers. This is the intended behavior specific to this provider.
Applied to files:
wren-ai-service/src/pipelines/generation/sql_generation.pywren-ai-service/src/pipelines/generation/sql_regeneration.pywren-ai-service/src/pipelines/generation/followup_sql_generation.py
🧬 Code graph analysis (8)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (4)
wren-ai-service/src/core/pipeline.py (1)
BasicPipeline(14-20)wren-ai-service/src/core/provider.py (1)
DocumentStoreProvider(34-41)wren-ai-service/src/pipelines/common.py (1)
retrieve_metadata(67-85)wren-ai-service/src/providers/engine/wren.py (1)
get_sql_knowledge(266-286)
wren-ai-service/src/web/v1/services/ask.py (4)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (1)
run(186-235)wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(159-195)wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(183-235)wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
run(189-226)
wren-ai-service/src/pipelines/generation/sql_generation.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(483-495)get_json_field_instructions(511-521)get_metric_instructions(498-508)get_sql_generation_system_prompt(524-548)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (3)
SqlKnowledge(20-58)metric_instructions(47-48)json_field_instructions(51-52)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)
wren-ai-service/src/pipelines/generation/sql_correction.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
get_text_to_sql_rules(470-480)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
SqlKnowledge(20-58)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (2)
SqlKnowledge(20-58)instructions(39-40)
wren-ai-service/src/web/v1/services/sql_corrections.py (6)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (1)
run(186-235)wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(159-195)wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(183-235)wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
run(189-226)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
run(116-137)wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)
run(499-519)
wren-ai-service/src/pipelines/generation/sql_regeneration.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(483-495)get_json_field_instructions(511-521)get_metric_instructions(498-508)get_text_to_sql_rules(470-480)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (3)
SqlKnowledge(20-58)metric_instructions(47-48)json_field_instructions(51-52)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (4)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(483-495)get_json_field_instructions(511-521)get_metric_instructions(498-508)get_sql_generation_system_prompt(524-548)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (3)
SqlKnowledge(20-58)metric_instructions(47-48)json_field_instructions(51-52)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)wren-ai-service/src/core/provider.py (1)
get_generator(8-9)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (3)
wren-ai-service/src/pipelines/generation/sql_regeneration.py (3)
30-50: Clean implementation of runtime system prompt generation.The
get_sql_regeneration_system_prompt()function correctly acceptssql_knowledgeand delegates toget_text_to_sql_rules(), enabling per-request customization of the system prompt. This is the right approach for avoiding global state contamination.
168-186: Good: Runtime generator initialization pattern.Storing
self._llm_providerinstead of creating the generator at initialization time allows each request to get a fresh generator with asql_knowledge-specific system prompt (created inrun()). This avoids the global state contamination issue and ensures per-request/per-project isolation.
188-226: Correct: Per-request generator with sql_knowledge context.The generator is now created at runtime (lines 205-208) with a system prompt that incorporates the request-specific
sql_knowledge. This ensures that concurrent requests for different projects don't interfere with each other. Thesql_knowledgeis also correctly propagated through the pipeline inputs (line 223).
…de clarity. Update SqlKnowledge class to initialize _data directly from the constructor, improving type consistency.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (1)
20-56: LGTM! Consider more specific type hint for_data.The SqlKnowledge class is well-structured with clear property accessors. The previous class-level mutable default issue has been resolved—
_datais now properly initialized as an instance variable in__init__.Minor suggestion: Consider using
Dict[str, Any]instead of bareDictfor better type specificity.self._data: Dict[str, Any] = sql_knowledge
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/pipelines/generation/sql_generation.py(7 hunks)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-01-27T08:55:44.919Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1236
File: wren-ai-service/src/providers/llm/litellm.py:49-49
Timestamp: 2025-01-27T08:55:44.919Z
Learning: In the LitellmLLMProvider class, model parameters (`self._model_kwargs`) are intentionally designed to take precedence over runtime parameters (`generation_kwargs`), which differs from other LLM providers. This is the intended behavior specific to this provider.
Applied to files:
wren-ai-service/src/pipelines/generation/sql_generation.py
🧬 Code graph analysis (2)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (4)
wren-ai-service/src/core/engine.py (1)
Engine(17-26)wren-ai-service/src/core/provider.py (1)
DocumentStoreProvider(34-41)wren-ai-service/src/pipelines/common.py (1)
retrieve_metadata(67-85)wren-ai-service/src/providers/engine/wren.py (1)
get_sql_knowledge(266-286)
wren-ai-service/src/pipelines/generation/sql_generation.py (4)
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
get_calculated_field_instructions(483-495)get_json_field_instructions(511-521)get_metric_instructions(498-508)get_sql_generation_system_prompt(524-548)wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (3)
SqlKnowledge(20-56)metric_instructions(45-46)json_field_instructions(49-50)wren-ai-service/src/providers/llm/litellm.py (1)
get_generator(62-161)wren-ai-service/src/core/provider.py (1)
get_generator(8-9)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (8)
wren-ai-service/src/pipelines/generation/sql_generation.py (5)
18-21: LGTM! Dynamic instruction imports support the new knowledge retrieval feature.The switch from static instruction imports to getter functions enables runtime-computed, project-aware SQL instructions. The SqlKnowledge import correctly surfaces the new type for downstream usage.
Also applies to: 24-24
86-120: LGTM! Prompt construction correctly uses dynamic getters.The sql_knowledge parameter is properly threaded through to the getter functions, and the conditional logic ensures instructions are only retrieved when the corresponding feature flags are enabled. This aligns with the project-aware, semantic-aware instruction retrieval described in the PR objectives.
168-168: LGTM! Storing the LLM provider enables deferred generator creation.Deferring generator instantiation allows the system prompt to be computed dynamically at runtime based on the sql_knowledge parameter, which can vary per request.
207-210: Verify the performance implications of per-request generator creation.The generator is now instantiated on every
run()call with a dynamically computed system prompt. While this design is necessary sincesql_knowledgevaries per request (and caching at the instance level isn't feasible), it does mean allocating resources and computing the prompt for each SQL generation request.Confirm that the performance overhead is acceptable for the flexibility gained, especially under high request load.
198-198: LGTM! The sql_knowledge parameter is correctly threaded through the pipeline.The parameter is properly added to the method signature and propagated to the pipeline inputs, ensuring downstream components (like the prompt builder) can access the dynamic SQL knowledge.
Also applies to: 229-229
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (3)
61-74: LGTM! Knowledge retrieval logic is correct.The function properly manages the HTTP session, validates the retrieved knowledge using
SqlKnowledge.empty(), and returns a well-typed result. Error handling is delegated to the engine layer (as seen in relevant snippets from wren.py).
78-86: Verify: ShouldNoneresults be cached, and isset_sql_knowledgestill needed?Two observations based on past review feedback:
Noneresults are not cached (line 83): When knowledge retrieval fails or returns empty, it's not stored in the TTL cache. This means every subsequent request for that data source will re-fetch from the engine until a successful retrieval occurs. If failures are expected to persist (e.g., connector doesn't provide knowledge), this could create unnecessary load.No
set_sql_knowledgecall: The past review mentioned thatset_sql_knowledgeshould always be invoked to "keep global SQL rules in sync" and to reset rules to defaults when retrieval fails. I don't see this function imported or called in the current implementation. Confirm whether:
- This global sync mechanism was removed intentionally, or
- It's handled elsewhere in the codebase.
If
Noneshould be cached to prevent repeated fetches:@observe(capture_input=False) def cache( data_source: str, get_knowledge: Optional[SqlKnowledge], ttl_cache: TTLCache, ) -> Optional[SqlKnowledge]: - if get_knowledge: - ttl_cache[data_source] = get_knowledge + # Cache both successful retrieval and None to prevent repeated failed fetches + ttl_cache[data_source] = get_knowledge return get_knowledge
92-135: LGTM! SqlKnowledges pipeline is well-structured.The class properly initializes the retrieval pipeline with TTL caching. The
run()method correctly handles both cache hits (warm path) and misses (cold path), returning a consistentOptional[SqlKnowledge]in both cases sincecache()now returns the knowledge value.The 24-hour TTL and 100-entry maxsize seem reasonable for typical usage patterns, though you may want to monitor cache hit rates in production.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
wren-ai-service/tools/dev/.env (1)
14-17: Version updates align with PR objectives.The version bumps for WREN_ENGINE_VERSION, WREN_AI_SERVICE_VERSION, IBIS_SERVER_VERSION, and WREN_UI_VERSION are consistent with the PR objective to update service dependencies to the latest versions.
Optional: Consider reordering environment keys alphabetically.
Static analysis flagged that the version keys (lines 14-16) are out of alphabetical order. If the project enforces alphabetical ordering in
.envfiles for consistency, you may reorder them:- WREN_PRODUCT_VERSION=development - WREN_ENGINE_VERSION=0.21.3 - WREN_AI_SERVICE_VERSION=0.29.0 - IBIS_SERVER_VERSION=0.21.3 - WREN_UI_VERSION=0.31.3 + IBIS_SERVER_VERSION=0.21.3 + WREN_AI_SERVICE_VERSION=0.29.0 + WREN_BOOTSTRAP_VERSION=0.1.5 + WREN_ENGINE_VERSION=0.21.3 + WREN_PRODUCT_VERSION=development + WREN_UI_VERSION=0.31.3This is a low-priority stylistic concern and can be addressed in a follow-up if the project has a convention for key ordering.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tools/dev/.env(1 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-02-03T06:25:14.605Z
Learnt from: paopa
Repo: Canner/WrenAI PR: 976
File: wren-ai-service/tools/dev/.env:17-17
Timestamp: 2025-02-03T06:25:14.605Z
Learning: The IBIS server version sha-ce21e44 refers to a commit in the internal wren-engine repository (https://github.com/Canner/wren-engine), not the upstream ibis-project/ibis repository. This version includes fixes for server startup issues.
Applied to files:
wren-ai-service/tools/dev/.env
📚 Learning: 2025-02-12T22:05:37.109Z
Learnt from: cyyeh
Repo: Canner/WrenAI PR: 1293
File: wren-ai-service/pyproject.toml:38-38
Timestamp: 2025-02-12T22:05:37.109Z
Learning: The qdrant-client package version in wren-ai-service must match the Qdrant server version (1.11.0) to maintain compatibility.
Applied to files:
wren-ai-service/tools/dev/.env
🪛 dotenv-linter (4.0.0)
wren-ai-service/tools/dev/.env
[warning] 14-14: [UnorderedKey] The WREN_ENGINE_VERSION key should go before the WREN_PRODUCT_VERSION key
(UnorderedKey)
[warning] 15-15: [UnorderedKey] The WREN_AI_SERVICE_VERSION key should go before the WREN_ENGINE_VERSION key
(UnorderedKey)
[warning] 16-16: [UnorderedKey] The IBIS_SERVER_VERSION key should go before the WREN_AI_SERVICE_VERSION key
(UnorderedKey)
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
464-499: Consider simplifying the getter functions.All four getter functions (
get_text_to_sql_rules,get_calculated_field_instructions,get_metric_instructions,get_json_field_instructions) checkif sql_knowledge is not Nonebefore calling_extract_from_sql_knowledge, but the helper already handles None at line 457. This pattern is redundant.Simplify each getter to eliminate the redundant check. For example:
def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: - if sql_knowledge is not None: - return _extract_from_sql_knowledge( - sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES - ) - - return _DEFAULT_TEXT_TO_SQL_RULES + return _extract_from_sql_knowledge( + sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES + )Apply the same simplification to the other three getters.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(7 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/config.py
🧰 Additional context used
🧬 Code graph analysis (1)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
wren-ai-service/src/pipelines/retrieval/sql_knowledge.py (2)
SqlKnowledge(20-56)instructions(37-38)
🔇 Additional comments (7)
wren-ai-service/tools/config/config.example.yaml (2)
177-179: Syntax and configuration look good.The new
sql_knowledge_retrievalpipeline step is correctly configured with the appropriate engine and document store references, and follows the established pattern of similar retrieval pipelines (e.g.,sql_functions_retrieval).
182-206: Verify that theallow_sql_knowledge_retrievalflag should be added to the settings section.The PR objectives state that a new configuration flag
allow_sql_knowledge_retrieval(default: true) controls whether to retrieve knowledge from the SQL knowledge pipeline. However, this flag is not present in the settings section of the config file. For consistency with other feature toggles in settings (e.g.,allow_intent_classification,allow_sql_functions_retrievalat lines 191, 193), this flag should likely be added here to allow users to configure the behavior via the YAML config file.Please confirm:
- Should
allow_sql_knowledge_retrievalbe added to the settings section with a default value (e.g.,allow_sql_knowledge_retrieval: true)?- Or is the flag controlled entirely through code defaults rather than the config file?
wren-ai-service/src/pipelines/generation/utils/sql.py (5)
14-14: LGTM! Global mutable state issue has been properly resolved.The refactoring successfully addresses the previous critical concurrency concern. The code now uses immutable default constants (
_DEFAULT_TEXT_TO_SQL_RULES,_DEFAULT_CALCULATED_FIELD_INSTRUCTIONS,_DEFAULT_METRIC_INSTRUCTIONS,_DEFAULT_JSON_FIELD_INSTRUCTIONS) instead of mutable globals. These constants provide safe fallback values whenSqlKnowledgeis not provided, eliminating the risk of cross-request data leakage in concurrent/multi-tenant scenarios.Also applies to: 166-399
454-461: LGTM! Helper function correctly handles edge cases.The
_extract_from_sql_knowledgehelper properly handles None inputs, missing attributes, and empty/whitespace-only values by falling back to the provided default. The use ofgetattrwith a default and the subsequentstrip()check ensures robust behavior.
502-526: LGTM! System prompt correctly threads SQL knowledge.The
get_sql_generation_system_promptfunction properly acceptsSqlKnowledgeas a parameter and threads it through toget_text_to_sql_rules, enabling per-request SQL knowledge injection. This is the correct pattern for avoiding global state in concurrent environments.
427-451: LGTM! Static prompt constant is safe for concurrent usage.The
sql_generation_reasoning_system_promptis an immutable module-level constant that doesn't change per-request, making it safe for concurrent/multi-tenant usage.
14-14: Excellent refactoring—concurrency safety achieved.This refactoring successfully eliminates the previous critical concurrency issue. The new architecture:
✅ Threads
SqlKnowledgethrough function parameters instead of mutating globals
✅ Uses immutable default constants (_DEFAULT_*) as safe fallbacks
✅ Enables per-request/per-project SQL knowledge without cross-request leakage
✅ Maintains backward compatibility via optionalsql_knowledgeparametersThe code is now safe for FastAPI's async concurrent request handling in multi-tenant environments. Well done!
Also applies to: 166-526
| def construct_instructions( | ||
| instructions: list[dict] | None = None, | ||
| ): | ||
| _instructions = [] | ||
| if instructions: | ||
| _instructions += [ | ||
| instruction.get("instruction") for instruction in instructions | ||
| ] | ||
|
|
||
| return _instructions |
There was a problem hiding this comment.
Handle missing "instruction" keys in dictionaries.
The function uses instruction.get("instruction") which returns None if the key is missing. This could result in None values in the returned list, potentially causing issues in downstream code that expects only strings.
Consider one of these approaches:
def construct_instructions(
instructions: list[dict] | None = None,
):
_instructions = []
if instructions:
_instructions += [
- instruction.get("instruction") for instruction in instructions
+ instruction.get("instruction", "") for instruction in instructions
]
return _instructionsOr filter out None values:
def construct_instructions(
instructions: list[dict] | None = None,
):
_instructions = []
if instructions:
_instructions += [
- instruction.get("instruction") for instruction in instructions
+ instruction.get("instruction")
+ for instruction in instructions
+ if instruction.get("instruction") is not None
]
return _instructions🤖 Prompt for AI Agents
In wren-ai-service/src/pipelines/generation/utils/sql.py around lines 544 to
553, the current construct_instructions builds a list using
instruction.get("instruction") which may append None when the key is missing;
update the comprehension to only include non-None string values (e.g., iterate
instructions, retrieve val = instruction.get("instruction") and append it only
if isinstance(val, str) and val != ""), or filter out falsy/None values after
the list is built so the returned list contains only valid instruction strings.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
wren-ai-service/src/force_update_config.py (2)
21-24: Correctly routesql_knowledge_retrievaltowren_ibisIncluding
"sql_knowledge_retrieval"alongside"sql_functions_retrieval"ensures both knowledge/function retrieval pipelines are consistently mapped towren_ibis, which matches the PR intent. As a small maintainability tweak, consider extracting these special pipeline names into a constantSETat the top of the file so future additions stay centralized.
33-33: Log message no longer matches behaviorThe message
"Successfully updated engine names to 'wren_ui' in all pipelines"is now inaccurate because some pipelines are explicitly set towren_ibis. Consider changing it to something like"Successfully updated pipeline engine names for local development", or similar, to avoid confusion when running this script.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/force_update_config.py(1 hunks)wren-ai-service/src/providers/llm/litellm.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: build-image (linux_arm64_runner, linux/arm64)
- GitHub Check: build-image (ubuntu-latest, linux/amd64)
- GitHub Check: Analyze (go)
| print(f"prompt: {openai_formatted_messages}") | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
Remove or replace with proper logging to avoid logging PII.
This print statement logs all messages sent to the LLM, which may include user queries, emails, and other PII, creating compliance risks (GDPR/CCPA). Additionally, using print() bypasses proper log-level controls and structured logging.
Consider:
- If this is temporary debug code, remove it before merging.
- If logging is needed, use a proper logging framework with appropriate log levels and redaction of sensitive fields.
Is this debug statement intentional, or should it be removed? If intentional, please implement structured logging with PII redaction:
- print(f"prompt: {openai_formatted_messages}")
-
+ # Use proper logging if needed
+ # logger.debug("LLM prompt prepared", extra={"message_count": len(openai_formatted_messages)})
---
🏁 Script executed:
```shell
#!/bin/bash
# View the litellm.py file around lines 97-98 with surrounding context
cat wren-ai-service/src/providers/llm/litellm.py | head -120 | tail -40
Length of output: 1592
🏁 Script executed:
#!/bin/bash
# Check if this is a debug-marked print or commented section
grep -n "print\|logger\|logging" wren-ai-service/src/providers/llm/litellm.py | head -20Length of output: 60
🏁 Script executed:
#!/bin/bash
# Check the imports and setup of the file to understand logging context
head -50 wren-ai-service/src/providers/llm/litellm.pyLength of output: 1697
🏁 Script executed:
#!/bin/bash
# Search for other print statements in the file to determine if this is pattern or isolated
grep -n "print(" wren-ai-service/src/providers/llm/litellm.pyLength of output: 60
Remove debug print statement before merging to prevent logging PII to stdout.
Verification confirms the print statement at lines 97-98 logs complete message content (including user queries, system prompts, and history) directly to stdout on every LLM call. This is the only print statement in the file and lacks both logging framework integration and conditional guards—characteristics of leftover debug code.
Issues confirmed:
- Privacy/Compliance Risk:
openai_formatted_messagescontains full message history that may include PII, emails, or sensitive queries; logging to stdout without redaction violates GDPR/CCPA. - Code Quality: Uses
print()instead of proper logging, with no log-level control or ability to disable without code changes. - High Execution Frequency: Runs on every LLM generation call, amplifying exposure risk.
Remove this statement before merging. If logging is genuinely needed for debugging, implement structured logging with appropriate log levels and PII redaction.
🤖 Prompt for AI Agents
In wren-ai-service/src/providers/llm/litellm.py around lines 97-98 there is a
leftover debug print: print(f"prompt: {openai_formatted_messages}"); remove this
print statement to prevent writing full message history (potential PII) to
stdout on every LLM call. If logging is required for debugging, replace the
print with a logger at a configurable level (e.g., debug) and ensure messages
are redacted or masked (strip PII) before logging, or wrap the log behind a
feature-flag/conditional so it’s not executed in production.
github-actions
bot
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
docker/config.example.yaml(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
| - name: sql_knowledge_retrieval | ||
| engine: wren_ibis | ||
| document_store: qdrant |
There was a problem hiding this comment.
Add missing allow_sql_knowledge_retrieval configuration flag.
Per the PR objectives, a new allow_sql_knowledge_retrieval flag (default: true) controls whether the pipeline retrieves SQL knowledge or falls back to Ibis prompts. This flag is not present in the settings section, but it should be added to align with the feature's design.
Add the configuration flag to the settings section:
settings:
doc_endpoint: https://docs.getwren.ai
is_oss: true
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
+ allow_sql_knowledge_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In docker/config.example.yaml around lines 166 to 168, the settings block is
missing the new allow_sql_knowledge_retrieval boolean flag; add
allow_sql_knowledge_retrieval: true under the settings section (aligning with
defaults) so the pipeline can toggle SQL knowledge retrieval vs Ibis prompts,
and ensure it is documented/commented the same way as nearby flags for clarity.
Description
This update introduces a new mechanism to replace SQL-generation prompts using the SQL Knowledge Retrieval Pipeline, enabling project-aware, schema-aware, and semantic-aware SQL instructions across the system. When enabled, services will dynamically obtain instruction sets from the semantic layer instead of relying solely on prompts provided by the Ibis engine.
Notion : https://www.notion.so/canner/Replace-SQL-prompt-with-v3-connector-data_source-knowledge-from-engine-2ae92976d0bf809d92d2dd2379984113?source=copy_link
Key Changes
1. Dynamic Prompt Replacement via Knowledge Pipeline
File updated:
src/pipelines/generation/utils/sql.pyThe following fields will now be populated using SQL knowledge retrieved from the pipeline:
These previously came from the Ibis engine. They can now be replaced by semantic-layer knowledge when enabled.
2. Integration Across Core Services
SQL knowledge retrieval is now integrated into the following services:
askask_feedbackquestion_recommendationsql_correctionsThese services can now use dynamically generated SQL instructions that reflect project modeling, calculated fields, metrics, JSON extraction rules, and semantic relationships.
3. Configuration Flag
A new setting enables or disables SQL knowledge retrieval:
Default: enabled
🆕 New Features
Chores
Summary by CodeRabbit
New Features
Chores