Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/byok_guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -307,7 +307,7 @@ rag:
- okp # include OKP context inline

# Tool RAG: the LLM can call file_search to retrieve context on demand
# Omit to use all registered BYOK stores (backward compatibility)
# If omitted, tool RAG is disabled. If both tool and inline are omitted, all registered stores are used as fallback
tool:
- my-docs # expose this BYOK store as the file_search tool
- okp # expose OKP as the file_search tool
Expand Down
17 changes: 5 additions & 12 deletions docs/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -11033,25 +11033,18 @@
"description": "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": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"items": {
"type": "string"
},
"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)."
}
},
"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 ``None`` which means all registered vector stores\n are used (identical to the previous ``tool.byok.enabled = True`` default)."
"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."
Comment on lines 11035 to +11047

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Resolve conflicting rag.tool behavior in the OpenAPI schema docs.

Line 11041 states that omitting tool enables all BYOK stores, but Line 11047 states tool defaults to [] and “all stores” is only a fallback when both inline and tool are empty. These semantics conflict and will mislead API/config consumers.

Suggested doc fix
                     "tool": {
                         "items": {
                             "type": "string"
                         },
                         "type": "array",
+                        "default": [],
                         "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. Defaults to an empty list (tool RAG disabled). If both rag.inline and rag.tool are empty, all stores available to llama-stack are registered as fallback."
                     }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
