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
57 changes: 34 additions & 23 deletions docs/devel_doc/openapi.json

Large diffs are not rendered by default.

8 changes: 4 additions & 4 deletions src/models/api/requests/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ class QueryRequest(BaseModel):
generate_topic_summary: Whether to generate topic summary for new conversations.
media_type: The optional media type for response format (application/json or text/plain).
vector_store_ids: The optional list of specific vector store IDs to query for RAG.
shield_ids: The optional list of safety shield IDs to apply.
shield_ids: The optional list of configured shield names to apply.
solr: Optional Solr inline RAG options (mode, filters) or legacy filter-only dict.
"""

Expand Down Expand Up @@ -105,9 +105,9 @@ class QueryRequest(BaseModel):

shield_ids: Optional[list[str]] = Field(
None,
description="Optional list of safety shield IDs to apply. "
"If None, all configured shields are used. ",
examples=["llama-guard", "custom-shield"],
description="Optional list of configured shield names to apply. "
"If None, all configured shields are used.",
examples=["topic-guard", "pii-redaction"],
)

solr: Optional[SolrVectorSearchRequest] = Field(
Expand Down
4 changes: 2 additions & 2 deletions src/models/api/requests/responses_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,8 @@ class ResponsesRequest(BaseModel):
calls, MCP tools). Defaults to all tools available to the model.
generate_topic_summary: LCORE-specific flag indicating whether to generate a
topic summary for new conversations. Defaults to True.
shield_ids: LCORE-specific list of safety shield IDs to apply. If None, all
configured shields are used.
shield_ids: LCORE-specific list of configured shield names to apply.
If None, all configured shields are used.
solr: Optional Solr inline RAG options (mode, filters) or legacy filter-only dict.
"""

Expand Down
6 changes: 4 additions & 2 deletions src/models/api/responses/successful/catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,8 @@ class ShieldsResponse(AbstractSuccessfulResponse):
"shields": [
{
"name": "question-validity",
"type": "question_validity",
"provider_id": "question_validity",
"type": "shield",
"config": {
"model_id": "openai/gpt-4o-mini",
"model_prompt": "Is this question valid?",
Expand All @@ -102,7 +103,8 @@ class ShieldsResponse(AbstractSuccessfulResponse):
},
{
"name": "pii-redaction",
"type": "redaction",
"provider_id": "redaction",
"type": "shield",
"config": {
"rules": [
{
Expand Down
14 changes: 11 additions & 3 deletions src/models/common/shields.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,20 @@
class CatalogShield(BaseModel):
"""Shield entry in the ``/shields`` catalog response.

Mirrors the LCS-owned ``ShieldConfiguration`` shape (name / type / config).
Attributes:
name: Unique, user-facing name identifying this shield instance.
provider_id: Shield provider / type discriminator.
type: Catalog entry type; always shield.
config: Type-specific shield configuration.
"""

name: str = Field(description="Unique, user-facing name of the shield instance")
type: Literal["question_validity", "redaction"] = Field(
description="Shield type discriminator",
provider_id: Literal["question_validity", "redaction"] = Field(
description="Shield provider / type discriminator",
)
type: Literal["shield"] = Field(
default="shield",
description="Catalog entry type; always shield",
)
config: dict[str, Any] = Field(
description="Type-specific shield configuration",
Expand Down
22 changes: 11 additions & 11 deletions src/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -2595,7 +2595,7 @@ class QuestionValidityShieldConfiguration(ConfigurationBase):

Attributes:
name: Unique, user-facing name identifying this shield instance.
type: Discriminator identifying this as a question-validity shield.
provider_id: Discriminator identifying this as a question-validity shield.
config: Question-validity-specific configuration.
"""

Expand All @@ -2605,9 +2605,9 @@ class QuestionValidityShieldConfiguration(ConfigurationBase):
description="Unique, user-facing name identifying this shield instance.",
)

type: Literal["question_validity"] = Field(
provider_id: Literal["question_validity"] = Field(
...,
title="Shield type",
title="Shield provider id",
description="Discriminator identifying this as a question-validity shield.",
)

Expand All @@ -2623,7 +2623,7 @@ class RedactionShieldConfiguration(ConfigurationBase):

Attributes:
name: Unique, user-facing name identifying this shield instance.
type: Discriminator identifying this as a redaction shield.
provider_id: Discriminator identifying this as a redaction shield.
config: Redaction-specific configuration.
"""

Expand All @@ -2633,9 +2633,9 @@ class RedactionShieldConfiguration(ConfigurationBase):
description="Unique, user-facing name identifying this shield instance.",
)

type: Literal["redaction"] = Field(
provider_id: Literal["redaction"] = Field(
...,
title="Shield type",
title="Shield provider id",
description="Discriminator identifying this as a redaction shield.",
)

Expand All @@ -2648,11 +2648,11 @@ class RedactionShieldConfiguration(ConfigurationBase):

ShieldConfiguration = Annotated[
QuestionValidityShieldConfiguration | RedactionShieldConfiguration,
Field(discriminator="type"),
Field(discriminator="provider_id"),
]
"""Configuration for a single named guardrail shield (question validity or redaction).

A discriminated union on ``type``: Pydantic selects
A discriminated union on ``provider_id``: Pydantic selects
``QuestionValidityShieldConfiguration`` or ``RedactionShieldConfiguration``
and validates ``config`` against the matching model.
"""
Expand Down Expand Up @@ -2847,9 +2847,9 @@ class Configuration(ConfigurationBase):
default_factory=list,
title="Shields configuration",
description="List of pydantic-ai-lightspeed agent guardrail shields "
"(question validity and PII redaction). Each entry has a unique 'name', "
"a 'type' ('question_validity' or 'redaction'), and a type-specific "
"'config'.",
"(question validity and PII redaction). Each entry has a unique "
"'name', a 'provider_id' ('question_validity' or 'redaction'), "
"and a type-specific 'config'.",
)

@model_validator(mode="after")
Expand Down
18 changes: 8 additions & 10 deletions src/utils/shields.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,22 +149,20 @@ def get_shields_for_request(
shields: list[ShieldConfiguration],
shield_ids: Optional[list[str]] = None,
) -> list[ShieldConfiguration]:
"""Return configured shields, optionally filtered by request ``shield_ids``.

Shield identifiers in the request map to each shield's configured ``name``.
"""Return configured shields, optionally filtered by request shield_ids.

Args:
shields: Configured LCS shields.
shield_ids: Optional list of shield names. If ``None``, all ``shields``
are returned. An empty list skips all shields. Otherwise only
shields whose ``name`` is in this list are returned.
shield_ids: Optional list of shield names. If None, all shields are
returned. An empty list skips all shields. Otherwise only shields
whose name is in this list are returned.

Returns:
list[ShieldConfiguration]: Shield configurations to run for this request.

Raises:
HTTPException: 404 if ``shield_ids`` is provided and any requested
shield name is not present in ``shields``.
HTTPException: 404 if shield_ids is provided and any requested shield
name is not present in shields.
"""
if shield_ids is None:
return list(shields)
Expand All @@ -173,8 +171,8 @@ def get_shields_for_request(
return []

requested = set(shield_ids)
configured_ids = {shield.name for shield in shields}
missing = requested - configured_ids
configured_names = {shield.name for shield in shields}
missing = requested - configured_names
if missing:
response = NotFoundResponse(
resource=f"Shield{'s' if len(missing) > 1 else ''}",
Expand Down
9 changes: 9 additions & 0 deletions tests/e2e/configuration/library-mode/lightspeed-stack.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -35,3 +35,12 @@ byok_rag:
rag:
tool:
- e2e-test-docs

shields:
- name: pii-redaction
provider_id: redaction
config:
rules:
- pattern: '\d+'
replacement: '[NUM]'

10 changes: 9 additions & 1 deletion tests/e2e/configuration/server-mode/lightspeed-stack.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,12 @@ byok_rag:

rag:
tool:
- e2e-test-docs
- e2e-test-docs

shields:
- name: pii-redaction
provider_id: redaction
config:
rules:
- pattern: '\d+'
replacement: '[NUM]'
1 change: 0 additions & 1 deletion tests/e2e/features/info.feature
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ Feature: Info tests
And The body of the response has proper name Lightspeed Core Service (LCS) and version 0.6.0rc2
And The body of the response has llama-stack version 1.0.2

@skip
Scenario: Check if shields endpoint is working
When I access REST API endpoint "shields" using HTTP GET method
Then The status code of the response is 200
Expand Down
12 changes: 4 additions & 8 deletions tests/e2e/features/steps/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,15 +58,11 @@ def check_shield_structure(context: Context) -> None:
# Validate structure and values
assert found_shield["type"] == "shield", "type should be 'shield'"
assert (
found_shield["provider_id"] == "llama-guard"
), "provider_id should be 'llama-guard'"
assert found_shield["provider_resource_id"] == "openai/gpt-4o-mini", (
f"provider_resource_id should be 'openai/gpt-4o-mini', "
f"but is '{found_shield['provider_resource_id']}'"
found_shield["provider_id"] == "redaction"
), "provider_id should be 'redaction'"
assert found_shield["name"] == "pii-redaction", (
f"name should be 'pii-redaction', " f"but is '{found_shield['name']}'"
)
assert (
found_shield["identifier"] == "llama-guard"
), f"identifier should be 'llama-guard', but is '{found_shield["identifier"]}'"


@then("The response contains {count:d} tools listed for provider {provider_name}")
Expand Down
10 changes: 6 additions & 4 deletions tests/unit/app/endpoints/test_shields.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ async def test_shields_endpoint_handler_configured_shields(
config_dict["shields"] = [
{
"name": "question-validity",
"type": "question_validity",
"provider_id": "question_validity",
"config": {
"model_id": "openai/gpt-4o-mini",
"model_prompt": "Is this question valid?",
Expand All @@ -109,7 +109,7 @@ async def test_shields_endpoint_handler_configured_shields(
},
{
"name": "pii-redaction",
"type": "redaction",
"provider_id": "redaction",
"config": {
"rules": [
{
Expand All @@ -132,8 +132,10 @@ async def test_shields_endpoint_handler_configured_shields(
assert isinstance(response, ShieldsResponse)
assert len(response.shields) == 2
assert response.shields[0].name == "question-validity"
assert response.shields[0].type == "question_validity"
assert response.shields[0].provider_id == "question_validity"
assert response.shields[0].type == "shield"
assert response.shields[0].config["model_id"] == "openai/gpt-4o-mini"
assert response.shields[1].name == "pii-redaction"
assert response.shields[1].type == "redaction"
assert response.shields[1].provider_id == "redaction"
assert response.shields[1].type == "shield"
assert response.shields[1].config["rules"][0]["replacement"] == "[REDACTED]"
Loading
Loading