Merge pull request #2215 from jrobertboos/lcore-3228

LCORE-3228: Wire Shields into `build_agent`

Author: tisnik

src/app/endpoints/a2a.py

@@ -372,7 +372,7 @@ async def _process_task_streaming(  # pylint: disable=too-many-locals
             )
             responses_params = compaction.params
 
-            agent = build_agent(client, responses_params, configuration.skills)
+            agent = build_agent(client, responses_params, configuration)
         except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e:
             error_response = map_agent_inference_error(e, query_request.model or "")
             logger.error("Error preparing A2A agent: %s", str(e), exc_info=True)

src/utils/agents/query.py

@@ -287,6 +287,7 @@ async def retrieve_agent_response(
     endpoint_path: str,
     _original_input: Optional[ResponseInput] = None,
     no_tools: bool = False,
+    shield_ids: Optional[list[str]] = None,
 ) -> TurnSummary:
     """Retrieve a turn summary from a blocking agent run.
 
@@ -297,6 +298,8 @@ async def retrieve_agent_response(
         endpoint_path: Endpoint path used for metric labeling.
         _original_input: Original user input before the explicit-input rewrite.
         no_tools: Whether to skip tool processing.
+        shield_ids: Optional list of shield names to run for this turn, mirroring
+            ``QueryRequest.shield_ids``. If ``None``, all configured shields run.
     Returns:
         Turn summary for the completed agent run.
 
@@ -316,7 +319,11 @@ async def retrieve_agent_response(
         )
     try:
         agent = build_agent(
-            client, responses_params, configuration.skills, no_tools=no_tools
+            client,
+            responses_params,
+            configuration,
+            shields=shield_ids,
+            no_tools=no_tools,
         )
         logger.debug("Starting agent non-streaming response processing")
         run_result = await agent.run(cast(str, responses_params.input))

src/utils/agents/streaming.py

@@ -121,7 +121,7 @@ async def retrieve_agent_response_generator(
             )
 
         agent = build_agent(
-            context.client, responses_params, configuration.skills, no_tools=no_tools
+            context.client, responses_params, configuration, no_tools=no_tools
         )
 
         return (

src/utils/pydantic_ai_helpers.py

@@ -11,12 +11,21 @@
 from pydantic_ai.capabilities import AbstractCapability, AgentCapability
 from pydantic_ai_skills import SkillsCapability
 
+from configuration import AppConfig
 from models.common.responses.responses_api_params import ResponsesApiParams
 from models.common.tools import CatalogTool, CatalogToolParameter
-from models.config import SkillsConfiguration
+from models.config import (
+    QuestionValidityConfig,
+    RedactionConfig,
+    ShieldConfiguration,
+    SkillsConfiguration,
+)
+from pydantic_ai_lightspeed.capabilities import QuestionValidity
+from pydantic_ai_lightspeed.capabilities.redaction import PiiRedactionCapability
 from pydantic_ai_lightspeed.llamastack import (
     OgxResponsesModel,
 )
+from utils.shields import get_shields_for_request
 
 _AGENT_SKILLS_PROVIDER_ID: Final[str] = "agent-skills"
 _AGENT_SKILLS_TOOLGROUP_ID: Final[str] = "builtin::agent-skills"
@@ -127,20 +136,51 @@ def get_agent_capability_tools(
     return tools
 
 
+def _shield_capability(shield: ShieldConfiguration) -> AgentCapability[object]:
+    """Build the pydantic-ai capability instance for a single configured shield.
+
+    Parameters:
+        shield: A single guardrail shield configuration entry.
+
+    Returns:
+        A ``QuestionValidity`` capability when ``shield.provider_id`` is
+        ``"question_validity"``, or a ``PiiRedactionCapability`` when it is
+        ``"redaction"``.
+
+    Raises:
+        ValueError: If ``shield.config`` doesn't match a known shield config type.
+    """
+    match shield.config:
+        case QuestionValidityConfig():
+            return QuestionValidity(config=shield.config)
+        case RedactionConfig():
+            return PiiRedactionCapability(config=shield.config)
+        case _:
+            raise ValueError(
+                f"Unsupported shield config type for shield '{shield.name}': "
+                f"{type(shield.config).__name__}"
+            )
+
+
 def _agent_capabilities(
     skills: Optional[SkillsConfiguration],
+    shields: Optional[list[ShieldConfiguration]] = None,
     no_tools: bool = False,
 ) -> Optional[list[AgentCapability[object]]]:
     """Assemble pydantic-ai capabilities for an LCS agent.
 
     Args:
         skills: Agent skills configuration from LCS, or None when skills are disabled.
