chore(wren-ai-service): fix relationship recommendation #1895
Conversation
WalkthroughUpdated relationship validation to require a model catalog input. The validated function now accepts mdl: dict, builds a set of non-relationship columns per model, and filters relationships to include only those with allowed relation types and existing models/columns. Return shape remains {"relationships": [...] }. Changes
Sequence Diagram(s)sequenceDiagram participant Client participant RelationshipValidator as relationship_recommendation.validated participant ModelCatalog as mdl Client->>RelationshipValidator: validated(normalized, mdl) RelationshipValidator->>ModelCatalog: Read models and columns RelationshipValidator->>RelationshipValidator: Build model_columns map loop For each candidate relationship RelationshipValidator->>RelationshipValidator: Check relation type is include RelationshipValidator->>RelationshipValidator: Verify fromModel/toModel exist RelationshipValidator->>RelationshipValidator: Verify fromColumn/toColumn exist end RelationshipValidator-->>Client: {"relationships": validated_list} Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ 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. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File ( |
cyyeh added module/ai-service 
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (3)
137-147: Align column filtering with cleaned_models; guard against None and simplify set constructionUse the same predicate as cleaned_models for detecting non-relationship columns, handle potential None columns safely, skip nameless models, and simplify with set comprehensions.
-def validated(normalized: dict, mdl: dict) -> dict: - model_columns = { - model["name"]: set( - [ - column["name"] - for column in model.get("columns", []) - if not column.get("relationship") - ] - ) - for model in mdl.get("models", []) - } +def validated(normalized: dict, mdl: dict) -> dict: + model_columns = { + model.get("name"): { + column.get("name") + for column in (model.get("columns") or []) + if column.get("name") and ("relationship" not in column) + } + for model in mdl.get("models", []) + if model.get("name") + }
149-151: Defensively normalize relationships from LLM outputThe LLM may return an unexpected shape; type-guard to avoid runtime issues.
- relationships = normalized.get("relationships", []) + relationships = normalized.get("relationships") or [] + if not isinstance(relationships, list): + relationships = []
154-159: Exclude same-model relationships and avoid None in membership checksEnforce the “no intra-model relationships” rule and make membership checks robust with a default empty set.
- and relationship.get("fromModel") in model_columns - and relationship.get("toModel") in model_columns - and relationship.get("fromColumn") - in model_columns.get(relationship.get("fromModel")) - and relationship.get("toColumn") - in model_columns.get(relationship.get("toModel")) + and relationship.get("fromModel") in model_columns + and relationship.get("toModel") in model_columns + and relationship.get("fromModel") != relationship.get("toModel") + and relationship.get("fromColumn") in model_columns.get(relationship.get("fromModel"), set()) + and relationship.get("toColumn") in model_columns.get(relationship.get("toModel"), set())
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (1)
wren-ai-service/src/web/v1/routers/semantics_description.py (1)
get(67-98)
⏰ 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: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (1)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (1)
137-137: Signature change — no direct callers foundVerified: the only occurrence of validated(...) is its definition at wren-ai-service/src/pipelines/generation/relationship_recommendation.py:137. No other call sites or rebindings were found in the repository.
- Location: wren-ai-service/src/pipelines/generation/relationship_recommendation.py:137
Summary by CodeRabbit
Bug Fixes
Refactor