fix: wrong key to get the recommendation

Canner · Oct 7, 2024 · c8ea912 · c8ea912
1 parent 88c778d
commit c8ea912
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py
@@ -61,7 +61,7 @@ async def recommend(self, request: Input, **kwargs) -> Resource:
             resp = await self._pipelines["relationship_recommendation"].run(**input)
 
             self._cache[request.id] = self.Resource(
-                id=request.id, status="finished", response=resp.get("recommendations")
+                id=request.id, status="finished", response=resp.get("validated")
             )
         except orjson.JSONDecodeError as e:
             self._handle_exception(request, f"Failed to parse MDL: {str(e)}")

diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py
@@ -19,7 +19,7 @@ def relationship_recommendation_service(mock_pipeline):
 @pytest.mark.asyncio
 async def test_recommend_success(relationship_recommendation_service, mock_pipeline):
     request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}')
-    mock_pipeline.run.return_value = {"recommendations": {"test": "data"}}
+    mock_pipeline.run.return_value = {"validated": {"test": "data"}}
 
     response = await relationship_recommendation_service.recommend(request)