diff --git a/src/utils/responses.py b/src/utils/responses.py index bdced0f71..9fd9ceef9 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -842,6 +842,16 @@ def _resolve_source_for_result( if len(vector_store_ids) > 1: attributes = getattr(result, "attributes", {}) or {} + + # Primary: read index name embedded directly by rag-content. + # This value is already the user-facing rag_id, not a vector_db_id, + # so no mapping is needed. + attr_source: Optional[str] = attributes.get("source") + if attr_source: + return attr_source + + # Fallback: if llama-stack ever populates vector_store_id in results, + # use it with the rag_id_mapping. attr_store_id: Optional[str] = attributes.get("vector_store_id") if attr_store_id: return rag_id_mapping.get(attr_store_id, attr_store_id) diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index 54ae158aa..f9b5f3a9e 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -2044,6 +2044,69 @@ def test_multiple_stores_attribute_not_in_mapping( ) assert source == "vs-unknown" + def test_multiple_stores_source_attribute_fallback( + self, mocker: MockerFixture + ) -> None: + """Test resolution falls back to source attribute when no vector_store_id.""" + mock_result = mocker.Mock() + mock_result.filename = "file-abc123" + mock_result.attributes = {"source": "ocp-documentation"} + + source = _resolve_source_for_result( + mock_result, + ["vs-001", "vs-002"], + {"vs-001": "ocp-4.18-docs"}, + ) + assert source == "ocp-documentation" + + def test_multiple_stores_source_attribute_ignores_mapping( + self, mocker: MockerFixture + ) -> None: + """Test source attribute is returned directly without rag_id_mapping lookup.""" + mock_result = mocker.Mock() + mock_result.filename = "file-abc123" + mock_result.attributes = {"source": "custom-index"} + + source = _resolve_source_for_result( + mock_result, + ["vs-001", "vs-002"], + {"custom-index": "should-not-be-used"}, + ) + assert source == "custom-index" + + def test_multiple_stores_source_preferred_over_vector_store_id( + self, mocker: MockerFixture + ) -> None: + """Test source attribute takes precedence over vector_store_id.""" + mock_result = mocker.Mock() + mock_result.filename = "file-abc123" + mock_result.attributes = { + "vector_store_id": "vs-002", + "source": "ocp-documentation", + } + + source = _resolve_source_for_result( + mock_result, + ["vs-001", "vs-002"], + {"vs-002": "rhel-9-docs"}, + ) + assert source == "ocp-documentation" + + def test_multiple_stores_no_vector_store_id_no_source( + self, mocker: MockerFixture + ) -> None: + """Test resolution returns None when neither vector_store_id nor source present.""" + mock_result = mocker.Mock() + mock_result.filename = "file-abc123" + mock_result.attributes = {"title": "some doc"} + + source = _resolve_source_for_result( + mock_result, + ["vs-001", "vs-002"], + {"vs-001": "ocp-docs"}, + ) + assert source is None + class TestBuildChunkAttributes: """Tests for _build_chunk_attributes function."""