"tool": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"items": {
"type": "string"
},
"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)."
}
},
"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 ``None`` which means all registered vector stores\n are used (identical to the previous ``tool.byok.enabled = True`` default)."
"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."
"tool": {
"items": {
"type": "string"
},
"type": "array",
"default": [],
"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. Defaults to an empty list (tool RAG disabled). If both rag.inline and rag.tool are empty, all stores available to llama-stack are registered as fallback."
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@docs/openapi.json` around lines 11035 - 11047, Update the RagConfiguration
description to resolve the contradiction about the "tool" field: clarify that
"tool" defaults to [] (no tool RAG) and does not by itself enable all BYOK
stores, and then explicitly state that the special-case fallback of registering
all stores occurs only when both "inline" and "tool" are empty; keep the note
about the special ID "okp" and references to byok_rag and constants "okp" intact
so readers can find the unique IDs and providers.

},
"ReadinessResponse": {
"properties": {
Expand Down
12 changes: 7 additions & 5 deletions src/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from enum import Enum
from functools import cached_property
from pathlib import Path
from typing import Any, Optional, Literal, Self
from re import Pattern
from typing import Any, Literal, Optional, Self

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check minimum Python version in project configuration
echo "=== Checking pyproject.toml for Python version ==="
cat pyproject.toml 2>/dev/null | grep -A5 "python" | head -20

echo ""
echo "=== Checking setup.py/setup.cfg for Python version ==="
cat setup.py 2>/dev/null | grep -i python | head -10
cat setup.cfg 2>/dev/null | grep -i python | head -10

echo ""
echo "=== Checking if typing_extensions.Self was previously used ==="
git log -1 --format="%H" -- src/models/config.py 2>/dev/null
git show HEAD:src/models/config.py 2>/dev/null | head -15

Repository: lightspeed-core/lightspeed-stack

Length of output: 865


🏁 Script executed:

# Check if Self import was added in this PR
echo "=== Check git diff for src/models/config.py ==="
git diff HEAD -- src/models/config.py 2>/dev/null | head -50

echo ""
echo "=== Check actual usage of Self in the file ==="
rg "Self" src/models/config.py -B2 -A2

Repository: lightspeed-core/lightspeed-stack

Length of output: 4681


Change Self import to use typing_extensions.Self for consistency with coding guidelines.

The file imports Self from typing (line 10), but the coding guidelines specify using typing_extensions.Self for model validators in type annotations. This file has 18+ @model_validator methods that all use Self as return type, so update the import to comply with the guideline.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/models/config.py` at line 10, Replace the import of Self from typing with
Self from typing_extensions: update the import line that currently reads "from
typing import Any, Literal, Optional, Self" to import Self from
typing_extensions instead, leaving other typing imports unchanged; then run a
quick search to ensure all `@model_validator` methods and their return annotations
(instances using Self) still type-check and update any import references if
necessary.


import jsonpath_ng
import yaml
Expand Down Expand Up @@ -1704,8 +1704,10 @@ class RagConfiguration(ConfigurationBase):

Backward compatibility:
- ``inline`` defaults to ``[]`` (no inline RAG).
- ``tool`` defaults to ``None`` which means all registered vector stores
are used (identical to the previous ``tool.byok.enabled = True`` default).
- ``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.
"""

inline: list[str] = Field(
Expand All @@ -1715,8 +1717,8 @@ class RagConfiguration(ConfigurationBase):
f"Use '{constants.OKP_RAG_ID}' to enable OKP inline RAG. Empty by default (no inline RAG).",
)

tool: Optional[list[str]] = Field(
default=None,
tool: list[str] = Field(
default_factory=list,
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. "
Expand Down
22 changes: 15 additions & 7 deletions src/utils/responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,18 +174,26 @@ async def prepare_tools( # pylint: disable=too-many-arguments,too-many-position
return None

toolgroups: list[InputTool] = []

# 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.
effective_ids: list[str] = []

# 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: 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(vector_store_ids, byok_rags)
elif is_tool_rag_enabled:
effective_ids = resolve_vector_store_ids(
configuration.configuration.rag.tool, byok_rags
)
else:
elif not is_inline_rag_enabled:
effective_ids = await get_vector_store_ids(client, None)

# Add RAG tools if vector stores are available
Expand Down
10 changes: 5 additions & 5 deletions tests/unit/models/config/test_dump_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ def test_dump_configuration(tmp_path: Path) -> None:
"azure_entra_id": None,
"rag": {
"inline": [],
"tool": None,
"tool": [],
},
"okp": {
"offline": True,
Expand Down Expand Up @@ -559,7 +559,7 @@ def test_dump_configuration_with_quota_limiters(tmp_path: Path) -> None:
"azure_entra_id": None,
"rag": {
"inline": [],
"tool": None,
"tool": [],
},
"okp": {
"offline": True,
Expand Down Expand Up @@ -788,7 +788,7 @@ def test_dump_configuration_with_quota_limiters_different_values(
"azure_entra_id": None,
"rag": {
"inline": [],
"tool": None,
"tool": [],
},
"okp": {
"offline": True,
Expand Down Expand Up @@ -992,7 +992,7 @@ def test_dump_configuration_byok(tmp_path: Path) -> None:
"azure_entra_id": None,
"rag": {
"inline": [],
"tool": None,
"tool": [],
},
"okp": {
"offline": True,
Expand Down Expand Up @@ -1181,7 +1181,7 @@ def test_dump_configuration_pg_namespace(tmp_path: Path) -> None:
"azure_entra_id": None,
"rag": {
"inline": [],
"tool": None,
"tool": [],
},
"okp": {
"offline": True,
Expand Down
10 changes: 5 additions & 5 deletions tests/unit/models/config/test_rag_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@ def test_default_values(self) -> None:
"""Test that RagConfiguration has correct default values."""
config = RagConfiguration()
assert config.inline == []
assert config.tool is None
assert config.tool == []

def test_inline_with_byok_ids(self) -> None:
"""Test inline list with BYOK rag IDs."""
config = RagConfiguration(inline=["store-1", "store-2"])
assert config.inline == ["store-1", "store-2"]
assert config.tool is None
assert config.tool == []

def test_inline_with_okp_rag(self) -> None:
"""Test inline list including the special OKP ID."""
Expand All @@ -45,10 +45,10 @@ def test_tool_empty_list(self) -> None:
config = RagConfiguration(tool=[])
assert config.tool == []

def test_tool_none_means_all_stores(self) -> None:
"""Test that tool=None (default) means all registered stores are used."""
def test_tool_default_is_empty_list(self) -> None:
"""Test that tool defaults to an empty list."""
config = RagConfiguration()
assert config.tool is None
assert config.tool == []

def test_no_unknown_fields_allowed(self) -> None:
"""Test that RagConfiguration rejects unknown fields."""
Expand Down
121 changes: 120 additions & 1 deletion tests/unit/utils/test_responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -1026,6 +1026,8 @@ async def test_translates_byok_ids_in_prepare_tools(
mock_byok_rag.vector_db_id = "vs-001"
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, ["ocp_docs"], False, "token")
Expand All @@ -1045,6 +1047,8 @@ async def test_passes_through_unknown_ids_in_prepare_tools(
# Configure empty BYOK RAG
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, ["raw-internal-id"], False, "token")
Expand Down Expand Up @@ -1073,7 +1077,8 @@ 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
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")
Expand All @@ -1082,6 +1087,120 @@ async def test_does_not_translate_when_ids_fetched_from_llama_stack(
assert result[0].vector_store_ids == ["vs-internal"]


class TestPrepareToolsVectorStoreResolution:
"""Tests for vector store ID resolution priority in prepare_tools."""

@pytest.mark.asyncio
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()
mock_config.configuration.byok_rag = []
mock_config.configuration.rag.tool = ["rag-tool-id-1", "rag-tool-id-2"]
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 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()
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]
mock_config.configuration.rag.tool = ["ocp_docs"]
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 result[0].vector_store_ids == ["vs-001"]
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."""
mock_client = mocker.AsyncMock()
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 = [
"inline-store-id"
] # inline is configured
mocker.patch("utils.responses.configuration", mock_config)

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()

@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()
mock_config.configuration.byok_rag = []
mock_config.configuration.rag.tool = ["config-id-1"]
mock_config.configuration.rag.inline = []
mocker.patch("utils.responses.configuration", mock_config)

result = await prepare_tools(mock_client, ["request-id-1"], False, "token")

assert result is not None
assert result[0].vector_store_ids == ["request-id-1"]
mock_client.vector_stores.list.assert_not_called()

@pytest.mark.asyncio
async def test_all_registered_dbs_used_when_neither_tool_nor_inline_configured(
self, mocker: MockerFixture
) -> None:
"""Test fallback to all registered vector DBs when neither rag.tool nor rag.inline are set."""
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 result[0].vector_store_ids == ["vs-registered"]
mock_client.vector_stores.list.assert_called_once()


class TestPrepareResponsesParams:
"""Tests for prepare_responses_params function."""

Expand Down
Loading