+        shields: Configured guardrail shields (question validity, redaction), or
+            None/empty when no shields are enabled.
         no_tools: When True, omit capabilities that expose a toolset via ``get_toolset()``.
 
     Returns:
         Configured capabilities, or None when no capabilities are enabled.
     """
     capabilities: list[AgentCapability[object]] = []
+    for shield in shields or []:
+        capabilities.append(_shield_capability(shield))
     if skills_capability := _skills_capability(skills):
         capabilities.append(skills_capability)
     if no_tools:
@@ -158,7 +198,8 @@ def _agent_capabilities(
 def build_agent(
     client: AsyncOgxClient | AsyncOGXAsLibraryClient,
     responses_params: ResponsesApiParams,
-    skills: Optional[SkillsConfiguration],
+    config: AppConfig,
+    shields: Optional[list[str]] = None,
     no_tools: bool = False,
 ) -> Agent[None, str]:
     """Build a Pydantic AI agent that mirrors ``responses_params`` on the Llama Stack backend.
@@ -171,14 +212,20 @@ def build_agent(
     Parameters:
         client: Initialized Llama Stack client from ``AsyncOgxClientHolder().get_client()``.
         responses_params: Parameters produced by ``prepare_responses_params`` for this turn.
-        skills: Agent skills configuration from LCS, or None when skills are disabled.
+        config: Application configuration. Agent skills (``config.skills``) and the
+            configured guardrail shields (``config.shields``) are extracted from it.
+        shields: Optional list of shield names to run for this turn, matching each
+            shield's configured ``name``. Mirrors ``QueryRequest.shield_ids``: if
+            ``None``, all shields configured in ``config.shields`` run; an empty
+            list disables all shields.
         no_tools: When True, omit capabilities that expose a toolset via ``get_toolset()``.
 
     Returns:
         ``Agent`` configured for ``await agent.run(...)`` (or streaming) against the same
         stack configuration as ``client.responses.create(**responses_params.model_dump())``.
     """
-    capabilities = _agent_capabilities(skills, no_tools=no_tools)
+    shield_configs = get_shields_for_request(config.shields, shields)
+    capabilities = _agent_capabilities(config.skills, shield_configs, no_tools=no_tools)
 
     model = OgxResponsesModel.from_ogx_client(
         responses_params.model, client, responses_params=responses_params

tests/unit/conftest.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from pathlib import Path
+from typing import Optional
 
 import httpx
 import pytest
@@ -14,7 +15,7 @@
 from configuration import AppConfig
 from constants import DEFAULT_LOGGER_NAME
 from models.common.responses.responses_api_params import ResponsesApiParams
-from models.config import SkillsConfiguration
+from models.config import ShieldConfiguration, SkillsConfiguration
 
 type AgentFixtures = Generator[
     tuple[
@@ -143,3 +144,26 @@ def mock_skills_configuration_fixture(tmp_path: Path) -> SkillsConfiguration:
         encoding="utf-8",
     )
     return SkillsConfiguration(paths=[skills_root])
+
+
+@pytest.fixture(name="make_agent_config")
+def make_agent_config_fixture(
+    mocker: MockerFixture,
+) -> Callable[..., AppConfig]:
+    """Return a factory building a duck-typed AppConfig stand-in for build_agent.
+
+    ``build_agent`` only reads ``config.skills`` and ``config.shields`` off the
+    config object it receives, so tests can pass a lightweight mock instead of
+    a fully-initialized ``AppConfig``.
+    """
+
+    def _make(
+        skills: Optional[SkillsConfiguration] = None,
+        shields: Optional[list[ShieldConfiguration]] = None,
+    ) -> AppConfig:
+        config = mocker.Mock()
+        config.skills = skills
+        config.shields = shields or []
+        return config
+
+    return _make