chore(wren-ai-service): minor updates (ai-env-changed) #1809
Conversation
cyyeh
added
module/ai-service
WalkthroughThis update standardizes type annotations throughout the codebase, removing the use of Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant WebAPI
participant Pipelines
participant MetadataStore
Client->>WebAPI: POST /pipelines/{pipeline_name} (with kwargs)
WebAPI->>Pipelines: pipeline.run(**kwargs)
alt use_dry_plan enabled
Pipelines->>MetadataStore: get_project_metadata(project_id)
MetadataStore-->>Pipelines: metadata dict
Pipelines->>Pipelines: extract data_source from metadata
else use_dry_plan disabled
Pipelines->>Pipelines: use empty metadata
end
Pipelines-->>WebAPI: result
WebAPI-->>Client: result
Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (3)
✅ Files skipped from review due to trivial changes (3)
⏰ Context from checks skipped due to timeout of 90000ms (3)
✨ Finishing Touches
🧪 Generate Unit Tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (1)
wren-ai-service/src/web/development.py (1)
67-92: Refactor to reduce duplication and add parameter validation.This endpoint has several issues:
- Duplicates the pipeline extraction logic from the GET endpoint
- No validation of request parameters against expected pipeline parameters
- Error messages could expose internal implementation details
Extract the pipeline collection logic and add parameter validation:
+def _get_all_pipelines(service_container) -> Dict[str, Any]: + """Extract all pipelines from the service container.""" + pipelines = {} + for attr_name, service in service_container.__dict__.items(): + if attr_name.startswith("_"): + continue + if hasattr(service, "_pipelines"): + pipelines.update(service._pipelines) + return pipelines + @router.post("/pipelines/{pipeline_name}") async def run_pipeline(pipeline_name: str, request_body: Dict[str, Any]) -> dict: + if os.getenv("ENVIRONMENT", "production") == "production": + raise HTTPException(status_code=403, detail="This endpoint is only available in development") + service_container = app.state.service_container - pipe_components = {} - - for _, service in service_container.__dict__.items(): - if hasattr(service, "_pipelines"): - for _pipeline_name, pipeline_instance in service._pipelines.items(): - pipe_components[_pipeline_name] = pipeline_instance + pipe_components = _get_all_pipelines(service_container) if pipeline_name not in pipe_components: logger.error(f"Pipeline {pipeline_name} not found") raise HTTPException( status_code=404, detail=f"Pipeline '{pipeline_name}' not found" ) try: pipeline = pipe_components[pipeline_name] + + # Validate parameters + expected_params = _extract_run_method_params(pipeline) + unexpected_params = set(request_body.keys()) - set(expected_params.keys()) + if unexpected_params: + logger.warning(f"Unexpected parameters for pipeline '{pipeline_name}': {unexpected_params}") + result = await pipeline.run(**request_body) return result except Exception as e: logger.error(f"Error running pipeline {pipeline_name}: {e}") raise HTTPException( - status_code=500, detail=f"Error running pipeline '{pipeline_name}': {e}" + status_code=500, detail=f"Error running pipeline '{pipeline_name}': {type(e).__name__}" )Also update the GET endpoint to use the new helper function:
@router.get("/pipelines") async def get_pipelines() -> dict: + if os.getenv("ENVIRONMENT", "production") == "production": + raise HTTPException(status_code=403, detail="This endpoint is only available in development") + service_container = app.state.service_container - pipeline_params = {} - - # Extract pipelines from all services - for _, service in service_container.__dict__.items(): - if hasattr(service, "_pipelines"): - for pipeline_name, pipeline_instance in service._pipelines.items(): - pipeline_params[pipeline_name] = _extract_run_method_params( - pipeline_instance - ) + pipelines = _get_all_pipelines(service_container) + + pipeline_params = { + name: _extract_run_method_params(pipeline) + for name, pipeline in pipelines.items() + } return pipeline_params
🧹 Nitpick comments (1)
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
188-217: Validate project_id when dry plan is enabled.The implementation looks good, but consider validating
project_idwhenuse_dry_planis True to avoid potential issues with metadata retrieval.logger.info("SQL Generation pipeline is running...") if use_dry_plan: + if not project_id: + logger.warning("use_dry_plan is True but project_id is not provided, metadata retrieval may fail") metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {}
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (24)
wren-ai-service/src/pipelines/generation/chart_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(6 hunks)wren-ai-service/src/pipelines/generation/relationship_recommendation.py(2 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(6 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(6 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(2 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(3 hunks)wren-ai-service/src/pipelines/indexing/db_schema.py(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(4 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(6 hunks)wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py(1 hunks)wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/src/pipelines/retrieval/sql_executor.py(1 hunks)wren-ai-service/src/pipelines/retrieval/sql_functions.py(1 hunks)wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py(1 hunks)wren-ai-service/src/providers/embedder/litellm.py(1 hunks)wren-ai-service/src/web/development.py(1 hunks)wren-ai-service/src/web/v1/routers/question_recommendation.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(6 hunks)wren-ai-service/src/web/v1/services/chart.py(1 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(2 hunks)wren-ai-service/tools/dev/.env(1 hunks)
🧰 Additional context used
🧠 Learnings (17)
wren-ai-service/tools/dev/.env (2)
Learnt from: paopa
PR: Canner/WrenAI#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: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/web/v1/services/chart.py (3)
Learnt from: andreashimin
PR: Canner/WrenAI#1117
File: wren-ui/src/apollo/server/backgrounds/chart.ts:86-86
Timestamp: 2025-01-15T03:50:47.868Z
Learning: In the WrenAI codebase, ChartType is maintained as a string in the web database, while the ChartType enum is specifically used for AI adaptor input/output interfaces. Therefore, string case conversions like toUpperCase() are intentional for maintaining data consistency in the web DB.
Learnt from: andreashimin
PR: Canner/WrenAI#1117
File: wren-ui/src/apollo/server/repositories/threadResponseRepository.ts:37-37
Timestamp: 2025-01-15T03:51:14.432Z
Learning: In the ThreadResponseRepository, the `chartType` property is intentionally typed as string in the database layer, while the ChartType enum is used specifically for AI adaptor input/output in the application layer.
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/web/v1/services/__init__.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/generation/chart_generation.py (1)
Learnt from: andreashimin
PR: Canner/WrenAI#1117
File: wren-ui/src/apollo/server/backgrounds/chart.ts:86-86
Timestamp: 2025-01-15T03:50:47.868Z
Learning: In the WrenAI codebase, ChartType is maintained as a string in the web database, while the ChartType enum is specifically used for AI adaptor input/output interfaces. Therefore, string case conversions like toUpperCase() are intentional for maintaining data consistency in the web DB.
wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/indexing/instructions.py (2)
Learnt from: paopa
PR: Canner/WrenAI#1376
File: wren-ai-service/src/pipelines/indexing/instructions.py:100-107
Timestamp: 2025-03-13T03:48:24.664Z
Learning: The `InstructionsCleaner.run()` method in the instructions indexing pipeline only accepts `instruction_ids` and `project_id` parameters, and doesn't support a `delete_all` parameter for unconditional deletion.
Learnt from: onlyjackfrost
PR: Canner/WrenAI#1392
File: wren-ui/src/apollo/server/services/instructionService.ts:107-146
Timestamp: 2025-03-17T09:16:53.878Z
Learning: In the WrenAI project, the instructionResolver merges the instruction ID (from args.where) and project ID (from the current project context) with instruction data before passing to the instructionService, allowing the GraphQL schema to have separate input types while the service layer works with complete objects.
wren-ai-service/src/web/v1/routers/question_recommendation.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
Learnt from: andreashimin
PR: Canner/WrenAI#1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:10.593Z
Learning: The typo in error code enum names ("IDENTIED_AS_GENERAL" and "IDENTIED_AS_MISLEADING_QUERY" instead of "IDENTIFIED_AS_GENERAL" and "IDENTIFIED_AS_MISLEADING_QUERY") in wren-ui/src/apollo/server/utils/error.ts was recognized but intentionally not fixed as it was considered out of scope for PR #1414.
wren-ai-service/src/web/v1/services/ask.py (3)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/web/v1/services/sql_question.py:17-28
Timestamp: 2025-01-14T07:45:32.410Z
Learning: In the SqlQuestionRequest model (Python/FastAPI), the query_id field is implemented using a protected _query_id attribute with getter/setter properties, rather than a direct field, as per the team's preference.
Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/web/v1/services/sql_question.py:101-109
Timestamp: 2025-01-14T07:45:56.117Z
Learning: Warning about accessing `query_id` during exception handling in SQL question service can be skipped as per team's decision. This applies to the pattern where `sql_question_request.query_id` is accessed within the exception handling block.
wren-ai-service/src/pipelines/generation/sql_regeneration.py (2)
Learnt from: wwwy3y3
PR: Canner/WrenAI#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.
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/web/v1/services/question_recommendation.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/providers/embedder/litellm.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/web/development.py (1)
Learnt from: wwwy3y3
PR: Canner/WrenAI#1388
File: wren-ui/src/apollo/server/services/askingService.ts:599-604
Timestamp: 2025-03-17T19:26:10.835Z
Learning: Error handling for non-existent asking tasks in WrenAI is handled at the resolver level rather than in the service implementation.
wren-ai-service/src/pipelines/generation/sql_correction.py (3)
Learnt from: andreashimin
PR: Canner/WrenAI#1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:10.593Z
Learning: The typo in error code enum names ("IDENTIED_AS_GENERAL" and "IDENTIED_AS_MISLEADING_QUERY" instead of "IDENTIFIED_AS_GENERAL" and "IDENTIFIED_AS_MISLEADING_QUERY") in wren-ui/src/apollo/server/utils/error.ts was recognized but intentionally not fixed as it was considered out of scope for PR #1414.
Learnt from: andreashimin
PR: Canner/WrenAI#1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:26.608Z
Learning: The enum names `IDENTIED_AS_GENERAL` and `IDENTIED_AS_MISLEADING_QUERY` in GeneralErrorCodes enum in wren-ui/src/apollo/server/utils/error.ts have typos (missing "FI" in "IDENTIFIED") but fixing them is not in scope for the current PR.
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (2)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
Learnt from: cyyeh
PR: Canner/WrenAI#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.
🧬 Code Graph Analysis (2)
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
wren-ai-service/src/pipelines/common.py (1)
retrieve_metadata(87-105)
wren-ai-service/src/pipelines/indexing/sql_pairs.py (4)
wren-ai-service/src/pipelines/indexing/instructions.py (3)
run(31-52)run(61-76)run(149-164)wren-ai-service/src/pipelines/retrieval/instructions.py (2)
run(23-34)run(177-187)wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py (2)
run(23-33)run(146-156)wren-ai-service/src/pipelines/indexing/__init__.py (3)
run(25-47)run(57-72)run(78-87)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
- GitHub Check: Analyze (javascript-typescript)
🔇 Additional comments (41)
wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)
440-441: LGTM! Type annotations improved for better clarity.The removal of
Optionaltyping from parameters with default values is a good practice. These parameters will never beNonein practice, so the updated type annotations (intinstead ofOptional[int]) provide better clarity and type safety.wren-ai-service/src/providers/embedder/litellm.py (1)
168-168: LGTM! Timeout parameter type annotation standardized.The change from
Optional[float]tofloatwith a default value is consistent with the codebase-wide effort to remove optionality from parameters that have defaults and should never beNone.wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py (1)
125-125: LGTM! Similarity threshold parameter properly typed.The change from
Optional[float]tofloatwith a default value of0.9improves type safety and aligns with the PR's goal of standardizing parameter typing by removing optionality where defaults are provided.wren-ai-service/src/web/v1/services/chart.py (1)
20-20: LGTM! Boolean flag properly typed.The change from
Optional[bool]toboolwith a default value ensures the flag will always be a boolean value, eliminating potentialNone-related issues and improving type clarity.wren-ai-service/src/pipelines/retrieval/sql_executor.py (1)
69-69: LGTM! Engine timeout parameter properly typed.The change from
Optional[float]tofloatwith a default value of30.0is consistent with the codebase-wide standardization effort and ensures the timeout parameter will always have a valid float value.wren-ai-service/src/web/v1/services/__init__.py (1)
29-30: LGTM! Type annotation improvement enhances type safety.Removing
Optionalin favor of non-optional types with default values ensures these configuration fields always have valid values and cannot beNone. This improves type safety while maintaining backward compatibility.wren-ai-service/src/pipelines/generation/relationship_recommendation.py (1)
190-190: LGTM! Consistent timeout parameter typing improvement.Changing from
Optional[float]tofloat = 30.0ensures the timeout parameter always has a valid numeric value, preventing potentialNone-related errors and making the timeout behavior more predictable.wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
153-153: LGTM! Consistent with timeout parameter standardization.This change aligns with the broader effort to standardize timeout parameters across pipeline classes, ensuring they always have valid float values instead of potentially being
None.wren-ai-service/src/pipelines/indexing/db_schema.py (1)
339-339: LGTM! Batch size parameter typing improvement.Changing from
Optional[int]toint = 50ensures the batch size parameter always has a valid integer value, preventing potential issues in column batching logic.wren-ai-service/src/web/v1/routers/question_recommendation.py (1)
78-80: LGTM! Request model typing improvements enhance API consistency.Removing
Optionaltypes in favor of non-optional types with sensible defaults (max_questions=5,max_categories=3,regenerate=False) ensures these API parameters always have valid values, improving predictability and reducing the need forNonechecks in downstream processing.wren-ai-service/src/pipelines/generation/chart_generation.py (1)
157-157: LGTM! Type annotation improvement.Removing
Optionalfrom the type annotation is correct since this parameter has a default value and should never beNone. This improves type safety while maintaining the same functionality.wren-ai-service/src/pipelines/retrieval/sql_functions.py (1)
97-98: LGTM! Type annotation improvements.Removing
Optionalfromengine_timeoutandttlparameters is correct since they have default values and should always be concrete numeric values for proper functionality.wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py (1)
121-122: LGTM! Type annotation improvements.Removing
Optionalfrom similarity threshold and retrieval size parameters is appropriate since they have default values and should always be concrete values for proper retrieval functionality.wren-ai-service/src/pipelines/retrieval/instructions.py (1)
153-154: LGTM! Type annotation improvements.Removing
Optionalfromsimilarity_thresholdandtop_kparameters is correct since they have default values and should always be concrete values for proper retrieval filtering.wren-ai-service/src/web/v1/services/question_recommendation.py (2)
25-25: LGTM! Type annotation improvement.Removing
Optionalfrom theresponseattribute is correct since it has a default dictionary value and should never beNone.
159-161: LGTM! Type annotation improvements.Removing
Optionalfrommax_questions,max_categories, andregenerateparameters is appropriate since they all have default values and should always be concrete values for proper request handling.wren-ai-service/src/pipelines/generation/utils/sql.py (3)
2-2: LGTM: Type annotation standardizationRemoving the
Optionalimport aligns with the broader effort to standardize type annotations by using non-optional types with default values instead of optional types.
32-32: LGTM: Consistent timeout parameter handlingThe change from
Optional[float]tofloat = 30.0standardizes timeout handling across SQL-related pipelines and eliminates potential null pointer issues while maintaining the same default behavior.
247-252: LGTM: Valuable SQL correction examplesThe addition of
SQL_CORRECTION_EXAMPLESprovides concrete guidance for SQL correction scenarios, specifically addressing UNION query syntax issues with parentheses. This will enhance the effectiveness of the SQL correction prompt.wren-ai-service/src/web/v1/services/ask.py (4)
28-31: LGTM: Well-designed dry plan configurationThe new boolean fields
use_dry_planandallow_dry_plan_fallbackprovide clear control over dry plan behavior with sensible defaults. Theuse_dry_plan=Falsedefault maintains backward compatibility, whileallow_dry_plan_fallback=Trueprovides a safe fallback mechanism.
83-83: LGTM: Type annotation consistencyChanging
is_followupfromOptional[bool]tobool = Falsealigns with the broader type annotation standardization effort and eliminates potential null handling issues.
213-214: LGTM: Proper parameter extractionThe extraction of
use_dry_planandallow_dry_plan_fallbackfrom the request follows the established pattern for other configuration parameters in this method.
501-502: LGTM: Consistent parameter propagationThe dry plan parameters are correctly passed to all relevant SQL pipelines (
followup_sql_generation,sql_generation, andsql_correction), ensuring consistent behavior across different execution paths.Also applies to: 518-519, 561-562
wren-ai-service/src/pipelines/indexing/instructions.py (1)
31-31: LGTM: Consistent type annotation standardizationThe changes from
Optional[str]tostr = ""forproject_idparameters across multiple functions maintain the same default behavior while improving type safety by eliminating potential null values. This aligns with the broader codebase standardization effort.Also applies to: 86-86, 104-104, 152-152
wren-ai-service/src/pipelines/generation/sql_correction.py (4)
3-3: LGTM: Type annotation standardizationRemoving the
Optionalimport and changingengine_timeoutfromOptional[float]tofloat = 30.0aligns with the codebase standardization effort and eliminates potential null handling issues.Also applies to: 115-115
16-16: LGTM: Enhanced SQL correction guidanceThe import and integration of
SQL_CORRECTION_EXAMPLESinto the system prompt provides concrete examples for SQL correction scenarios, which should improve the LLM's ability to fix syntax errors in SQL queries.Also applies to: 32-39
91-91: LGTM: Proper data source parameter integrationAdding the
data_sourceparameter to thepost_processfunction signature enables the pipeline to handle different data sources appropriately during SQL validation and correction.
153-156: LGTM: Efficient conditional metadata retrievalThe conditional metadata retrieval based on
use_dry_planis implemented efficiently - only fetching metadata when needed for dry plan validation, with proper fallback to empty dict when not required.wren-ai-service/src/pipelines/generation/followup_sql_generation.py (4)
3-3: LGTM: Type annotation consistencyRemoving the
Optionalimport and changingengine_timeouttofloat = 30.0maintains consistency with other SQL pipeline modules and improves type safety.Also applies to: 153-153
12-13: LGTM: Proper dependency integrationThe addition of
DocumentStoreProviderdependency and retriever initialization enables metadata retrieval for dry plan functionality. The implementation follows the established pattern from other SQL generation pipelines.Also applies to: 151-151, 156-158
129-129: LGTM: Consistent parameter handlingThe addition of
data_source,use_dry_plan, andallow_dry_plan_fallbackparameters to thepost_processfunction maintains consistency with the updatedSQLGenPostProcessorinterface.Also applies to: 131-132, 138-140
194-195: LGTM: Complete dry plan integrationThe implementation properly integrates dry plan functionality with conditional metadata retrieval and parameter propagation through the pipeline execution. The pattern matches the implementation in other SQL generation pipelines, ensuring consistent behavior.
Also applies to: 199-202, 218-220
wren-ai-service/src/web/development.py (1)
1-3: LGTM!The imports are appropriate for the introspection functionality.
wren-ai-service/src/pipelines/generation/sql_generation.py (3)
3-13: LGTM!The import changes are appropriate for the dry plan functionality.
120-136: LGTM!The function correctly propagates the dry plan parameters to the post processor.
143-154: LGTM!The constructor changes properly initialize the document store retriever for metadata access.
wren-ai-service/src/pipelines/indexing/sql_pairs.py (5)
31-31: LGTM!The type annotation change is consistent with the codebase standardization.
111-111: LGTM!The type annotation standardization is consistent.
129-129: LGTM!The type annotation change follows the established pattern.
171-171: LGTM!Removing
Optionalfrom a parameter with a non-None default value improves type accuracy.
196-209: Excellent fix for the mutable default argument!These changes improve the code by:
- Standardizing the
project_idtype annotation- Fixing the mutable default argument anti-pattern for
external_pairs- Properly handling the None case with
(external_pairs or {})
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tests/data/config.test.yaml(2 hunks)
🧰 Additional context used
🧠 Learnings (1)
wren-ai-service/tests/data/config.test.yaml (2)
Learnt from: cyyeh
PR: Canner/WrenAI#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.
Learnt from: cyyeh
PR: Canner/WrenAI#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.
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (1)
wren-ai-service/tests/data/config.test.yaml (1)
2-2: Double-check provider IDs
provider: litellm_llmandprovider: litellm_embeddermust exactly match the identifiers registered in the service’s provider registry. A mismatch will make pipeline resolution fail at runtime.Please verify the canonical names in
src/providers/__init__.py(or equivalent).Also applies to: 17-17
github-actions
bot
changed the title
Summary by CodeRabbit
New Features
Improvements
Chores
Refactor