From 47f1be30e6d58732010fb37a7a552241be2c277e Mon Sep 17 00:00:00 2001 From: Jazzcort Date: Tue, 14 Jul 2026 11:49:13 -0400 Subject: [PATCH 1/3] Remove fallback to all registered vector stores when tool RAG is unconfigured MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, when neither rag.tool nor rag.inline were configured, prepare_tools would fetch and register all vector stores from llama-stack as a backward-compatibility fallback. This implicit behavior is removed — tool RAG now requires explicit configuration via rag.tool. Updated docs, config model, example YAML, and tests to reflect the new behavior. --- docs/devel_doc/openapi.json | 4 +- docs/devel_doc/openapi.md | 10 +- docs/models/responses_succ.md | 10 +- docs/user_doc/byok_guide.md | 2 +- docs/user_doc/config.md | 10 +- examples/lightspeed-stack-byok-okp-rag.yaml | 2 +- src/models/config.py | 10 +- src/utils/responses.py | 48 +----- tests/unit/utils/test_responses.py | 168 +------------------- 9 files changed, 26 insertions(+), 238 deletions(-) diff --git a/docs/devel_doc/openapi.json b/docs/devel_doc/openapi.json index ba4df6865..f1acabe1e 100644 --- a/docs/devel_doc/openapi.json +++ b/docs/devel_doc/openapi.json @@ -18084,13 +18084,13 @@ }, "type": "array", "title": "Tool RAG IDs", - "description": "RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, all registered BYOK vector stores are used (backward compatibility)." + "description": "RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, tool RAG is disabled." } }, "additionalProperties": false, "type": "object", "title": "RagConfiguration", - "description": "RAG strategy configuration.\n\nControls which RAG sources are used for inline and tool-based retrieval.\n\nEach strategy lists RAG IDs to include. The special ID ``\"okp\"`` defined in constants,\nactivates the OKP provider; all other IDs refer to entries in ``byok_rag``.\n\nBackward compatibility:\n - ``inline`` defaults to ``[]`` (no inline RAG).\n - ``tool`` defaults to ``[]`` (no tool RAG).\n\nIf no RAG strategy is defined (inline and tool are empty),\nthe RAG tool will register all stores available to llama-stack." + "description": "RAG strategy configuration.\n\nControls which RAG sources are used for inline and tool-based retrieval.\n\nEach strategy lists RAG IDs to include. The special ID ``\"okp\"`` defined in constants,\nactivates the OKP provider; all other IDs refer to entries in ``byok_rag``.\n\nBoth ``inline`` and ``tool`` default to ``[]`` (disabled).\nEach must be explicitly configured to activate its respective RAG strategy." }, "ReadinessResponse": { "properties": { diff --git a/docs/devel_doc/openapi.md b/docs/devel_doc/openapi.md index aee97bdd8..34e1e9f1c 100644 --- a/docs/devel_doc/openapi.md +++ b/docs/devel_doc/openapi.md @@ -8155,18 +8155,14 @@ Controls which RAG sources are used for inline and tool-based retrieval. Each strategy lists RAG IDs to include. The special ID ``"okp"`` defined in constants, activates the OKP provider; all other IDs refer to entries in ``byok_rag``. -Backward compatibility: - - ``inline`` defaults to ``[]`` (no inline RAG). - - ``tool`` defaults to ``[]`` (no tool RAG). - -If no RAG strategy is defined (inline and tool are empty), -the RAG tool will register all stores available to llama-stack. +Both ``inline`` and ``tool`` default to ``[]`` (disabled). +Each must be explicitly configured to activate its respective RAG strategy. | Field | Type | Description | |-------|------|-------------| | inline | array | RAG IDs whose sources are injected as context before the LLM call. Use 'okp' to enable OKP inline RAG. Empty by default (no inline RAG). | -| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, all registered BYOK vector stores are used (backward compatibility). | +| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, tool RAG is disabled. | ## ReadinessResponse diff --git a/docs/models/responses_succ.md b/docs/models/responses_succ.md index 9e3fc8295..8a4918102 100644 --- a/docs/models/responses_succ.md +++ b/docs/models/responses_succ.md @@ -1840,18 +1840,14 @@ Controls which RAG sources are used for inline and tool-based retrieval. Each strategy lists RAG IDs to include. The special ID ``"okp"`` defined in constants, activates the OKP provider; all other IDs refer to entries in ``byok_rag``. -Backward compatibility: - - ``inline`` defaults to ``[]`` (no inline RAG). - - ``tool`` defaults to ``[]`` (no tool RAG). - -If no RAG strategy is defined (inline and tool are empty), -the RAG tool will register all stores available to llama-stack. +Both ``inline`` and ``tool`` default to ``[]`` (disabled). +Each must be explicitly configured to activate its respective RAG strategy. | Field | Type | Description | |-------|------|-------------| | inline | array | RAG IDs whose sources are injected as context before the LLM call. Use 'okp' to enable OKP inline RAG. Empty by default (no inline RAG). | -| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, all registered BYOK vector stores are used (backward compatibility). | +| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, tool RAG is disabled. | ## ReadinessResponse diff --git a/docs/user_doc/byok_guide.md b/docs/user_doc/byok_guide.md index dafdc9355..aaecfccfc 100644 --- a/docs/user_doc/byok_guide.md +++ b/docs/user_doc/byok_guide.md @@ -261,7 +261,7 @@ rag: - okp # include OKP context inline # Tool RAG: the LLM can call file_search to retrieve context on demand - # If omitted, tool RAG is disabled. If both tool and inline are omitted, all registered stores are used as fallback + # If omitted, tool RAG is disabled tool: - my-docs # expose this BYOK store as the file_search tool - okp # expose OKP as the file_search tool diff --git a/docs/user_doc/config.md b/docs/user_doc/config.md index 752dcc0dc..33e2d7b71 100644 --- a/docs/user_doc/config.md +++ b/docs/user_doc/config.md @@ -606,18 +606,14 @@ Controls which RAG sources are used for inline and tool-based retrieval. Each strategy lists RAG IDs to include. The special ID ``"okp"`` defined in constants, activates the OKP provider; all other IDs refer to entries in ``byok_rag``. -Backward compatibility: - - ``inline`` defaults to ``[]`` (no inline RAG). - - ``tool`` defaults to ``[]`` (no tool RAG). - -If no RAG strategy is defined (inline and tool are empty), -the RAG tool will register all stores available to llama-stack. +Both ``inline`` and ``tool`` default to ``[]`` (disabled). +Each must be explicitly configured to activate its respective RAG strategy. | Field | Type | Description | |--------|-------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | inline | array | RAG IDs whose sources are injected as context before the LLM call. Use 'okp' to enable OKP inline RAG. Empty by default (no inline RAG). | -| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, all registered BYOK vector stores are used (backward compatibility). | +| tool | array | RAG IDs made available to the LLM as a file_search tool. Use 'okp' to include the OKP vector store. When omitted, tool RAG is disabled. | ## RerankerConfiguration diff --git a/examples/lightspeed-stack-byok-okp-rag.yaml b/examples/lightspeed-stack-byok-okp-rag.yaml index 7cbb36fc8..b99fd1e88 100644 --- a/examples/lightspeed-stack-byok-okp-rag.yaml +++ b/examples/lightspeed-stack-byok-okp-rag.yaml @@ -60,7 +60,7 @@ rag: - okp # Tool RAG: LLM can call file_search on demand to retrieve context # List rag_ids from byok_rag, or 'okp' to include OKP - # Omit to use all registered BYOK stores (backward compatibility) + # Omit to disable tool RAG tool: - ocp-docs - knowledge-base diff --git a/src/models/config.py b/src/models/config.py index 28ff23264..6ab7f6947 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -2258,12 +2258,8 @@ class RagConfiguration(ConfigurationBase): Each strategy lists RAG IDs to include. The special ID ``"okp"`` defined in constants, activates the OKP provider; all other IDs refer to entries in ``byok_rag``. - Backward compatibility: - - ``inline`` defaults to ``[]`` (no inline RAG). - - ``tool`` defaults to ``[]`` (no tool RAG). - - If no RAG strategy is defined (inline and tool are empty), - the RAG tool will register all stores available to llama-stack. + Both ``inline`` and ``tool`` default to ``[]`` (disabled). + Each must be explicitly configured to activate its respective RAG strategy. """ inline: list[str] = Field( @@ -2278,7 +2274,7 @@ class RagConfiguration(ConfigurationBase): title="Tool RAG IDs", description="RAG IDs made available to the LLM as a file_search tool. " f"Use '{constants.OKP_RAG_ID}' to include the OKP vector store. " - "When omitted, all registered BYOK vector stores are used (backward compatibility).", + "When omitted, tool RAG is disabled.", ) diff --git a/src/utils/responses.py b/src/utils/responses.py index 5e5916e0b..87bf5eb67 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -127,44 +127,6 @@ logger = get_logger(__name__) -async def get_vector_store_ids( - client: AsyncLlamaStackClient, - vector_store_ids: Optional[list[str]] = None, -) -> list[str]: - """Get vector store IDs for querying. - - If vector_store_ids are provided, returns them. Otherwise fetches all - available vector stores from Llama Stack. - - Args: - client: The AsyncLlamaStackClient to use for fetching stores - vector_store_ids: Optional list of vector store IDs. If provided, - returns this list. If None, fetches all available vector stores. - - Returns: - List of vector store IDs to query - - Raises: - HTTPException: With ServiceUnavailableResponse if connection fails, - or InternalServerErrorResponse if API returns an error status - """ - if vector_store_ids is not None: - return vector_store_ids - - try: - vector_stores = await client.vector_stores.list() - return [vector_store.id for vector_store in vector_stores.data] - except APIConnectionError as e: - error_response = ServiceUnavailableResponse( - backend_name="Llama Stack", - cause=str(e), - ) - raise HTTPException(**error_response.model_dump()) from e - except APIStatusError as e: - error_response = InternalServerErrorResponse.generic() - raise HTTPException(**error_response.model_dump()) from e - - async def get_topic_summary( # pylint: disable=too-many-nested-blocks question: str, client: AsyncLlamaStackClient, model_id: str ) -> str: @@ -228,7 +190,7 @@ async def maybe_get_topic_summary( return await get_topic_summary(input_text, client, model_id) -async def prepare_tools( # pylint: disable=too-many-arguments,too-many-positional-arguments +async def prepare_tools( # pylint: disable=too-many-arguments,too-many-positional-arguments,unused-argument client: AsyncLlamaStackClient, vector_store_ids: Optional[list[str]], no_tools: Optional[bool], @@ -241,7 +203,7 @@ async def prepare_tools( # pylint: disable=too-many-arguments,too-many-position Args: client: The Llama Stack client instance vector_store_ids: The list of vector store IDs to use for RAG tools - or None if all vector stores should be used + or None to fall back to rag.tool configuration no_tools: Whether to skip tool preparation token: Authentication token for MCP tools mcp_headers: Per-request headers for MCP servers @@ -259,13 +221,9 @@ async def prepare_tools( # pylint: disable=too-many-arguments,too-many-position # Vector store ID resolution priority: # 1. Per-request IDs: highest prio; customer-facing rag_ids are translated to vector_db_ids. # 2. rag.tool config IDs: used when no per-request IDs provided, and rag.tool is configured. - # If rag.inline is configured, but not rag.tool, tool RAG is disabled. - # 3. All registered vector DBs: fallback when neither rag.tool nor rag.inline are configured. - # IDs fetched from llama-stack are already internal and need no translation. byok_rags = configuration.configuration.byok_rag is_tool_rag_enabled = len(configuration.configuration.rag.tool) > 0 - is_inline_rag_enabled = len(configuration.configuration.rag.inline) > 0 if vector_store_ids is not None: effective_ids = resolve_vector_store_ids(vector_store_ids, byok_rags) @@ -273,8 +231,6 @@ async def prepare_tools( # pylint: disable=too-many-arguments,too-many-position effective_ids = resolve_vector_store_ids( configuration.configuration.rag.tool, byok_rags ) - elif not is_inline_rag_enabled: - effective_ids = await get_vector_store_ids(client, None) # Add RAG tools if vector stores are available rag_tools = get_rag_tools(effective_ids) diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index eb9c4b1b1..8e587eb97 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -78,7 +78,6 @@ get_mcp_tools, get_rag_tools, get_topic_summary, - get_vector_store_ids, is_server_deployed_output, parse_arguments_string, parse_referenced_documents, @@ -1614,43 +1613,6 @@ async def test_prepare_tools_with_vector_store_ids( assert result[0].type == "file_search" assert result[0].vector_store_ids == ["vs1", "vs2"] - @pytest.mark.asyncio - async def test_prepare_tools_fetch_vector_stores( - self, mocker: MockerFixture - ) -> None: - """Test prepare_tools fetches vector stores when not specified.""" - mock_client = mocker.AsyncMock() - mock_vector_store1 = mocker.Mock() - mock_vector_store1.id = "vs1" - mock_vector_store2 = mocker.Mock() - mock_vector_store2.id = "vs2" - mock_vector_stores = mocker.Mock() - mock_vector_stores.data = [mock_vector_store1, mock_vector_store2] - mock_client.vector_stores.list = mocker.AsyncMock( - return_value=mock_vector_stores - ) - mocker.patch("utils.responses.get_mcp_tools", return_value=None) - - result = await prepare_tools(mock_client, None, False, "token") - assert result is not None - assert len(result) == 1 - assert isinstance(result[0], InputToolFileSearch) - assert result[0].vector_store_ids == ["vs1", "vs2"] - - @pytest.mark.asyncio - async def test_prepare_tools_connection_error(self, mocker: MockerFixture) -> None: - """Test prepare_tools raises HTTPException on connection error.""" - mock_client = mocker.AsyncMock() - mock_client.vector_stores.list = mocker.AsyncMock( - side_effect=APIConnectionError( - message="Connection failed", request=mocker.Mock() - ) - ) - - with pytest.raises(HTTPException) as exc_info: - await prepare_tools(mock_client, None, False, "token") - assert exc_info.value.status_code == 503 - @pytest.mark.asyncio async def test_prepare_tools_with_mcp_servers(self, mocker: MockerFixture) -> None: """Test prepare_tools includes MCP tools.""" @@ -1666,20 +1628,6 @@ async def test_prepare_tools_with_mcp_servers(self, mocker: MockerFixture) -> No assert len(result) == 2 # RAG tool + MCP tool assert any(tool.type == "mcp" for tool in result) - @pytest.mark.asyncio - async def test_prepare_tools_api_status_error(self, mocker: MockerFixture) -> None: - """Test prepare_tools raises HTTPException on API status error when fetching vector stores.""" - mock_client = mocker.AsyncMock() - mock_client.vector_stores.list = mocker.AsyncMock( - side_effect=APIStatusError( - message="API error", response=mocker.Mock(request=None), body=None - ) - ) - - with pytest.raises(HTTPException) as exc_info: - await prepare_tools(mock_client, None, False, "token") - assert exc_info.value.status_code == 500 - @pytest.mark.asyncio async def test_prepare_tools_empty_toolgroups(self, mocker: MockerFixture) -> None: """Test prepare_tools returns None when no tools are available.""" @@ -1812,38 +1760,6 @@ async def test_passes_through_unknown_ids_in_prepare_tools( assert isinstance(result[0], InputToolFileSearch) 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] - mock_config.configuration.rag.tool = [] - mock_config.configuration.rag.inline = [] - 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 isinstance(result[0], InputToolFileSearch) - assert result[0].vector_store_ids == ["vs-internal"] - class TestPrepareToolsVectorStoreResolution: """Tests for vector store ID resolution priority in prepare_tools.""" @@ -1895,8 +1811,10 @@ async def test_rag_tool_config_ids_are_translated( mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio - async def test_inline_rag_disables_tool_rag(self, mocker: MockerFixture) -> None: - """Test that configuring rag.inline without rag.tool disables tool RAG.""" + async def test_inline_rag_config_does_not_affect_tool_rag( + self, mocker: MockerFixture + ) -> None: + """Test that configuring rag.inline have no effect on tool rag.""" mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) @@ -1910,7 +1828,6 @@ async def test_inline_rag_disables_tool_rag(self, mocker: MockerFixture) -> None result = await prepare_tools(mock_client, None, False, "token") - # Tool RAG should be disabled — no RAG tool in result, no llama-stack fetch assert result is None mock_client.vector_stores.list.assert_not_called() @@ -1936,30 +1853,22 @@ async def test_per_request_ids_override_rag_tool_config( mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio - async def test_all_registered_dbs_used_when_neither_tool_nor_inline_configured( + async def test_tool_rag_disabled_when_tool_not_configured( self, mocker: MockerFixture ) -> None: - """Test fallback to all registered vector DBs when neither rag.tool nor rag.inline are set.""" + """Test no rag tools is returned when rag.tool is not set and there is no per-request vector ids.""" mock_client = mocker.AsyncMock() - mock_vs = mocker.Mock() - mock_vs.id = "vs-registered" - mock_list = mocker.Mock() - mock_list.data = [mock_vs] - mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_list) mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_config = mocker.Mock() mock_config.configuration.byok_rag = [] mock_config.configuration.rag.tool = [] - mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) result = await prepare_tools(mock_client, None, False, "token") - assert result is not None - assert isinstance(result[0], InputToolFileSearch) - assert result[0].vector_store_ids == ["vs-registered"] - mock_client.vector_stores.list.assert_called_once() + assert result is None + mock_client.vector_stores.list.assert_not_called() class TestPrepareResponsesParams: @@ -3145,67 +3054,6 @@ def test_multiple_stores_source_is_none(self, mocker: MockerFixture) -> None: assert docs[0].source is None -class TestGetVectorStoreIds: - """Tests for get_vector_store_ids utility function.""" - - @pytest.mark.asyncio - async def test_returns_provided_ids_directly(self, mocker: MockerFixture) -> None: - """Test that provided vector_store_ids are returned without fetching.""" - client_mock = mocker.AsyncMock() - result = await get_vector_store_ids(client_mock, ["vs1", "vs2"]) - assert result == ["vs1", "vs2"] - client_mock.vector_stores.list.assert_not_called() - - @pytest.mark.asyncio - async def test_fetches_all_when_no_ids_provided( - self, mocker: MockerFixture - ) -> None: - """Test that all vector stores are fetched when no IDs provided.""" - mock_store1 = mocker.Mock() - mock_store1.id = "vs-fetched-1" - mock_store2 = mocker.Mock() - mock_store2.id = "vs-fetched-2" - - mock_list_result = mocker.Mock() - mock_list_result.data = [mock_store1, mock_store2] - - client_mock = mocker.AsyncMock() - client_mock.vector_stores.list.return_value = mock_list_result - - result = await get_vector_store_ids(client_mock, None) - assert result == ["vs-fetched-1", "vs-fetched-2"] - client_mock.vector_stores.list.assert_called_once() - - @pytest.mark.asyncio - async def test_raises_on_connection_error(self, mocker: MockerFixture) -> None: - """Test that APIConnectionError raises HTTPException 503.""" - client_mock = mocker.AsyncMock() - client_mock.vector_stores.list.side_effect = APIConnectionError.__new__( - APIConnectionError - ) - - with pytest.raises(HTTPException) as exc_info: - await get_vector_store_ids(client_mock, None) - assert exc_info.value.status_code == 503 - - @pytest.mark.asyncio - async def test_raises_on_api_status_error(self, mocker: MockerFixture) -> None: - """Test that APIStatusError raises HTTPException 500.""" - mock_response = mocker.Mock() - mock_response.status_code = 500 - mock_response.headers = {} - mock_response.text = "error" - - client_mock = mocker.AsyncMock() - client_mock.vector_stores.list.side_effect = APIStatusError( - "error", response=mock_response, body=None - ) - - with pytest.raises(HTTPException) as exc_info: - await get_vector_store_ids(client_mock, None) - assert exc_info.value.status_code == 500 - - class TestGetRAGToolsWithConfig: """Tests for get_rag_tools with configuration checks.""" From 9330701c50ba3008b3db537738f3dee4d6d7abb2 Mon Sep 17 00:00:00 2001 From: Jazzcort Date: Wed, 15 Jul 2026 13:53:03 -0400 Subject: [PATCH 2/3] Remove unused client parameter from prepare_tools --- src/utils/responses.py | 12 +-- .../endpoints/test_query_integration.py | 17 +--- .../test_responses_byok_integration.py | 1 - .../endpoints/test_responses_integration.py | 1 - tests/unit/app/endpoints/test_responses.py | 6 -- tests/unit/utils/test_responses.py | 90 +++---------------- 6 files changed, 17 insertions(+), 110 deletions(-) diff --git a/src/utils/responses.py b/src/utils/responses.py index 87bf5eb67..7af8ce8dd 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -81,7 +81,6 @@ from llama_stack_client import APIConnectionError, APIStatusError, AsyncLlamaStackClient import constants -from client import AsyncLlamaStackClientHolder from configuration import configuration from constants import DEFAULT_RAG_TOOL from log import get_logger @@ -190,8 +189,7 @@ async def maybe_get_topic_summary( return await get_topic_summary(input_text, client, model_id) -async def prepare_tools( # pylint: disable=too-many-arguments,too-many-positional-arguments,unused-argument - client: AsyncLlamaStackClient, +async def prepare_tools( # pylint: disable=too-many-arguments,too-many-positional-arguments vector_store_ids: Optional[list[str]], no_tools: Optional[bool], token: str, @@ -201,7 +199,6 @@ async def prepare_tools( # pylint: disable=too-many-arguments,too-many-position """Prepare tools for Responses API including RAG and MCP tools. Args: - client: The Llama Stack client instance vector_store_ids: The list of vector store IDs to use for RAG tools or None to fall back to rag.tool configuration no_tools: Whether to skip tool preparation @@ -331,7 +328,6 @@ async def prepare_responses_params( # pylint: disable=too-many-arguments,too-ma # Prepare tools for responses API tools = await prepare_tools( - client, query_request.vector_store_ids, query_request.no_tools, token, @@ -1710,9 +1706,7 @@ async def _resolve_client_tools( # Optionally merge server-configured tools (RAG, MCP) with client tools if merge_server_tools: - client = AsyncLlamaStackClientHolder().get_client() server_tools = await prepare_tools( - client=client, vector_store_ids=vector_store_ids, no_tools=False, token=token, @@ -1740,9 +1734,7 @@ async def _resolve_server_tools( Returns: List of server-configured tools, or None if none are configured. """ - client = AsyncLlamaStackClientHolder().get_client() return await prepare_tools( - client=client, vector_store_ids=None, # allow all vector stores configured no_tools=False, token=token, @@ -1786,9 +1778,7 @@ async def resolve_tool_choice( if tools is None: # Register all tools configured in LCORE configuration - client = AsyncLlamaStackClientHolder().get_client() prepared_tools = await prepare_tools( - client=client, vector_store_ids=None, # allow all vector stores configured no_tools=False, token=token, diff --git a/tests/integration/endpoints/test_query_integration.py b/tests/integration/endpoints/test_query_integration.py index d29bd228c..a1e29cadf 100644 --- a/tests/integration/endpoints/test_query_integration.py +++ b/tests/integration/endpoints/test_query_integration.py @@ -4,7 +4,6 @@ # pylint: disable=too-many-arguments # Integration tests need many fixtures # pylint: disable=too-many-positional-arguments # Integration tests need many fixtures - import pytest from fastapi import HTTPException, Request, status from llama_stack_client import APIConnectionError @@ -676,14 +675,13 @@ async def test_query_v2_endpoint_bypasses_tools_when_no_tools_true( @pytest.mark.asyncio -async def test_query_v2_endpoint_uses_tools_when_available( +async def test_query_v2_endpoint_uses_tools_when_available( # pylint: disable=unused-argument test_config: AppConfig, mock_llama_stack_client: AsyncMockType, mock_query_agent: AsyncMockType, test_request: Request, test_auth: AuthTuple, patch_db_session: Session, - mocker: MockerFixture, ) -> None: """Test that tools are used when no_tools=False and vector stores are available. @@ -707,18 +705,11 @@ async def test_query_v2_endpoint_uses_tools_when_available( ------- None """ - _ = test_config + # prepare_tools does not require llama-stack client anymore so the way to + # enable RAG tools is through config + test_config.rag.tool = ["vs-test-123"] _ = patch_db_session - # Mock vector stores to be available (simulating RAG tools) - mock_vector_store = mocker.MagicMock() - mock_vector_store.id = "vs-test-123" - - mock_list_result = mocker.MagicMock() - mock_list_result.data = [mock_vector_store] - - mock_llama_stack_client.vector_stores.list.return_value = mock_list_result - query_request = QueryRequest(query="What is Ansible?", no_tools=False) response = await query_endpoint_handler( diff --git a/tests/integration/endpoints/test_responses_byok_integration.py b/tests/integration/endpoints/test_responses_byok_integration.py index d316af6a6..8a702dd9e 100644 --- a/tests/integration/endpoints/test_responses_byok_integration.py +++ b/tests/integration/endpoints/test_responses_byok_integration.py @@ -79,7 +79,6 @@ def _patch_all_client_holders(mocker: MockerFixture, mock_client: Any) -> None: for module in ( "app.endpoints.responses", "utils.endpoints", - "utils.responses", ): holder = mocker.patch(f"{module}.AsyncLlamaStackClientHolder") holder.return_value.get_client.return_value = mock_client diff --git a/tests/integration/endpoints/test_responses_integration.py b/tests/integration/endpoints/test_responses_integration.py index 19a7237f6..96016e118 100644 --- a/tests/integration/endpoints/test_responses_integration.py +++ b/tests/integration/endpoints/test_responses_integration.py @@ -118,7 +118,6 @@ def _patch_client_holders(mocker: MockerFixture, mock_client: Any) -> None: for module in ( "app.endpoints.responses", "utils.endpoints", - "utils.responses", ): holder = mocker.patch(f"{module}.AsyncLlamaStackClientHolder") holder.return_value.get_client.return_value = mock_client diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index 34fede253..5c0c851fa 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -122,12 +122,6 @@ def _patch_base(mocker: MockerFixture, config: AppConfig) -> None: mocker.patch(f"{MODULE}.check_configuration_loaded") mocker.patch(f"{MODULE}.check_tokens_available") mocker.patch(f"{MODULE}.validate_model_provider_override") - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mocker.Mock() - mocker.patch( - f"{UTILS_RESPONSES_MODULE}.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) mocker.patch( f"{UTILS_RESPONSES_MODULE}.prepare_tools", new=mocker.AsyncMock(return_value=None), diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index 8e587eb97..5bb1bddbc 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -1034,7 +1034,6 @@ async def test_tool_choice_none_returns_none_tuple( self, mocker: MockerFixture, tools_arg: Optional[list[InputTool]] ) -> None: """ToolChoiceMode.none always yields (None, None).""" - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") mocker.patch("utils.responses.prepare_tools", new_callable=mocker.AsyncMock) out = await resolve_tool_choice( tools_arg, @@ -1118,7 +1117,6 @@ async def test_tool_choice_object_implicit_prepare_empty_returns_none_tuple( new_callable=mocker.AsyncMock, return_value=None, ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") tool_choice_obj = ToolChoiceFileSearch() prepared, choice = await resolve_tool_choice( None, @@ -1249,7 +1247,6 @@ async def test_implicit_tool_choice_uses_prepare_tools( new_callable=mocker.AsyncMock, return_value=[fs], ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") prepared, choice = await resolve_tool_choice( None, mode_choice, @@ -1268,7 +1265,6 @@ async def test_implicit_prepare_tools_returns_none_clears_tool_choice( new_callable=mocker.AsyncMock, return_value=None, ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") prepared, choice = await resolve_tool_choice( None, ToolChoiceMode.auto, @@ -1289,7 +1285,6 @@ async def test_allowed_tools_applies_after_prepare_tools( new_callable=mocker.AsyncMock, return_value=[fs, mcp], ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") allowed = OpenAIResponseInputToolChoiceAllowedTools( mode="auto", tools=[{"type": "mcp", "server_label": "s1"}], @@ -1315,7 +1310,6 @@ async def test_allowed_tools_implicit_filter_excludes_all_tools( new_callable=mocker.AsyncMock, return_value=[mcp], ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") allowed = OpenAIResponseInputToolChoiceAllowedTools( mode="auto", tools=[{"type": "file_search"}], @@ -1335,7 +1329,6 @@ async def test_allowed_tools_implicit_required_mode_after_prepare( new_callable=mocker.AsyncMock, return_value=[mcp], ) - mocker.patch("utils.responses.AsyncLlamaStackClientHolder.get_client") allowed = OpenAIResponseInputToolChoiceAllowedTools( mode="required", tools=[{"type": "mcp"}], @@ -1604,10 +1597,9 @@ async def test_prepare_tools_with_vector_store_ids( self, mocker: MockerFixture ) -> None: """Test prepare_tools with specified vector store IDs.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) - result = await prepare_tools(mock_client, ["vs1", "vs2"], False, "token") + result = await prepare_tools(["vs1", "vs2"], False, "token") assert result is not None assert len(result) == 1 assert result[0].type == "file_search" @@ -1616,14 +1608,13 @@ async def test_prepare_tools_with_vector_store_ids( @pytest.mark.asyncio async def test_prepare_tools_with_mcp_servers(self, mocker: MockerFixture) -> None: """Test prepare_tools includes MCP tools.""" - mock_client = mocker.AsyncMock() mock_mcp_tool = InputToolMCP( server_label="test-server", server_url="http://test", ) mocker.patch("utils.responses.get_mcp_tools", return_value=[mock_mcp_tool]) - result = await prepare_tools(mock_client, ["vs1"], False, "token") + result = await prepare_tools(["vs1"], False, "token") assert result is not None assert len(result) == 2 # RAG tool + MCP tool assert any(tool.type == "mcp" for tool in result) @@ -1631,28 +1622,18 @@ async def test_prepare_tools_with_mcp_servers(self, mocker: MockerFixture) -> No @pytest.mark.asyncio async def test_prepare_tools_empty_toolgroups(self, mocker: MockerFixture) -> None: """Test prepare_tools returns None when no tools are available.""" - mock_client = mocker.AsyncMock() - mock_vector_stores = mocker.Mock() - mock_vector_stores.data = [] # No vector stores - mock_client.vector_stores.list = mocker.AsyncMock( - return_value=mock_vector_stores - ) mocker.patch("utils.responses.get_mcp_tools", return_value=None) - result = await prepare_tools(mock_client, None, False, "token") + result = await prepare_tools(None, False, "token") assert result is None @pytest.mark.asyncio - async def test_prepare_tools_no_tools_true(self, mocker: MockerFixture) -> None: + async def test_prepare_tools_no_tools_true(self) -> None: """Test prepare_tools returns None when no_tools=True.""" - mock_client = mocker.AsyncMock() - # Should return None immediately without fetching vector stores or MCP tools - result = await prepare_tools(mock_client, ["vs1", "vs2"], True, "token") + result = await prepare_tools(["vs1", "vs2"], True, "token") assert result is None - # Verify that vector_stores.list was not called - mock_client.vector_stores.list.assert_not_called() class TestResolveVectorStoreIds: @@ -1721,7 +1702,6 @@ 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 @@ -1734,7 +1714,7 @@ async def test_translates_byok_ids_in_prepare_tools( mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, ["ocp_docs"], False, "token") + result = await prepare_tools(["ocp_docs"], False, "token") assert result is not None assert len(result) == 1 assert result[0].type == "file_search" @@ -1745,7 +1725,6 @@ 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 @@ -1755,7 +1734,7 @@ async def test_passes_through_unknown_ids_in_prepare_tools( mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, ["raw-internal-id"], False, "token") + result = await prepare_tools(["raw-internal-id"], False, "token") assert result is not None assert isinstance(result[0], InputToolFileSearch) assert result[0].vector_store_ids == ["raw-internal-id"] @@ -1769,7 +1748,6 @@ async def test_uses_rag_tool_config_when_no_per_request_ids( self, mocker: MockerFixture ) -> None: """Test that rag.tool config IDs are used when no per-request IDs are provided.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_config = mocker.Mock() @@ -1778,20 +1756,18 @@ async def test_uses_rag_tool_config_when_no_per_request_ids( mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, None, False, "token") + result = await prepare_tools(None, False, "token") assert result is not None assert len(result) == 1 assert result[0].type == "file_search" assert result[0].vector_store_ids == ["rag-tool-id-1", "rag-tool-id-2"] - mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio async def test_rag_tool_config_ids_are_translated( self, mocker: MockerFixture ) -> None: """Test that rag.tool config IDs are translated from rag_ids to vector_db_ids.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_byok_rag = mocker.Mock() @@ -1803,19 +1779,17 @@ async def test_rag_tool_config_ids_are_translated( mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, None, False, "token") + result = await prepare_tools(None, False, "token") assert result is not None assert isinstance(result[0], InputToolFileSearch) assert result[0].vector_store_ids == ["vs-001"] - mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio async def test_inline_rag_config_does_not_affect_tool_rag( self, mocker: MockerFixture ) -> None: """Test that configuring rag.inline have no effect on tool rag.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_config = mocker.Mock() @@ -1826,17 +1800,15 @@ async def test_inline_rag_config_does_not_affect_tool_rag( ] # inline is configured mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, None, False, "token") + result = await prepare_tools(None, False, "token") assert result is None - mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio async def test_per_request_ids_override_rag_tool_config( self, mocker: MockerFixture ) -> None: """Test that per-request vector_store_ids take priority over rag.tool config.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_config = mocker.Mock() @@ -1845,19 +1817,17 @@ async def test_per_request_ids_override_rag_tool_config( mock_config.configuration.rag.inline = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, ["request-id-1"], False, "token") + result = await prepare_tools(["request-id-1"], False, "token") assert result is not None assert isinstance(result[0], InputToolFileSearch) assert result[0].vector_store_ids == ["request-id-1"] - mock_client.vector_stores.list.assert_not_called() @pytest.mark.asyncio async def test_tool_rag_disabled_when_tool_not_configured( self, mocker: MockerFixture ) -> None: """Test no rag tools is returned when rag.tool is not set and there is no per-request vector ids.""" - mock_client = mocker.AsyncMock() mocker.patch("utils.responses.get_mcp_tools", return_value=None) mock_config = mocker.Mock() @@ -1865,10 +1835,9 @@ async def test_tool_rag_disabled_when_tool_not_configured( mock_config.configuration.rag.tool = [] mocker.patch("utils.responses.configuration", mock_config) - result = await prepare_tools(mock_client, None, False, "token") + result = await prepare_tools(None, False, "token") assert result is None - mock_client.vector_stores.list.assert_not_called() class TestPrepareResponsesParams: @@ -3244,13 +3213,6 @@ async def test_client_tools_without_merge_header( self, mocker: MockerFixture ) -> None: """Test client tools used as-is without merge header.""" - mock_client = mocker.AsyncMock() - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mock_client - mocker.patch( - "utils.responses.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) mock_config = mocker.Mock() mock_config.configuration.byok_rag = [] mock_config.mcp_servers = [] @@ -3269,13 +3231,6 @@ async def test_client_tools_without_merge_header( @pytest.mark.asyncio async def test_client_tools_with_merge_header(self, mocker: MockerFixture) -> None: """Test client tools merged with server tools when header is set.""" - mock_client = mocker.AsyncMock() - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mock_client - mocker.patch( - "utils.responses.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) mock_config = mocker.Mock() mock_config.configuration.byok_rag = [] mock_config.mcp_servers = [] @@ -3306,13 +3261,6 @@ async def test_merge_header_conflict_raises_409( self, mocker: MockerFixture ) -> None: """Test 409 when merge header is set and tools conflict.""" - mock_client = mocker.AsyncMock() - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mock_client - mocker.patch( - "utils.responses.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) mock_config = mocker.Mock() mock_config.configuration.byok_rag = [] mock_config.mcp_servers = [] @@ -3340,13 +3288,6 @@ async def test_merge_header_conflict_raises_409( @pytest.mark.asyncio async def test_no_tools_uses_prepare_tools(self, mocker: MockerFixture) -> None: """Test that no client tools falls through to prepare_tools.""" - mock_client = mocker.AsyncMock() - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mock_client - mocker.patch( - "utils.responses.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) server_tool = InputToolFileSearch(type="file_search", vector_store_ids=["vs1"]) mock_prepare = mocker.AsyncMock(return_value=[server_tool]) mocker.patch("utils.responses.prepare_tools", new=mock_prepare) @@ -3365,13 +3306,6 @@ async def test_merge_header_no_server_tools_returns_client_only( self, mocker: MockerFixture ) -> None: """Test merge header with no server tools returns client tools unchanged.""" - mock_client = mocker.AsyncMock() - mock_holder = mocker.Mock() - mock_holder.get_client.return_value = mock_client - mocker.patch( - "utils.responses.AsyncLlamaStackClientHolder", - return_value=mock_holder, - ) mock_config = mocker.Mock() mock_config.configuration.byok_rag = [] mock_config.mcp_servers = [] From 4230eea0b44f5c9cee3e7a875457ac98700c03e6 Mon Sep 17 00:00:00 2001 From: Jazzcort Date: Thu, 16 Jul 2026 12:59:18 -0400 Subject: [PATCH 3/3] Add teardown process for container lifecycle test --- .../test_container_lifecycle.py | 49 +++++++++++++++++-- 1 file changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/integration/container_lifecycle/test_container_lifecycle.py b/tests/integration/container_lifecycle/test_container_lifecycle.py index 59bc4a775..273ab9c8d 100644 --- a/tests/integration/container_lifecycle/test_container_lifecycle.py +++ b/tests/integration/container_lifecycle/test_container_lifecycle.py @@ -8,6 +8,7 @@ import time import urllib.error import urllib.request +import warnings from collections.abc import Generator import pytest @@ -18,6 +19,8 @@ CONTAINER_START_TIMEOUT = 300 # 5 minutes for container start CONTAINER_STOP_TIMEOUT = 15 CONTAINER_CLEANUP_TIMEOUT = 10 +IMAGE_CLEANUP_TIMEOUT = 30 +DANGLING_IMAGES_CLEANUP_TIMEOUT = 300 # 5 minutes for dangling images cleanup HEALTH_CHECK_TIMEOUT = 5 PORT_QUERY_TIMEOUT = 5 @@ -25,6 +28,8 @@ HEALTH_CHECK_MAX_ATTEMPTS = 30 NETWORK_BINDING_MAX_ATTEMPTS = 5 +DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME = "lightspeed-llama-stack:local" + @pytest.fixture(scope="session") def container_runtime() -> str: @@ -52,6 +57,39 @@ def container_runtime() -> str: pytest.skip("No container runtime available") +@pytest.fixture(scope="session", autouse=True) +def cleanup_container_artifacts(container_runtime: str) -> Generator[None]: + """Remove container images and dangling layers after all tests complete. + + Parameters + ---------- + container_runtime (str): Container runtime to use. + + Yields + ------ + None + """ + yield + + try: + subprocess.run( + [container_runtime, "rmi", "-f", DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME], + capture_output=True, + timeout=IMAGE_CLEANUP_TIMEOUT, + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: + warnings.warn(f"Image cleanup failed: {e}") + + try: + subprocess.run( + [container_runtime, "image", "prune", "-f"], + capture_output=True, + timeout=DANGLING_IMAGES_CLEANUP_TIMEOUT, + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: + warnings.warn(f"Dangling image cleanup failed: {e}") + + @pytest.fixture(scope="class") def managed_container(container_runtime: str) -> Generator[str, None, None]: """Start container once for entire test class with strict cleanup. @@ -102,7 +140,7 @@ class TestContainerBuild: """Test container image building with idempotency checks.""" def _get_image_id( - self, runtime: str, image_name: str = "lightspeed-llama-stack:local" + self, runtime: str, image_name: str = DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME ) -> str: """Get the unique, immutable Image ID (SHA256). @@ -145,7 +183,7 @@ def test_build_llama_stack_image(self, container_runtime: str) -> None: # Verify image is listed with correct tag result = subprocess.run( - [container_runtime, "images", "lightspeed-llama-stack:local"], + [container_runtime, "images", DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME], capture_output=True, text=True, timeout=PORT_QUERY_TIMEOUT, @@ -548,7 +586,12 @@ def test_clean_removes_image_and_container(self, container_runtime: str) -> None # Verify image is removed result = subprocess.run( - [container_runtime, "images", "-q", "lightspeed-llama-stack:local"], + [ + container_runtime, + "images", + "-q", + DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME, + ], capture_output=True, text=True, timeout=PORT_QUERY_TIMEOUT,