Improve source fallback semantics and use keyword arguments

Author: max-svistunov

src/app/endpoints/streaming_query.py

@@ -532,8 +532,8 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
                 tool_call, tool_result = build_tool_call_summary(
                     output_item_done_chunk.item,
                     turn_summary.rag_chunks,
-                    context.vector_store_ids,
-                    context.rag_id_mapping,
+                    vector_store_ids=context.vector_store_ids,
+                    rag_id_mapping=context.rag_id_mapping,
                 )
                 if tool_call:
                     turn_summary.tool_calls.append(tool_call)

src/utils/responses.py

@@ -483,11 +483,12 @@ def _resolve_single_store_source(
         rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
 
     Returns:
-        The resolved rag_id if exactly one store is used, None otherwise.
+        The resolved rag_id (or raw store_id as fallback) if exactly one
+        store is used, None otherwise.
     """
     if len(vector_store_ids) == 1:
         store_id = vector_store_ids[0]
-        return rag_id_mapping.get(store_id)
+        return rag_id_mapping.get(store_id, store_id)
     return None
 
 
@@ -807,15 +808,15 @@ def _resolve_source_for_result(
     """
     if len(vector_store_ids) == 1:
         store_id = vector_store_ids[0]
-        return rag_id_mapping.get(store_id, result.filename)
+        return rag_id_mapping.get(store_id, store_id)
 
     if len(vector_store_ids) > 1:
         attributes = getattr(result, "attributes", {}) or {}
         attr_store_id: Optional[str] = attributes.get("vector_store_id")
-        if attr_store_id and attr_store_id in rag_id_mapping:
-            return rag_id_mapping[attr_store_id]
+        if attr_store_id:
+            return rag_id_mapping.get(attr_store_id, attr_store_id)
 
-    return result.filename
+    return None
 
 
 def _build_chunk_attributes(result: Any) -> Optional[dict[str, Any]]:

tests/unit/utils/test_responses.py

@@ -1387,7 +1387,7 @@ def test_extract_rag_chunks_with_results(self, mocker: MockerFixture) -> None:
 
         assert len(rag_chunks) == 2
         assert rag_chunks[0].content == "chunk 1"
-        assert rag_chunks[0].source == "doc1.pdf"
+        assert rag_chunks[0].source is None
         assert rag_chunks[0].score == 0.9
 
     def test_extract_rag_chunks_no_results(self, mocker: MockerFixture) -> None:
@@ -1596,12 +1596,12 @@ def test_single_store_mapped(self, mocker: MockerFixture) -> None:
         assert source == "ocp-4.18-docs"
 
     def test_single_store_unmapped(self, mocker: MockerFixture) -> None:
-        """Test resolution with single vector store without mapping falls back to filename."""
+        """Test resolution with single vector store without mapping returns raw store ID."""
         mock_result = mocker.Mock()
         mock_result.filename = "file-abc123"
 
         source = _resolve_source_for_result(mock_result, ["vs-unknown"], {})
-        assert source == "file-abc123"
+        assert source == "vs-unknown"
 
     def test_multiple_stores_with_attribute(self, mocker: MockerFixture) -> None:
         """Test resolution with multiple stores using result attributes."""
@@ -1617,7 +1617,7 @@ def test_multiple_stores_with_attribute(self, mocker: MockerFixture) -> None:
         assert source == "rhel-9-docs"
 
     def test_multiple_stores_no_attribute(self, mocker: MockerFixture) -> None:
-        """Test resolution with multiple stores and no vector_store_id attribute."""
+        """Test resolution with multiple stores and no vector_store_id attribute returns None."""
         mock_result = mocker.Mock()
         mock_result.filename = "file-abc123"
         mock_result.attributes = {}
@@ -1627,20 +1627,20 @@ def test_multiple_stores_no_attribute(self, mocker: MockerFixture) -> None:
             ["vs-001", "vs-002"],
             {"vs-001": "ocp-4.18-docs", "vs-002": "rhel-9-docs"},
         )
-        assert source == "file-abc123"
+        assert source is None
 
     def test_no_stores(self, mocker: MockerFixture) -> None:
-        """Test resolution with no vector stores falls back to filename."""
+        """Test resolution with no vector stores returns None."""
         mock_result = mocker.Mock()
         mock_result.filename = "file-abc123"
 
         source = _resolve_source_for_result(mock_result, [], {})
-        assert source == "file-abc123"
+        assert source is None
 
     def test_multiple_stores_attribute_not_in_mapping(
         self, mocker: MockerFixture
     ) -> None:
-        """Test resolution when attribute store ID is not in mapping."""
+        """Test resolution when attribute store ID is not in mapping returns raw store ID."""
         mock_result = mocker.Mock()
         mock_result.filename = "file-abc123"
         mock_result.attributes = {"vector_store_id": "vs-unknown"}
@@ -1650,7 +1650,7 @@ def test_multiple_stores_attribute_not_in_mapping(
             ["vs-001", "vs-002"],
             {"vs-001": "ocp-docs"},
         )
-        assert source == "file-abc123"
+        assert source == "vs-unknown"
 
 
 class TestBuildChunkAttributes:
@@ -1759,7 +1759,7 @@ def test_chunks_no_mapping_falls_back(self, mocker: MockerFixture) -> None:
         extract_rag_chunks_from_file_search_item(mock_item, rag_chunks)
 
         assert len(rag_chunks) == 1
-        assert rag_chunks[0].source == "file-abc"
+        assert rag_chunks[0].source is None
         assert rag_chunks[0].attributes is None
 
     def test_chunks_multiple_stores_attribute_resolution(
@@ -1853,7 +1853,7 @@ def test_file_search_without_mapping(self, mocker: MockerFixture) -> None:
         call_summary, _ = build_tool_call_summary(mock_item, rag_chunks)
 
         assert len(rag_chunks) == 1
-        assert rag_chunks[0].source == "doc.pdf"
+        assert rag_chunks[0].source is None
         assert rag_chunks[0].attributes is None
         assert call_summary is not None