Rebased, solved conflicts

Author: are-ces

src/app/endpoints/streaming_query.py

@@ -56,7 +56,6 @@
     UnauthorizedResponse,
     UnprocessableEntityResponse,
 )
-from utils.types import ReferencedDocument
 from utils.endpoints import (
     check_configuration_loaded,
     validate_and_retrieve_conversation,
@@ -186,29 +185,10 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 
     client = AsyncLlamaStackClientHolder().get_client()
 
-<<<<<<< HEAD
-<<<<<<< HEAD
-    _, _, doc_ids_from_chunks, pre_rag_chunks = await perform_vector_search(
-        client, query_request.query, query_request.solr
-    )
-
-    rag_context = format_rag_context_for_injection(pre_rag_chunks)
-    if rag_context:
-=======
-    # Build RAG context from BYOK and Solr sources
-    rag_context = await build_rag_context(client, query_request, configuration)
-
-    # Inject RAG context into query
-    if rag_context.context_text:
-        # Mutate a local copy to avoid surprising other logic
->>>>>>> 2ace88f7 (Add chunk prioritization and always RAG support)
-        query_request = query_request.model_copy(deep=True)
-        query_request.query = query_request.query + rag_context.context_text
-=======
     # Build RAG context from Inline RAG sources
-    inline_rag_context = await build_rag_context(client, query_request, configuration)
->>>>>>> a4075c6d (Address review: rename always RAG to inline RAG, Solr config to OKP, fix query mutation)
-
+    inline_rag_context = await build_rag_context(
+        client, query_request.query, query_request.vector_store_ids, query_request.solr
+    )
     # Prepare API request parameters
     responses_params = await prepare_responses_params(
         client=client,
@@ -291,15 +271,7 @@ async def retrieve_response_generator(
     Args:
         responses_params: The Responses API parameters
         context: The response generator context
-<<<<<<< HEAD
-<<<<<<< HEAD
-        doc_ids_from_chunks: List of ReferencedDocument objects extracted from static RAG
-=======
-        pre_rag_documents: Referenced documents from pre-query RAG (BYOK + Solr)
->>>>>>> 2ace88f7 (Add chunk prioritization and always RAG support)
-=======
-        inline_rag_documents: Referenced documents from pre-query RAG (BYOK + Solr)
->>>>>>> a4075c6d (Address review: rename always RAG to inline RAG, Solr config to OKP, fix query mutation)
+        inline_rag_documents: Referenced documents from inline  RAG (BYOK + Solr)
 
     Returns:
         tuple[AsyncIterator[str], TurnSummary]: The response generator and turn summary
@@ -771,25 +743,10 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
         rag_id_mapping=context.rag_id_mapping,
     )
 
-<<<<<<< HEAD
+    # Merge pre-RAG documents with tool-based documents and deduplicate
     turn_summary.referenced_documents = deduplicate_referenced_documents(
-        tool_based_documents + turn_summary.pre_rag_documents
+        turn_summary.inline_rag_documents + tool_based_documents
     )
-=======
-    # Merge pre-RAG documents with tool-based documents (similar to query.py)
-    if turn_summary.inline_rag_documents:
-        all_documents = turn_summary.inline_rag_documents + tool_based_documents
-        seen = set()
-        deduplicated_documents = []
-        for doc in all_documents:
-            key = (doc.doc_url, doc.doc_title)
-            if key not in seen:
-                seen.add(key)
-                deduplicated_documents.append(doc)
-        turn_summary.referenced_documents = deduplicated_documents
-    else:
-        turn_summary.referenced_documents = tool_based_documents
->>>>>>> a4075c6d (Address review: rename always RAG to inline RAG, Solr config to OKP, fix query mutation)
 
 
 def stream_http_error_event(

src/configuration.py

@@ -427,63 +427,6 @@ def inline_solr_enabled(self) -> bool:
             raise LogicError("logic error: configuration is not loaded")
         return constants.OKP_RAG_ID in self._configuration.rag.inline
 
-    @property
-    def inline_byok_vector_store_ids(self) -> list[str]:
-        """Return vector store IDs for the BYOK sources listed in rag.inline.
-
-        Maps non-okp rag_ids in rag.inline to their corresponding vector_db_ids
-        from the byok_rag configuration. IDs that are not found in byok_rag are
-        silently skipped.
-
-        Returns:
-            list[str]: Ordered list of vector_db_ids for inline BYOK RAG.
-
-        Raises:
-            LogicError: If the configuration has not been loaded.
-        """
-        if self._configuration is None:
-            raise LogicError("logic error: configuration is not loaded")
-        inline_ids = [
-            rid for rid in self._configuration.rag.inline if rid != constants.OKP_RAG_ID
-        ]
-        rag_to_vdb = {
-            brag.rag_id: brag.vector_db_id for brag in self._configuration.byok_rag
-        }
-        return [rag_to_vdb[rid] for rid in inline_ids if rid in rag_to_vdb]
-
-    @property
-    def tool_vector_store_ids(self) -> Optional[list[str]]:
-        """Return vector store IDs for tool RAG, or None to use all registered stores.
-
-        When rag.tool is None (default), returns None to signal that all
-        registered vector stores should be used (backward compatibility).
-
-        When rag.tool is an explicit list, maps rag_ids to vector_db_ids and
-        includes the OKP vector store ID for the special 'okp-rag' entry.
-
-        Returns:
-            Optional[list[str]]: List of vector_db_ids for tool RAG, or None
-            when all registered stores should be used.
-
-        Raises:
-            LogicError: If the configuration has not been loaded.
-        """
-        if self._configuration is None:
-            raise LogicError("logic error: configuration is not loaded")
-        tool_ids = self._configuration.rag.tool
-        if tool_ids is None:
-            return None
-        rag_to_vdb = {
-            brag.rag_id: brag.vector_db_id for brag in self._configuration.byok_rag
-        }
-        result = []
-        for rid in tool_ids:
-            if rid == constants.OKP_RAG_ID:
-                result.append(constants.SOLR_DEFAULT_VECTOR_STORE_ID)
-            elif rid in rag_to_vdb:
-                result.append(rag_to_vdb[rid])
-        return result
-
     def resolve_index_name(
         self, vector_store_id: str, rag_id_mapping: Optional[dict[str, str]] = None
     ) -> str:

src/utils/responses.py

@@ -167,14 +167,19 @@ async def prepare_tools(  # pylint: disable=too-many-arguments,too-many-position
         return None
 
     toolgroups: list[InputTool] = []
-    # Per-request vector_store_ids override takes priority.
-    # When not provided, use config-based tool list (or None = all stores).
-    effective_ids = (
-        vector_store_ids
-        if vector_store_ids is not None
-        else configuration.tool_vector_store_ids
-    )
-    effective_ids = await get_vector_store_ids(client, effective_ids)
+
+    # Priority: per-request IDs > rag.tool config > all registered stores.
+    # In all cases, customer-facing rag_ids are translated to internal vector_db_ids.
+    # IDs fetched from llama-stack are already internal and need no translation.
+    byok_rags = configuration.configuration.byok_rag
+    if vector_store_ids is not None:
+        effective_ids: list[str] = resolve_vector_store_ids(vector_store_ids, byok_rags)
+    elif configuration.configuration.rag.tool is not None:
+        effective_ids = resolve_vector_store_ids(
+            configuration.configuration.rag.tool, byok_rags
+        )
+    else:
+        effective_ids = await get_vector_store_ids(client, None)
 
     # Add RAG tools if vector stores are available
     rag_tools = get_rag_tools(effective_ids)
@@ -350,10 +355,11 @@ def extract_vector_store_ids_from_tools(
 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.
+    """Translate customer-facing 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.
+    The special ``okp-rag`` ID is mapped to the Solr vector store ID.
     Otherwise the ID is passed through unchanged (assumed to already be a
     llama-stack vector store ID).
 
@@ -366,6 +372,9 @@ def resolve_vector_store_ids(
         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}
+    rag_id_to_vector_db_id[constants.OKP_RAG_ID] = (
+        constants.SOLR_DEFAULT_VECTOR_STORE_ID
+    )
     return [rag_id_to_vector_db_id.get(vs_id, vs_id) for vs_id in vector_store_ids]
 
 

src/utils/vector_search.py

@@ -16,6 +16,7 @@
 from configuration import configuration
 from log import get_logger
 from models.responses import ReferencedDocument
+from utils.responses import resolve_vector_store_ids
 from utils.types import RAGChunk, RAGContext
 
 logger = get_logger(__name__)
@@ -45,10 +46,10 @@ def _build_query_params(solr: Optional[dict[str, Any]] = None) -> dict[str, Any]
         "mode": constants.SOLR_VECTOR_SEARCH_DEFAULT_MODE,
     }
     logger.debug("Initial params: %s", params)
-    logger.debug("query_request.solr: %s", query_request.solr)
+    logger.debug("query_request.solr: %s", solr)
 
-    if query_request.solr:
-        params["solr"] = query_request.solr
+    if solr:
+        params["solr"] = solr
         logger.debug("Final params with solr filters: %s", params)
     else:
         logger.debug("No solr filters provided")
@@ -316,7 +317,6 @@ def _process_solr_chunks_for_documents(
 async def _fetch_byok_rag(
     client: AsyncLlamaStackClient,
     query: str,
-    configuration: AppConfig,
     vector_store_ids: Optional[list[str]] = None,
 ) -> tuple[list[RAGChunk], list[ReferencedDocument]]:
     """Fetch chunks and documents from BYOK RAG sources.
@@ -347,7 +347,14 @@ async def _fetch_byok_rag(
             if vs_id != constants.SOLR_DEFAULT_VECTOR_STORE_ID
         ]
     else:
-        vector_store_ids_to_query = configuration.inline_byok_vector_store_ids
+        inline_rag_ids = [
+            rid
+            for rid in configuration.configuration.rag.inline
+            if rid != constants.OKP_RAG_ID
+        ]
+        vector_store_ids_to_query = resolve_vector_store_ids(
+            inline_rag_ids, configuration.configuration.byok_rag
+        )
 
     # If inline byok stores are not defined, we disable the inline RAG for backward compatibility
     if not vector_store_ids_to_query:
@@ -410,8 +417,8 @@ async def _fetch_byok_rag(
 
 async def _fetch_solr_rag(
     client: AsyncLlamaStackClient,
-    query_request: QueryRequest,
-    configuration: AppConfig,
+    query: str,
+    solr: Optional[dict[str, Any]] = None,
 ) -> tuple[list[RAGChunk], list[ReferencedDocument]]:
     """Fetch chunks and documents from Solr RAG source.
 
@@ -486,8 +493,9 @@ async def _fetch_solr_rag(
 
 async def build_rag_context(
     client: AsyncLlamaStackClient,
-    query_request: QueryRequest,
-    configuration: AppConfig,
+    query: str,
+    vector_store_ids: Optional[list[str]],
+    solr: Optional[dict[str, Any]] = None,
 ) -> RAGContext:
     """Build RAG context by fetching and merging chunks from all enabled sources.
 
@@ -512,7 +520,7 @@ async def build_rag_context(
     # Merge chunks from all sources (BYOK + Solr)
     context_chunks = byok_chunks + solr_chunks
 
-    context_text = _format_rag_context(context_chunks, query_request.query)
+    context_text = _format_rag_context(context_chunks, query)
 
     logger.debug(
         "Inline RAG context built: %d chunks, %d characters",

tests/unit/models/config/test_byok_rag.py

@@ -28,7 +28,7 @@ def test_byok_rag_configuration_default_values() -> None:
     assert byok_rag.embedding_model == DEFAULT_EMBEDDING_MODEL
     assert byok_rag.embedding_dimension == DEFAULT_EMBEDDING_DIMENSION
     assert byok_rag.vector_db_id == "vector_db_id"
-    assert byok_rag.db_path == Path("tests/configuration/rag.txt")
+    assert byok_rag.db_path == "tests/configuration/rag.txt"
     assert byok_rag.score_multiplier == DEFAULT_SCORE_MULTIPLIER
 
 

tests/unit/utils/test_responses.py

@@ -24,7 +24,6 @@
 from pytest_mock import MockerFixture
 
 import constants
-from configuration import AppConfig
 from models.config import ByokRag, ModelContextProtocolServer
 from models.requests import QueryRequest
 from utils.responses import (
@@ -1122,6 +1121,7 @@ async def test_does_not_translate_when_ids_fetched_from_llama_stack(
         mock_byok_rag.vector_db_id = "vs-translated"
         mock_config = mocker.Mock()
         mock_config.configuration.byok_rag = [mock_byok_rag]
+        mock_config.configuration.rag.tool = None
         mocker.patch("utils.responses.configuration", mock_config)
 
         result = await prepare_tools(mock_client, None, False, "token")

tests/unit/utils/test_vector_search.py

@@ -351,7 +351,8 @@ class TestFetchByokRag:
     async def test_byok_no_inline_ids(self, mocker) -> None:  # type: ignore[no-untyped-def]
         """Test when no inline BYOK sources are configured."""
         config_mock = mocker.Mock(spec=AppConfig)
-        config_mock.inline_byok_vector_store_ids = []
+        config_mock.configuration.rag.inline = []
+        config_mock.configuration.byok_rag = []
         mocker.patch("utils.vector_search.configuration", config_mock)
 
         client_mock = mocker.AsyncMock()
@@ -366,7 +367,11 @@ async def test_byok_enabled_success(self, mocker) -> None:  # type: ignore[no-un
         """Test successful BYOK RAG fetch when inline IDs are configured."""
         # Mock configuration
         config_mock = mocker.Mock(spec=AppConfig)
-        config_mock.inline_byok_vector_store_ids = ["vs_1"]
+        byok_rag_mock = mocker.Mock()
+        byok_rag_mock.rag_id = "rag_1"
+        byok_rag_mock.vector_db_id = "vs_1"
+        config_mock.configuration.rag.inline = ["rag_1"]
+        config_mock.configuration.byok_rag = [byok_rag_mock]
         config_mock.score_multiplier_mapping = {"vs_1": 1.5}
         config_mock.rag_id_mapping = {"vs_1": "rag_1"}
         mocker.patch("utils.vector_search.configuration", config_mock)
@@ -451,7 +456,8 @@ class TestBuildRagContext:
     async def test_both_sources_disabled(self, mocker) -> None:  # type: ignore[no-untyped-def]
         """Test when both BYOK inline and Solr inline are not configured."""
         config_mock = mocker.Mock(spec=AppConfig)
-        config_mock.inline_byok_vector_store_ids = []
+        config_mock.configuration.rag.inline = []
+        config_mock.configuration.byok_rag = []
         config_mock.inline_solr_enabled = False
         mocker.patch("utils.vector_search.configuration", config_mock)
 
@@ -467,7 +473,11 @@ async def test_byok_enabled_only(self, mocker) -> None:  # type: ignore[no-untyp
         """Test when only inline BYOK is configured."""
         # Mock configuration
         config_mock = mocker.Mock(spec=AppConfig)
-        config_mock.inline_byok_vector_store_ids = ["vs_1"]
+        byok_rag_mock = mocker.Mock()
+        byok_rag_mock.rag_id = "rag_1"
+        byok_rag_mock.vector_db_id = "vs_1"
+        config_mock.configuration.rag.inline = ["rag_1"]
+        config_mock.configuration.byok_rag = [byok_rag_mock]
         config_mock.inline_solr_enabled = False
         config_mock.score_multiplier_mapping = {"vs_1": 1.0}
         config_mock.rag_id_mapping = {"vs_1": "rag_1"}