Merge pull request #1248 from max-svistunov/lcore-1376-translate-byok-vector-store-ids

LCORE-1376 RAG: resolve BYOK rag_ids to llama-stack vector_db_ids in query path

Author: tisnik

src/utils/responses.py

@@ -34,6 +34,7 @@
 import metrics
 from configuration import configuration
 from constants import DEFAULT_RAG_TOOL
+from models.config import ByokRag
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
 from models.responses import (
@@ -142,6 +143,11 @@ async def prepare_tools(  # pylint: disable=too-many-arguments,too-many-position
         except APIStatusError as e:
             error_response = InternalServerErrorResponse.generic()
             raise HTTPException(**error_response.model_dump()) from e
+    else:
+        # Translate customer-facing BYOK rag_ids to llama-stack vector_db_ids
+        vector_store_ids = resolve_vector_store_ids(
+            vector_store_ids, configuration.configuration.byok_rag
+        )
 
     # Add RAG tools if vector stores are available
     rag_tools = get_rag_tools(vector_store_ids)
@@ -309,6 +315,28 @@ def extract_vector_store_ids_from_tools(
     return vector_store_ids
 
 
+def resolve_vector_store_ids(
+    vector_store_ids: list[str], byok_rags: list[ByokRag]
+) -> list[str]:
+    """Translate customer-facing BYOK rag_ids to llama-stack vector_db_ids.
+
+    Each ID is looked up against the BYOK RAG configuration. If a matching
+    ``rag_id`` is found, the corresponding ``vector_db_id`` is returned.
+    Otherwise the ID is passed through unchanged (assumed to already be a
+    llama-stack vector store ID).
+
+    Parameters:
+        vector_store_ids: List of IDs from the client request (may be
+            customer-facing rag_ids or raw llama-stack vector_db_ids).
+        byok_rags: BYOK RAG configuration entries.
+
+    Returns:
+        List of llama-stack vector_db_ids ready for the Llama Stack API.
+    """
+    rag_id_to_vector_db_id = {brag.rag_id: brag.vector_db_id for brag in byok_rags}
+    return [rag_id_to_vector_db_id.get(vs_id, vs_id) for vs_id in vector_store_ids]
+
+
 def get_rag_tools(vector_store_ids: list[str]) -> Optional[list[InputToolFileSearch]]:
     """Convert vector store IDs to tools format for Responses API.
 

tests/unit/utils/test_responses.py

@@ -23,7 +23,7 @@
 from pydantic import AnyUrl
 from pytest_mock import MockerFixture
 
-from models.config import ModelContextProtocolServer
+from models.config import ByokRag, ModelContextProtocolServer
 from models.requests import QueryRequest
 from utils.responses import (
     build_mcp_tool_call_from_arguments_done,
@@ -41,6 +41,7 @@
     parse_referenced_documents,
     prepare_responses_params,
     prepare_tools,
+    resolve_vector_store_ids,
     _build_chunk_attributes,
     _increment_llm_call_metric,
     _resolve_source_for_result,
@@ -1000,6 +1001,132 @@ async def test_prepare_tools_no_tools_true(self, mocker: MockerFixture) -> None:
         mock_client.vector_stores.list.assert_not_called()
 
 
+class TestResolveVectorStoreIds:
+    """Tests for resolve_vector_store_ids function."""
+
+    @staticmethod
+    def _make_byok_rag(rag_id: str, vector_db_id: str) -> ByokRag:
+        """Create a ByokRag instance for testing."""
+        return ByokRag(
+            rag_id=rag_id,
+            vector_db_id=vector_db_id,
+            db_path="tests/configuration/rag.txt",
+        )
+
+    def test_translates_customer_facing_ids_to_internal(self) -> None:
+        """Test that customer-facing rag_ids are translated to vector_db_ids."""
+        byok_rags = [
+            self._make_byok_rag("ocp_docs", "vs-001"),
+            self._make_byok_rag("knowledge_base", "vs-002"),
+        ]
+        result = resolve_vector_store_ids(["ocp_docs", "knowledge_base"], byok_rags)
+        assert result == ["vs-001", "vs-002"]
+
+    def test_passes_through_unknown_ids(self) -> None:
+        """Test that IDs not matching any rag_id are passed through unchanged."""
+        byok_rags = [self._make_byok_rag("ocp_docs", "vs-001")]
+        result = resolve_vector_store_ids(["unknown-id"], byok_rags)
+        assert result == ["unknown-id"]
+
+    def test_mixed_known_and_unknown_ids(self) -> None:
+        """Test mix of customer-facing IDs and raw llama-stack IDs."""
+        byok_rags = [self._make_byok_rag("ocp_docs", "vs-001")]
+        result = resolve_vector_store_ids(["ocp_docs", "already-internal"], byok_rags)
+        assert result == ["vs-001", "already-internal"]
+
+    def test_empty_vector_store_ids(self) -> None:
+        """Test that empty input returns empty output."""
+        byok_rags = [self._make_byok_rag("ocp_docs", "vs-001")]
+        result = resolve_vector_store_ids([], byok_rags)
+        assert result == []
+
+    def test_empty_byok_rags(self) -> None:
+        """Test that all IDs pass through when no BYOK RAGs are configured."""
+        result = resolve_vector_store_ids(["vs-001", "vs-002"], [])
+        assert result == ["vs-001", "vs-002"]
+
+    def test_preserves_order(self) -> None:
+        """Test that output order matches input order."""
+        byok_rags = [
+            self._make_byok_rag("b_rag", "vs-b"),
+            self._make_byok_rag("a_rag", "vs-a"),
+        ]
+        result = resolve_vector_store_ids(["a_rag", "b_rag"], byok_rags)
+        assert result == ["vs-a", "vs-b"]
+
+
+class TestPrepareToolsTranslatesVectorStoreIds:
+    """Tests that prepare_tools translates BYOK IDs before building RAG tools."""
+
+    @pytest.mark.asyncio
+    async def test_translates_byok_ids_in_prepare_tools(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test that prepare_tools translates customer-facing IDs to internal IDs."""
+        mock_client = mocker.AsyncMock()
+        mocker.patch("utils.responses.get_mcp_tools", return_value=None)
+
+        # Configure BYOK RAG mapping
+        mock_byok_rag = mocker.Mock()
+        mock_byok_rag.rag_id = "ocp_docs"
+        mock_byok_rag.vector_db_id = "vs-001"
+        mock_config = mocker.Mock()
+        mock_config.configuration.byok_rag = [mock_byok_rag]
+        mocker.patch("utils.responses.configuration", mock_config)
+
+        result = await prepare_tools(mock_client, ["ocp_docs"], False, "token")
+        assert result is not None
+        assert len(result) == 1
+        assert result[0].type == "file_search"
+        assert result[0].vector_store_ids == ["vs-001"]
+
+    @pytest.mark.asyncio
+    async def test_passes_through_unknown_ids_in_prepare_tools(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test that prepare_tools passes through IDs not in BYOK config."""
+        mock_client = mocker.AsyncMock()
+        mocker.patch("utils.responses.get_mcp_tools", return_value=None)
+
+        # Configure empty BYOK RAG
+        mock_config = mocker.Mock()
+        mock_config.configuration.byok_rag = []
+        mocker.patch("utils.responses.configuration", mock_config)
+
+        result = await prepare_tools(mock_client, ["raw-internal-id"], False, "token")
+        assert result is not None
+        assert result[0].vector_store_ids == ["raw-internal-id"]
+
+    @pytest.mark.asyncio
+    async def test_does_not_translate_when_ids_fetched_from_llama_stack(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test that IDs fetched from llama-stack (None path) are not translated."""
+        mock_client = mocker.AsyncMock()
+        mock_vector_store = mocker.Mock()
+        mock_vector_store.id = "vs-internal"
+        mock_vector_stores = mocker.Mock()
+        mock_vector_stores.data = [mock_vector_store]
+        mock_client.vector_stores.list = mocker.AsyncMock(
+            return_value=mock_vector_stores
+        )
+        mocker.patch("utils.responses.get_mcp_tools", return_value=None)
+
+        # Configure BYOK RAG whose rag_id matches the fetched ID so that
+        # accidental translation would change the result and fail the assertion
+        mock_byok_rag = mocker.Mock()
+        mock_byok_rag.rag_id = "vs-internal"
+        mock_byok_rag.vector_db_id = "vs-translated"
+        mock_config = mocker.Mock()
+        mock_config.configuration.byok_rag = [mock_byok_rag]
+        mocker.patch("utils.responses.configuration", mock_config)
+
+        result = await prepare_tools(mock_client, None, False, "token")
+        assert result is not None
+        # The IDs from llama-stack should be used as-is (no BYOK translation on None path)
+        assert result[0].vector_store_ids == ["vs-internal"]
+
+
 class TestPrepareResponsesParams:
     """Tests for prepare_responses_params function."""