Skip to content
Open
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
4 changes: 2 additions & 2 deletions docs/devel_doc/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down
10 changes: 3 additions & 7 deletions docs/devel_doc/openapi.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 3 additions & 7 deletions docs/models/responses_succ.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/user_doc/byok_guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 3 additions & 7 deletions docs/user_doc/config.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion examples/lightspeed-stack-byok-okp-rag.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 3 additions & 7 deletions src/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.",
)


Expand Down
56 changes: 1 addition & 55 deletions src/utils/responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -127,44 +126,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:
Expand Down Expand Up @@ -229,7 +190,6 @@ async def maybe_get_topic_summary(


async def prepare_tools( # pylint: disable=too-many-arguments,too-many-positional-arguments
client: AsyncLlamaStackClient,
vector_store_ids: Optional[list[str]],
no_tools: Optional[bool],
token: str,
Expand All @@ -239,9 +199,8 @@ 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 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
Expand All @@ -259,22 +218,16 @@ 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)
elif is_tool_rag_enabled:
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)
Expand Down Expand Up @@ -375,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,
Expand Down Expand Up @@ -1754,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,
Expand Down Expand Up @@ -1784,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,
Expand Down Expand Up @@ -1830,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,
Expand Down
49 changes: 46 additions & 3 deletions tests/integration/container_lifecycle/test_container_lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import time
import urllib.error
import urllib.request
import warnings
from collections.abc import Generator

import pytest
Expand All @@ -18,13 +19,17 @@
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

# Retry constants
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:
Expand Down Expand Up @@ -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}")
Comment on lines +74 to +81

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.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Missing check=True suppresses CalledProcessError for cleanup commands.

Both cleanup steps attempt to catch subprocess.CalledProcessError to log a warning on failure, but subprocess.run will never raise this exception unless check=True is passed. Without it, failed commands silently return a non-zero exit code and the warnings are never triggered.

  • tests/integration/container_lifecycle/test_container_lifecycle.py#L74-L81: Add check=True to the rmi subprocess call.
  • tests/integration/container_lifecycle/test_container_lifecycle.py#L83-L90: Add check=True to the image prune subprocess call.
🛠️ Proposed fix
     try:
         subprocess.run(
             [container_runtime, "rmi", "-f", DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME],
             capture_output=True,
             timeout=IMAGE_CLEANUP_TIMEOUT,
+            check=True,
         )
     except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
         warnings.warn(f"Image cleanup failed: {e}")
 
     try:
         subprocess.run(
             [container_runtime, "image", "pr prune", "-f"],
             capture_output=True,
             timeout=DANGLING_IMAGES_CLEANUP_TIMEOUT,
+            check=True,
         )
     except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
         warnings.warn(f"Dangling image cleanup failed: {e}")
🧰 Tools
🪛 ast-grep (0.44.1)

[error] 74-78: Command coming from incoming request
Context: subprocess.run(
[container_runtime, "rmi", "-f", DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME],
capture_output=True,
timeout=IMAGE_CLEANUP_TIMEOUT,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').

(subprocess-from-request)

📍 Affects 1 file
  • tests/integration/container_lifecycle/test_container_lifecycle.py#L74-L81 (this comment)
  • tests/integration/container_lifecycle/test_container_lifecycle.py#L83-L90
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/container_lifecycle/test_container_lifecycle.py` around
lines 74 - 81, Add check=True to both cleanup subprocess.run calls in
tests/integration/container_lifecycle/test_container_lifecycle.py at lines 74-81
(the rmi cleanup) and 83-90 (the image prune cleanup), so non-zero exits raise
CalledProcessError and are handled by the existing warnings.warn logic.


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}")


Comment on lines +60 to +92

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.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

Move session-scoped autouse fixture to conftest.py.

Defining a scope="session", autouse=True fixture in a specific test file is an anti-pattern. It will only execute if this specific test file is discovered and included in the test run. If developers run other integration tests in isolation (e.g., pytest tests/integration/endpoints/), this cleanup fixture will not be registered or run. As per coding guidelines, shared fixtures must be placed in conftest.py.

🧰 Tools
🪛 ast-grep (0.44.1)

[error] 74-78: Command coming from incoming request
Context: subprocess.run(
[container_runtime, "rmi", "-f", DEFAULT_LIGHTSPEED_LLAMA_STACK_IMAGE_NAME],
capture_output=True,
timeout=IMAGE_CLEANUP_TIMEOUT,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').

(subprocess-from-request)


[error] 83-87: Command coming from incoming request
Context: subprocess.run(
[container_runtime, "image", "prune", "-f"],
capture_output=True,
timeout=DANGLING_IMAGES_CLEANUP_TIMEOUT,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').

(subprocess-from-request)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/container_lifecycle/test_container_lifecycle.py` around
lines 60 - 92, Move the session-scoped autouse fixture
cleanup_container_artifacts from the test module into the applicable
integration-test conftest.py so it is registered for all integration test runs.
Preserve its container_runtime dependency, teardown cleanup commands, timeouts,
and warning behavior unchanged.

Source: Coding guidelines

@pytest.fixture(scope="class")
def managed_container(container_runtime: str) -> Generator[str, None, None]:
"""Start container once for entire test class with strict cleanup.
Expand Down Expand Up @@ -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).

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading