diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index 9ac1b7163..ff8cda209 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -60,7 +60,7 @@ from log import get_logger from models.api.requests import QueryRequest from models.config import Action -from utils.agents.query import map_agent_inference_error +from utils.agents.error_handler import map_agent_inference_error from utils.conversation_compaction import apply_compaction_blocking from utils.mcp_headers import McpHeaders, mcp_headers_dependency from utils.pydantic_ai_helpers import build_agent diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 0b8a6cc65..47035eb9b 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -105,7 +105,7 @@ select_model_for_responses, ) from utils.rh_identity import get_rh_identity_context -from utils.shields import run_shield_moderation +from utils.shields import run_shield_moderation_v2 from utils.suid import ( normalize_conversation_id, ) @@ -424,11 +424,11 @@ async def responses_endpoint_handler( attachments_text = extract_attachments_text(original_request.input) endpoint_path = ENDPOINT_PATH_RESPONSES - moderation_result = await run_shield_moderation( - client, + + moderation_result = await run_shield_moderation_v2( input_text + "\n\n" + attachments_text, - endpoint_path, - original_request.shield_ids, + configuration.configuration.shields, + responses_request.shield_ids, ) filter_server_tools = ( diff --git a/src/utils/agents/error_handler.py b/src/utils/agents/error_handler.py new file mode 100644 index 000000000..4d2fa8cbc --- /dev/null +++ b/src/utils/agents/error_handler.py @@ -0,0 +1,103 @@ +"""Error mapping for agent inference failures to structured API error responses.""" + +from typing import TypeAlias + +from ogx_client import APIConnectionError, APIStatusError +from pydantic_ai.exceptions import ( + AgentRunError, + ContentFilterError, + IncompleteToolCall, + ModelAPIError, + ModelHTTPError, + UnexpectedModelBehavior, + UsageLimitExceeded, +) + +from log import get_logger +from models.api.responses.error import ( + AbstractErrorResponse, + InternalServerErrorResponse, + PromptTooLongResponse, + QuotaExceededResponse, + ServiceUnavailableResponse, +) +from utils.query import ( + handle_known_apistatus_errors, + is_context_length_error, +) + +AgentInferenceError: TypeAlias = ( + AgentRunError | APIStatusError | APIConnectionError | RuntimeError +) + +logger = get_logger(__name__) + + +def map_agent_inference_error( + exc: AgentInferenceError, + model_id: str, +) -> AbstractErrorResponse: + """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response. + + Args: + exc: Agent, HTTP status, connection, or context-length runtime error. + model_id: Model identifier in provider/model format. + + Returns: + Structured error response for HTTP or SSE error events. + + Raises: + RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is + not a recognized context-length failure. + """ + match exc: + case AgentRunError() as agent_exc: + return map_pydantic_agent_run_error(agent_exc, model_id) + case APIStatusError() as status_exc: + return handle_known_apistatus_errors(status_exc, model_id) + case APIConnectionError() as connection_exc: + return ServiceUnavailableResponse( + backend_name="OGX", + cause=str(connection_exc), + ) + case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)): + return PromptTooLongResponse(model=model_id) + case _: + return InternalServerErrorResponse.generic() + + +def map_pydantic_agent_run_error( + exc: AgentRunError, model_id: str +) -> AbstractErrorResponse: + """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses. + + Args: + exc: Agent exception to map. + model_id: Model identifier in provider/model format. + + Returns: + Structured error response for HTTP or SSE error events. + """ + match exc: + case ContentFilterError() as filter_exc: + return InternalServerErrorResponse.query_failed(str(filter_exc)) + case IncompleteToolCall(): + return PromptTooLongResponse(model=model_id) + case UnexpectedModelBehavior(): + logger.error("Unexpected model behavior: %s", exc, exc_info=True) + return InternalServerErrorResponse.generic() + case UsageLimitExceeded(): + return QuotaExceededResponse.model(model_id) + case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)): + return PromptTooLongResponse(model=model_id) + case ModelHTTPError(status_code=429): + return QuotaExceededResponse.model(model_id) + case ModelHTTPError(): + return InternalServerErrorResponse.generic() + case ModelAPIError() as api_exc: + return ServiceUnavailableResponse( + backend_name="OGX", + cause=str(api_exc), + ) + case _: + return InternalServerErrorResponse.query_failed(str(exc)) diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index 16fa60100..b239dabfe 100644 --- a/src/utils/agents/query.py +++ b/src/utils/agents/query.py @@ -9,12 +9,6 @@ from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient from pydantic_ai.exceptions import ( AgentRunError, - ContentFilterError, - IncompleteToolCall, - ModelAPIError, - ModelHTTPError, - UnexpectedModelBehavior, - UsageLimitExceeded, ) from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart from pydantic_ai.run import AgentRunResult @@ -27,14 +21,13 @@ AbstractErrorResponse, InternalServerErrorResponse, PromptTooLongResponse, - QuotaExceededResponse, - ServiceUnavailableResponse, ) from models.common.agents import AgentTurnAccumulator from models.common.moderation import ShieldModerationResult from models.common.responses.responses_api_params import ResponsesApiParams from models.common.responses.types import ResponseInput from models.common.turn_summary import TurnSummary +from utils.agents.error_handler import map_agent_inference_error from utils.agents.tool_processor import ( process_function_tool_call, process_function_tool_result, @@ -45,8 +38,6 @@ from utils.pydantic_ai_helpers import build_agent from utils.query import ( extract_provider_and_model_from_model_id, - handle_known_apistatus_errors, - is_context_length_error, ) from utils.responses import extract_vector_store_ids_from_tools from utils.token_counter import TokenCounter @@ -68,76 +59,6 @@ class AgentFinishReason(str, Enum): ERROR = "error" -def map_agent_inference_error( - exc: AgentInferenceError, - model_id: str, -) -> AbstractErrorResponse: - """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response. - - Args: - exc: Agent, HTTP status, connection, or context-length runtime error. - model_id: Model identifier in provider/model format. - - Returns: - Structured error response for HTTP or SSE error events. - - Raises: - RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is - not a recognized context-length failure. - """ - match exc: - case AgentRunError() as agent_exc: - return map_pydantic_agent_run_error(agent_exc, model_id) - case APIStatusError() as status_exc: - return handle_known_apistatus_errors(status_exc, model_id) - case APIConnectionError() as connection_exc: - return ServiceUnavailableResponse( - backend_name="OGX", - cause=str(connection_exc), - ) - case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)): - return PromptTooLongResponse(model=model_id) - case _: - return InternalServerErrorResponse.generic() - - -def map_pydantic_agent_run_error( - exc: AgentRunError, model_id: str -) -> AbstractErrorResponse: - """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses. - - Args: - exc: Agent exception to map. - model_id: Model identifier in provider/model format. - - Returns: - Structured error response for HTTP or SSE error events. - """ - match exc: - case ContentFilterError() as filter_exc: - return InternalServerErrorResponse.query_failed(str(filter_exc)) - case IncompleteToolCall(): - return PromptTooLongResponse(model=model_id) - case UnexpectedModelBehavior(): - logger.error("Unexpected model behavior: %s", exc, exc_info=True) - return InternalServerErrorResponse.generic() - case UsageLimitExceeded(): - return QuotaExceededResponse.model(model_id) - case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)): - return PromptTooLongResponse(model=model_id) - case ModelHTTPError(status_code=429): - return QuotaExceededResponse.model(model_id) - case ModelHTTPError(): - return InternalServerErrorResponse.generic() - case ModelAPIError() as api_exc: - return ServiceUnavailableResponse( - backend_name="OGX", - cause=str(api_exc), - ) - case _: - return InternalServerErrorResponse.query_failed(str(exc)) - - def get_agent_finish_reason(response: ModelResponse) -> AgentFinishReason: """Get the finish reason from a completed agent model response. diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 63c2d0375..7c36e54a1 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -43,12 +43,12 @@ from models.common.responses.contexts import ResponseGeneratorContext from models.common.responses.responses_api_params import ResponsesApiParams from models.common.turn_summary import TurnSummary +from utils.agents.error_handler import map_agent_inference_error from utils.agents.query import ( AgentFinishReason, extract_agent_token_usage, get_agent_finish_reason, get_finish_reason_error, - map_agent_inference_error, ) from utils.agents.tool_processor import ( process_function_tool_call, diff --git a/src/utils/shields.py b/src/utils/shields.py index a4470b4dd..4d16dcc20 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -3,15 +3,11 @@ from typing import Optional from fastapi import HTTPException -from ogx_client import ( - APIConnectionError, - AsyncOgxClient, -) -from ogx_client import ( - APIStatusError as LLSApiStatusError, -) +from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient +from pydantic_ai.exceptions import AgentRunError from configuration import AppConfig +from log import get_logger from models.api.requests import QueryRequest from models.api.responses.error import ( InternalServerErrorResponse, @@ -19,8 +15,21 @@ ServiceUnavailableResponse, UnprocessableEntityResponse, ) -from models.common import ShieldModerationPassed, ShieldModerationResult -from models.config import ShieldConfiguration +from models.common.moderation import ( + ShieldModerationPassed, + ShieldModerationResult, +) +from models.config import QuestionValidityConfig, RedactionConfig, ShieldConfiguration +from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability +from pydantic_ai_lightspeed.capabilities.question_validity._capability import ( + QuestionValidity, +) +from pydantic_ai_lightspeed.capabilities.redaction._capability import ( + PiiRedactionCapability, +) +from utils.agents.error_handler import map_agent_inference_error + +logger = get_logger(__name__) def validate_shield_ids_override( @@ -59,6 +68,63 @@ def validate_shield_ids_override( raise HTTPException(**response.model_dump()) +async def run_shield_moderation_v2( + input_text: str, + shield_configs: list[ShieldConfiguration], + selected_shield_ids: Optional[list[str]] = None, +) -> ShieldModerationResult: + """Run v2 shield moderation on input text. + + Iterates through configured shields and runs moderation checks. + + Parameters: + input_text: The text to moderate. + shield_configs: List of shield configurations to evaluate. + selected_shield_ids: Optional list of shield names to filter by. + + Returns: + Result indicating if content was blocked or passed. + """ + selected_shield_configs = get_shields_for_request( + shield_configs, selected_shield_ids + ) + + for shield_config in selected_shield_configs: + shield = build_shield(shield_config) + + try: + shield_result = await shield.run(input_text) + # APIConnectionError and APIStatusError from ogx should not be raised from model_request, + # because they will be caught inside AsyncOpenAI and transferred into openai's + # APIConnectionError. The openai's exceptions will further transferred into ModelHTTPError + # or ModelAPIError by _map_api_errors in OpenAIResponseModel. + except (AgentRunError, RuntimeError) as exc: + model_id = getattr(shield_config.config, "model_id", "unknown-shield-model") + response = map_agent_inference_error(exc, model_id) + raise HTTPException(**response.model_dump()) from exc + + if shield_result.decision == "blocked": + return shield_result + + return ShieldModerationPassed() + + +def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability: + """Build a safety capability instance from a shield configuration. + + Parameters: + shield_config: The shield configuration to build from. + + Returns: + The constructed safety capability. + """ + match shield_config.config: + case QuestionValidityConfig(): + return QuestionValidity(shield_config.config) + case RedactionConfig(): + return PiiRedactionCapability(shield_config.config) + + async def run_shield_moderation( _client: AsyncOgxClient, _input_text: str, @@ -124,7 +190,7 @@ async def append_turn_to_conversation( cause=str(e), ) raise HTTPException(**error_response.model_dump()) from e - except LLSApiStatusError as e: + except APIStatusError as e: error_response = InternalServerErrorResponse.generic() raise HTTPException(**error_response.model_dump()) from e diff --git a/tests/integration/endpoints/test_responses_integration.py b/tests/integration/endpoints/test_responses_integration.py index aae7b9f26..9ef34711f 100644 --- a/tests/integration/endpoints/test_responses_integration.py +++ b/tests/integration/endpoints/test_responses_integration.py @@ -173,7 +173,7 @@ def _configure_shield_blocked( moderation_id=moderation_id, ) mocker.patch( - "app.endpoints.responses.run_shield_moderation", + "app.endpoints.responses.run_shield_moderation_v2", return_value=blocked, ) diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index 57020665b..63b4c203f 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -193,7 +193,7 @@ def _patch_moderation(mocker: MockerFixture, decision: str = "passed") -> Any: else: moderation_result = ShieldModerationPassed() mocker.patch( - f"{MODULE}.run_shield_moderation", + f"{MODULE}.run_shield_moderation_v2", new=mocker.AsyncMock(return_value=moderation_result), ) return moderation_result diff --git a/tests/unit/utils/agents/test_query.py b/tests/unit/utils/agents/test_query.py index 94f72cf48..f98f4c65a 100644 --- a/tests/unit/utils/agents/test_query.py +++ b/tests/unit/utils/agents/test_query.py @@ -477,7 +477,7 @@ async def test_api_status_error_raises_http_exception( "detail": {"response": "Quota exceeded", "cause": "quota exceeded"}, } mocker.patch( - "utils.agents.query.handle_known_apistatus_errors", + "utils.agents.error_handler.handle_known_apistatus_errors", return_value=mock_error, ) diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index 3190feb26..132e4ad13 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -605,7 +605,7 @@ async def test_agent_error_raises_http_exception( "detail": {"response": "Error", "cause": "agent failed"}, } mocker.patch( - "utils.agents.streaming.map_agent_inference_error", + "utils.agents.error_handler.map_agent_inference_error", return_value=mock_error, ) @@ -742,7 +742,7 @@ async def inner() -> AsyncIterator[str]: mock_error.detail.response = "Quota exceeded" mock_error.detail.cause = "quota exceeded" mocker.patch( - "utils.agents.streaming.map_agent_inference_error", + "utils.agents.error_handler.map_agent_inference_error", return_value=mock_error, ) mocker.patch( diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py index 7183e3e89..a1c22a8aa 100644 --- a/tests/unit/utils/test_shields.py +++ b/tests/unit/utils/test_shields.py @@ -2,17 +2,24 @@ import pytest from fastapi import HTTPException, status +from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError from pytest_mock import MockerFixture -from models.config import QuestionValidityConfig, QuestionValidityShieldConfiguration +from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.config import ( + QuestionValidityConfig, + QuestionValidityShieldConfiguration, + ShieldConfiguration, +) from utils.shields import ( append_turn_to_conversation, get_shields_for_request, + run_shield_moderation_v2, validate_shield_ids_override, ) -def _shield(name: str) -> QuestionValidityShieldConfiguration: +def _shield_config(name: str) -> QuestionValidityShieldConfiguration: """Build a minimal question-validity shield configuration for tests.""" return QuestionValidityShieldConfiguration( name=name, @@ -133,12 +140,156 @@ def test_raises_422_when_empty_list_shield_ids_and_override_disabled( assert exc_info.value.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY +class TestRunShieldModerationV2: + """Tests for run_shield_moderation_v2 function.""" + + @pytest.mark.asyncio + async def test_returns_passed_when_no_shields(self) -> None: + """Return ShieldModerationPassed when shield list is empty.""" + result = await run_shield_moderation_v2("test input", []) + assert isinstance(result, ShieldModerationPassed) + + @pytest.mark.asyncio + async def test_returns_passed_when_all_shields_pass( + self, mocker: MockerFixture + ) -> None: + """Return ShieldModerationPassed when every shield passes.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationPassed) + assert mock_shield.run.call_count == 2 + + @pytest.mark.asyncio + async def test_returns_blocked_on_first_block(self, mocker: MockerFixture) -> None: + """Return blocked result from first shield that blocks.""" + blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-123") + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=blocked) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationBlocked) + assert result.message == "rejected" + mock_shield.run.assert_called_once() + + @pytest.mark.asyncio + async def test_filters_by_selected_shield_ids(self, mocker: MockerFixture) -> None: + """Only run shields matching the selected IDs.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + _shield_config("s3"), + ] + result = await run_shield_moderation_v2( + "test input", shields, selected_shield_ids=["s2"] + ) + + assert isinstance(result, ShieldModerationPassed) + mock_shield.run.assert_called_once() + + @pytest.mark.asyncio + async def test_shields_stops_on_first_block(self, mocker: MockerFixture) -> None: + """Stop at the first blocking shield.""" + blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-789") + mock_qv_shield = mocker.Mock() + mock_qv_shield.run = mocker.AsyncMock(return_value=blocked) + + mock_redact_shield = mocker.Mock() + mock_redact_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + + mocker.patch( + "utils.shields.build_shield", + side_effect=[mock_qv_shield, mock_redact_shield], + ) + + shields: list[ShieldConfiguration] = [ + _shield_config("s-1"), + _shield_config("s-2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationBlocked) + mock_qv_shield.run.assert_called_once() + mock_redact_shield.run.assert_not_called() + + @pytest.mark.asyncio + async def test_raise_503_on_model_api_error(self, mocker: MockerFixture) -> None: + """Raise HTTP 503 when a shield raises ModelAPIError.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelAPIError("test", "Incompatible mode") + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert "OGX" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_raise_429_when_exceeds_quota(self, mocker: MockerFixture) -> None: + """Raise HTTP 429 when a shield raises ModelHTTPError with status 429.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelHTTPError(429, "openai/gpt-4o-mini", "Quota exceeded") + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert "test-model" in str(exc_info.value.detail) + assert "The model quota has been exceeded" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_raise_413_when_exceeds_context_length( + self, mocker: MockerFixture + ) -> None: + """Raise HTTP 413 when a shield raises ModelHTTPError due to context length exceeded.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelHTTPError( + 413, "openai/gpt-4o-mini", "Context length exceeded" + ) + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_413_CONTENT_TOO_LARGE + assert "test-model" in str(exc_info.value.detail) + assert "Prompt is too long" in str(exc_info.value.detail) + + class TestGetShieldsForRequest: """Tests for get_shields_for_request function.""" def test_returns_all_shields_when_shield_ids_none(self) -> None: """Return all configured shields when shield_ids is None.""" - shields = [_shield("shield-1"), _shield("shield-2")] + shields = [ + _shield_config("shield-1"), + _shield_config("shield-2"), + ] result = get_shields_for_request(shields, shield_ids=None) @@ -146,7 +297,10 @@ def test_returns_all_shields_when_shield_ids_none(self) -> None: def test_returns_empty_list_when_shield_ids_empty(self) -> None: """Return no shields when an empty shield_ids list is provided.""" - shields = [_shield("shield-1"), _shield("shield-2")] + shields = [ + _shield_config("shield-1"), + _shield_config("shield-2"), + ] result = get_shields_for_request(shields, shield_ids=[]) @@ -154,9 +308,9 @@ def test_returns_empty_list_when_shield_ids_empty(self) -> None: def test_filters_to_requested_shields_when_all_exist(self) -> None: """Return only shields whose names appear in shield_ids.""" - shield1 = _shield("shield-1") - shield2 = _shield("shield-2") - shield3 = _shield("shield-3") + shield1 = _shield_config("shield-1") + shield2 = _shield_config("shield-2") + shield3 = _shield_config("shield-3") result = get_shields_for_request( [shield1, shield2, shield3], shield_ids=["shield-1", "shield-3"] @@ -168,7 +322,8 @@ def test_raises_404_when_requested_shield_not_configured(self) -> None: """Raise 404 when a requested shield name is not configured.""" with pytest.raises(HTTPException) as exc_info: get_shields_for_request( - [_shield("shield-1")], shield_ids=["shield-1", "missing-shield"] + [_shield_config("shield-1")], + shield_ids=["shield-1", "missing-shield"], ) assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND