Merge pull request #2256 from asimurka/feature/shields-config

tisnik · web-flow · commit 226e4ef20b34 · 2026-07-27T13:00:06.000+02:00
LCORE-3269: Restructure LCORE shield configuration
diff --git a/docs/devel_doc/openapi.json b/docs/devel_doc/openapi.json
diff --git a/src/models/api/requests/query.py b/src/models/api/requests/query.py
@@ -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.
     """
 
@@ -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(
diff --git a/src/models/api/requests/responses_openai.py b/src/models/api/requests/responses_openai.py
@@ -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.
     """
 
diff --git a/src/models/api/responses/successful/catalog.py b/src/models/api/responses/successful/catalog.py
@@ -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?",
@@ -102,7 +103,8 @@ class ShieldsResponse(AbstractSuccessfulResponse):
                         },
                         {
                             "name": "pii-redaction",
-                            "type": "redaction",
+                            "provider_id": "redaction",
+                            "type": "shield",
                             "config": {
                                 "rules": [
                                     {
diff --git a/src/models/common/shields.py b/src/models/common/shields.py
@@ -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",
diff --git a/src/models/config.py b/src/models/config.py
@@ -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.
     """
 
@@ -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.",
     )
 
@@ -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.
     """
 
@@ -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.",
     )
 
@@ -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.
 """
@@ -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")
diff --git a/src/utils/shields.py b/src/utils/shields.py
@@ -133,22 +133,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)
@@ -157,8 +155,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 ''}",
diff --git a/tests/e2e/configuration/library-mode/lightspeed-stack.yaml b/tests/e2e/configuration/library-mode/lightspeed-stack.yaml
@@ -35,3 +35,12 @@ byok_rag:
 rag:
   tool:
     - e2e-test-docs
+
+shields:
+  - name: pii-redaction
+    provider_id: redaction
+    config:
+      rules:
+        - pattern: '\d+'
+          replacement: '[NUM]'
+
diff --git a/tests/e2e/configuration/server-mode/lightspeed-stack.yaml b/tests/e2e/configuration/server-mode/lightspeed-stack.yaml
@@ -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]'
diff --git a/tests/e2e/features/info.feature b/tests/e2e/features/info.feature
@@ -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
diff --git a/tests/e2e/features/steps/info.py b/tests/e2e/features/steps/info.py
@@ -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}")
diff --git a/tests/unit/app/endpoints/test_shields.py b/tests/unit/app/endpoints/test_shields.py
@@ -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?",
@@ -109,7 +109,7 @@ async def test_shields_endpoint_handler_configured_shields(
         },
         {
             "name": "pii-redaction",
-            "type": "redaction",
+            "provider_id": "redaction",
             "config": {
                 "rules": [
                     {
@@ -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]"
diff --git a/tests/unit/models/config/test_shields_configuration.py b/tests/unit/models/config/test_shields_configuration.py
diff --git a/tests/unit/models/responses/test_successful_responses.py b/tests/unit/models/responses/test_successful_responses.py
diff --git a/tests/unit/test_configuration.py b/tests/unit/test_configuration.py
diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py

Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,8 @@ class ShieldsResponse(AbstractSuccessfulResponse):`
`91`	`91`	`"shields": [`
`92`	`92`	`{`
`93`	`93`	`"name": "question-validity",`
`94`		`- "type": "question_validity",`
	`94`	`+ "provider_id": "question_validity",`
	`95`	`+ "type": "shield",`
`95`	`96`	`"config": {`
`96`	`97`	`"model_id": "openai/gpt-4o-mini",`
`97`	`98`	`"model_prompt": "Is this question valid?",`
`@@ -102,7 +103,8 @@ class ShieldsResponse(AbstractSuccessfulResponse):`
`102`	`103`	`},`
`103`	`104`	`{`
`104`	`105`	`"name": "pii-redaction",`
`105`		`- "type": "redaction",`
	`106`	`+ "provider_id": "redaction",`
	`107`	`+ "type": "shield",`
`106`	`108`	`"config": {`
`107`	`109`	`"rules": [`
`108`	`110`	`{`