diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index 16fa60100..9c8f5e07c 100644 --- a/src/utils/agents/query.py +++ b/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)) diff --git a/src/utils/pydantic_ai_helpers.py b/src/utils/pydantic_ai_helpers.py index bc9515d69..49e2b9736 100644 --- a/src/utils/pydantic_ai_helpers.py +++ b/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 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 5ed7d2e61..b16ea951e 100644 --- a/tests/unit/conftest.py +++ b/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 diff --git a/tests/unit/utils/test_pydantic_ai.py b/tests/unit/utils/test_pydantic_ai.py index 4f2716d35..bc477ce8e 100644 --- a/tests/unit/utils/test_pydantic_ai.py +++ b/tests/unit/utils/test_pydantic_ai.py @@ -2,21 +2,46 @@ # pylint: disable=protected-access +from collections.abc import Callable + import httpx +import pytest +from fastapi import HTTPException from ogx.core.library_client import AsyncOGXAsLibraryClient from ogx_client import AsyncOgxClient from pydantic_ai_skills import SkillsCapability from pytest_mock import MockerFixture +from configuration import AppConfig from models.common.responses.responses_api_params import ResponsesApiParams -from models.config import SkillsConfiguration +from models.config import ( + QuestionValidityConfig, + QuestionValidityShieldConfiguration, + RedactionConfig, + RedactionShieldConfiguration, + SkillsConfiguration, +) +from pydantic_ai_lightspeed.capabilities import QuestionValidity +from pydantic_ai_lightspeed.capabilities.redaction import PiiRedactionCapability from utils.pydantic_ai_helpers import ( _agent_capabilities, + _shield_capability, _skills_capability, build_agent, get_agent_capability_tools, ) +_QUESTION_VALIDITY_MODULE = ( + "pydantic_ai_lightspeed.capabilities.question_validity._capability" +) + + +@pytest.fixture(autouse=True) +def _mock_question_validity_model(mocker: MockerFixture) -> None: + """Avoid constructing a real client/model when building QuestionValidity.""" + mocker.patch(f"{_QUESTION_VALIDITY_MODULE}.AsyncOgxClientHolder") + mocker.patch(f"{_QUESTION_VALIDITY_MODULE}.OgxResponsesModel.from_ogx_client") + class TestSkillsCapability: """Tests for _skills_capability.""" @@ -39,6 +64,49 @@ def test_returns_capability_for_configured_paths( assert list(capability.toolset.skills) == ["test-skill"] +class TestShieldCapability: + """Tests for _shield_capability.""" + + def test_question_validity_shield_builds_question_validity_capability( + self, + ) -> None: + """Test that a question_validity shield builds a QuestionValidity capability.""" + shield = QuestionValidityShieldConfiguration( + name="topic-guard", + provider_id="question_validity", + config=QuestionValidityConfig(model_id="test-model"), + ) + + capability = _shield_capability(shield) + + assert isinstance(capability, QuestionValidity) + assert capability.config is shield.config + + def test_redaction_shield_builds_pii_redaction_capability(self) -> None: + """Test that a redaction shield builds a PiiRedactionCapability.""" + shield = RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ) + + capability = _shield_capability(shield) + + assert isinstance(capability, PiiRedactionCapability) + assert capability.config is shield.config + + def test_unsupported_config_type_raises_value_error( + self, mocker: MockerFixture + ) -> None: + """Test that an unrecognized shield config type raises ValueError.""" + shield = mocker.Mock(name="bad-shield") + shield.name = "bad-shield" + shield.config = object() + + with pytest.raises(ValueError, match="Unsupported shield config type"): + _shield_capability(shield) + + class TestAgentCapabilities: """Tests for _agent_capabilities.""" @@ -46,6 +114,7 @@ def test_returns_none_when_no_capabilities_configured(self) -> None: """Test that missing configuration yields None for Agent construction.""" assert _agent_capabilities(None) is None assert _agent_capabilities(SkillsConfiguration(paths=[])) is None + assert _agent_capabilities(None, shields=[]) is None def test_returns_skills_capability_when_configured( self, mock_skills_configuration: SkillsConfiguration @@ -56,11 +125,55 @@ def test_returns_skills_capability_when_configured( assert len(capabilities) == 1 assert isinstance(capabilities[0], SkillsCapability) + def test_returns_shield_capabilities_when_configured(self) -> None: + """Test that configured shields are included in the capability list.""" + shields = [ + QuestionValidityShieldConfiguration( + name="topic-guard", + provider_id="question_validity", + config=QuestionValidityConfig(model_id="test-model"), + ), + RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ), + ] + + capabilities = _agent_capabilities(None, shields=shields) or [] + + assert len(capabilities) == 2 + assert isinstance(capabilities[0], QuestionValidity) + assert isinstance(capabilities[1], PiiRedactionCapability) + + def test_combines_shields_and_skills( + self, mock_skills_configuration: SkillsConfiguration + ) -> None: + """Test that shield and skill capabilities are both included together.""" + shields = [ + RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ), + ] + + capabilities = ( + _agent_capabilities(mock_skills_configuration, shields=shields) or [] + ) + + capability_types = {type(capability) for capability in capabilities} + assert capability_types == {PiiRedactionCapability, SkillsCapability} + class TestBuildAgent: """Tests for the build_agent factory function.""" - def test_returns_agent_with_correct_model(self, mocker: MockerFixture) -> None: + def test_returns_agent_with_correct_model( + self, + mocker: MockerFixture, + make_agent_config: Callable[..., AppConfig], + ) -> None: """Test that build_agent returns an Agent with the specified model name.""" mock_client = mocker.Mock() mock_client.base_url = "http://localhost:8321" @@ -82,11 +195,15 @@ def test_returns_agent_with_correct_model(self, mocker: MockerFixture) -> None: mock_params.store = False mock_params.previous_response_id = None - agent = build_agent(mock_client, mock_params, None) + agent = build_agent(mock_client, mock_params, make_agent_config()) assert agent is not None - def test_agent_has_instructions(self, mocker: MockerFixture) -> None: + def test_agent_has_instructions( + self, + mocker: MockerFixture, + make_agent_config: Callable[..., AppConfig], + ) -> None: """Test that build_agent passes instructions to the Agent.""" mock_client = mocker.Mock() mock_client.base_url = "http://localhost:8321" @@ -105,11 +222,15 @@ def test_agent_has_instructions(self, mocker: MockerFixture) -> None: mock_params.store = False mock_params.previous_response_id = None - agent = build_agent(mock_client, mock_params, None) + agent = build_agent(mock_client, mock_params, make_agent_config()) assert "You are a helpful assistant." in agent._instructions - def test_agent_with_library_client(self, mocker: MockerFixture) -> None: + def test_agent_with_library_client( + self, + mocker: MockerFixture, + make_agent_config: Callable[..., AppConfig], + ) -> None: """Test that build_agent works with a library client.""" mock_lib_client = mocker.Mock(spec=AsyncOGXAsLibraryClient) mock_lib_client.provider_data = None @@ -128,7 +249,7 @@ def test_agent_with_library_client(self, mocker: MockerFixture) -> None: mock_params.store = True mock_params.previous_response_id = None - agent = build_agent(mock_lib_client, mock_params, None) + agent = build_agent(mock_lib_client, mock_params, make_agent_config()) assert agent is not None @@ -137,12 +258,13 @@ def test_agent_includes_skills_capability_when_configured( mock_client: AsyncOgxClient, mock_params: ResponsesApiParams, mock_skills_configuration: SkillsConfiguration, + make_agent_config: Callable[..., AppConfig], ) -> None: - """Test that build_agent attaches SkillsCapability when skills are passed.""" + """Test that build_agent attaches SkillsCapability when skills are configured.""" agent = build_agent( mock_client, mock_params, - mock_skills_configuration, + make_agent_config(skills=mock_skills_configuration), ) capability_types = { @@ -154,26 +276,73 @@ def test_agent_has_no_skills_capability_when_not_configured( self, mock_client: AsyncOgxClient, mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], ) -> None: - """Test that build_agent omits SkillsCapability when skills are not passed.""" - agent = build_agent(mock_client, mock_params, None) + """Test that build_agent omits SkillsCapability when skills are not configured.""" + agent = build_agent(mock_client, mock_params, make_agent_config()) capability_types = { type(capability) for capability in agent._root_capability.capabilities } assert SkillsCapability not in capability_types + def test_agent_includes_shield_capabilities_when_configured( + self, + mock_client: AsyncOgxClient, + mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], + ) -> None: + """Test that build_agent attaches shield capabilities configured for the app.""" + shields = [ + QuestionValidityShieldConfiguration( + name="topic-guard", + provider_id="question_validity", + config=QuestionValidityConfig(model_id="test-model"), + ), + RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ), + ] + + agent = build_agent( + mock_client, mock_params, make_agent_config(shields=shields) + ) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert QuestionValidity in capability_types + assert PiiRedactionCapability in capability_types + + def test_agent_has_no_shield_capabilities_when_not_configured( + self, + mock_client: AsyncOgxClient, + mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], + ) -> None: + """Test that build_agent omits shield capabilities when none are configured.""" + agent = build_agent(mock_client, mock_params, make_agent_config()) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert QuestionValidity not in capability_types + assert PiiRedactionCapability not in capability_types + def test_agent_excludes_tool_capabilities_when_no_tools( self, mock_client: AsyncOgxClient, mock_params: ResponsesApiParams, mock_skills_configuration: SkillsConfiguration, + make_agent_config: Callable[..., AppConfig], ) -> None: """Test that build_agent omits tool-bearing capabilities when no_tools=True.""" agent = build_agent( mock_client, mock_params, - mock_skills_configuration, + make_agent_config(skills=mock_skills_configuration), no_tools=True, ) @@ -182,6 +351,76 @@ def test_agent_excludes_tool_capabilities_when_no_tools( } assert SkillsCapability not in capability_types + def test_agent_filters_shields_by_name( + self, + mock_client: AsyncOgxClient, + mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], + ) -> None: + """Test that the shields param filters configured shields by name. + + Mirrors ``QueryRequest.shield_ids``: only shields whose ``name`` is in + the requested list should be attached to the agent. + """ + shields = [ + QuestionValidityShieldConfiguration( + name="topic-guard", + provider_id="question_validity", + config=QuestionValidityConfig(model_id="test-model"), + ), + RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ), + ] + config = make_agent_config(shields=shields) + + agent = build_agent(mock_client, mock_params, config, shields=["pii-guard"]) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert PiiRedactionCapability in capability_types + assert QuestionValidity not in capability_types + + def test_agent_disables_all_shields_with_empty_list( + self, + mock_client: AsyncOgxClient, + mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], + ) -> None: + """Test that an empty shields list disables all configured shields.""" + shields = [ + RedactionShieldConfiguration( + name="pii-guard", + provider_id="redaction", + config=RedactionConfig(rules=[]), + ), + ] + config = make_agent_config(shields=shields) + + agent = build_agent(mock_client, mock_params, config, shields=[]) + + capability_types = { + type(capability) for capability in agent._root_capability.capabilities + } + assert PiiRedactionCapability not in capability_types + + def test_agent_raises_not_found_for_unknown_shield_name( + self, + mock_client: AsyncOgxClient, + mock_params: ResponsesApiParams, + make_agent_config: Callable[..., AppConfig], + ) -> None: + """Test that requesting an unconfigured shield name raises HTTPException.""" + config = make_agent_config(shields=[]) + + with pytest.raises(HTTPException) as exc_info: + build_agent(mock_client, mock_params, config, shields=["missing-shield"]) + + assert exc_info.value.status_code == 404 + class TestGetAgentCapabilityTools: """Tests for get_agent_capability_tools."""