feat(wren-ai-service): add custom instructions #1756
Conversation
cyyeh
added
module/ai-service
WalkthroughSupport for an optional Changes
Suggested labels
Suggested reviewers
Poem
✨ 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: 1
♻️ Duplicate comments (1)
wren-ai-service/src/web/v1/services/sql_answer.py (1)
21-21: Consider adding validation for custom instructions.Similar to the chart service, the
custom_instructionfield accepts arbitrary user input without validation. This could potentially be exploited for prompt injection. Consider implementing the same validation approach across both services.Apply similar validation as suggested for the chart service:
+ from pydantic import Field, validator + class SqlAnswerRequest(BaseRequest): query: str sql: str sql_data: Dict - custom_instruction: Optional[str] = None + custom_instruction: Optional[str] = Field(None, max_length=1000) + + @validator('custom_instruction') + def validate_custom_instruction(cls, v): + if v is not None and len(v.strip()) == 0: + return None + return v
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
wren-ai-service/src/pipelines/generation/chart_generation.py(6 hunks)wren-ai-service/src/pipelines/generation/sql_answer.py(5 hunks)wren-ai-service/src/web/v1/services/chart.py(2 hunks)wren-ai-service/src/web/v1/services/sql_answer.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (10)
wren-ai-service/src/web/v1/services/chart.py (1)
128-128: LGTM: Custom instruction properly passed to pipeline.The parameter is correctly propagated to the chart generation pipeline.
wren-ai-service/src/web/v1/services/sql_answer.py (1)
103-103: LGTM: Custom instruction properly passed to pipeline.The parameter is correctly propagated to the SQL answer pipeline.
wren-ai-service/src/pipelines/generation/chart_generation.py (4)
30-30: LGTM: Appropriate system prompt instruction for custom instructions.The system prompt correctly instructs the model to strictly follow custom instructions to modify response style.
50-50: LGTM: User prompt template properly includes custom instruction field.The template correctly includes the custom_instruction field for prompt generation.
70-70: LGTM: Function signature and parameter passing implemented correctly.The
promptfunction signature is properly updated to acceptcustom_instructionand passes it to the prompt builder.Also applies to: 82-82
158-158: LGTM: Pipeline run method properly handles custom instruction parameter.The
runmethod correctly accepts the optionalcustom_instructionparameter and defaults it to an empty string, preventing None values in the prompt context.Also applies to: 169-169
wren-ai-service/src/pipelines/generation/sql_answer.py (4)
32-32: LGTM: Consistent system prompt instruction for custom instructions.The system prompt correctly mirrors the chart generation approach, instructing the model to strictly follow custom instructions.
45-45: LGTM: User prompt template consistent with chart generation.The template correctly includes the custom_instruction field, maintaining consistency across pipelines.
58-58: LGTM: Function signature and parameter passing implemented correctly.The
promptfunction signature is properly updated and parameter passing is consistent with the chart generation pipeline.Also applies to: 66-66
148-148: LGTM: Pipeline run method properly handles custom instruction parameter.The
runmethod implementation is consistent with the chart generation pipeline, correctly handling the optional parameter with appropriate defaulting.Also applies to: 159-159
| @@ -18,6 +18,7 @@ class ChartRequest(BaseRequest): | |||
| sql: str | |||
| data: Optional[Dict[str, Any]] = None | |||
| remove_data_from_chart_schema: Optional[bool] = True | |||
| custom_instruction: Optional[str] = None | |||
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider adding validation for custom instructions.
The custom_instruction field accepts arbitrary user input without validation, which could potentially be used for prompt injection attacks. Consider adding length limits, content validation, or sanitization to prevent misuse.
+ from pydantic import Field, validator
+
class ChartRequest(BaseRequest):
query: str
sql: str
data: Optional[Dict[str, Any]] = None
remove_data_from_chart_schema: Optional[bool] = True
- custom_instruction: Optional[str] = None
+ custom_instruction: Optional[str] = Field(None, max_length=1000)
+
+ @validator('custom_instruction')
+ def validate_custom_instruction(cls, v):
+ if v is not None and len(v.strip()) == 0:
+ return None
+ return v📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| custom_instruction: Optional[str] = None | |
| from pydantic import Field, validator | |
| class ChartRequest(BaseRequest): | |
| query: str | |
| sql: str | |
| data: Optional[Dict[str, Any]] = None | |
| remove_data_from_chart_schema: Optional[bool] = True | |
| custom_instruction: Optional[str] = Field(None, max_length=1000) | |
| @validator('custom_instruction') | |
| def validate_custom_instruction(cls, v): | |
| if v is not None and len(v.strip()) == 0: | |
| return None | |
| return v |
🤖 Prompt for AI Agents
In wren-ai-service/src/web/v1/services/chart.py at line 21, the
custom_instruction field currently accepts any user input without validation,
posing a security risk. Add validation logic to enforce length limits and
sanitize or filter the content to prevent prompt injection attacks. Implement
checks that reject or clean inputs containing suspicious or disallowed patterns
before processing or storing the custom_instruction value.
Summary by CodeRabbit