diff --git a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py index 55a9ae526..950ed5c74 100644 --- a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py +++ b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py @@ -19,7 +19,13 @@ from pydantic_ai._agent_graph import GraphAgentState from pydantic_ai.capabilities import WrapRunHandler from pydantic_ai.direct import model_request -from pydantic_ai.messages import ModelRequest, TextContent, UserContent +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextContent, + TextPart, + UserContent, +) from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIResponsesModelSettings @@ -35,6 +41,7 @@ ) from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability from pydantic_ai_lightspeed.llamastack import OgxResponsesModel +from utils.shields import append_turn_to_conversation logger = get_logger(__name__) @@ -62,6 +69,47 @@ def _extract_message_str_from_user_content(user_content: Sequence[UserContent]) return "\n".join(str_arr) +def _message_to_str(message: Optional[str | Sequence[UserContent]]) -> str: + """Convert a user message (string, content sequence, or None) to plain text. + + Parameters: + message: The user input as a string, sequence of user content, or None. + + Returns: + A plain-text representation of the message, or an empty string for None. + """ + match message: + case str() as s: + return s + case Sequence() as seq: + return _extract_message_str_from_user_content(seq) + case None: + return "" + + +def _extract_conversation_id(model: Model) -> Optional[str]: + """Extract the Llama Stack conversation ID from the agent's model settings. + + The main agent's model is built with ``conversation`` in its + ``extra_body`` model settings (see ``OgxResponsesModel.from_ogx_client``). + This pulls it back out so the capability can persist the rejected turn + to the same conversation. + + Parameters: + model: The model bound to the current agent run (``ctx.model``). + + Returns: + The conversation ID, or None if the model has no such setting + (e.g. when used outside a Llama Stack-backed agent). + """ + extra_body = (model.settings or {}).get("extra_body") + if not isinstance(extra_body, dict): + return None + + conversation_id = extra_body.get("conversation") + return conversation_id if isinstance(conversation_id, str) else None + + @dataclass class QuestionValidity(AbstractSafetyCapability): """Block or modify user input based on a guardrail check. @@ -100,16 +148,10 @@ def _build_prompt(self, message: Optional[str | Sequence[UserContent]]) -> str: Returns: The rendered prompt string ready to send to the validity model. """ - match message: - case str() as s: - _message = s - case Sequence() as seq: - _message = _extract_message_str_from_user_content(seq) - case None: - _message = "" - return Template(self.config.model_prompt).substitute( - message=_message, allowed=SUBJECT_ALLOWED, rejected=SUBJECT_REJECTED + message=_message_to_str(message), + allowed=SUBJECT_ALLOWED, + rejected=SUBJECT_REJECTED, ) async def wrap_run( @@ -143,7 +185,32 @@ async def wrap_run( return await handler() # proceed with the real run # short-circuit: return the rejection message with shield usage tracked - state = GraphAgentState(usage=ctx.usage) + user_message = _message_to_str(ctx.prompt) + state = GraphAgentState( + usage=ctx.usage, + message_history=[ + ModelRequest.user_text_prompt(user_message), + ModelResponse( + [TextPart(self.config.invalid_question_response)], + finish_reason="stop", + ), + ], + ) + + conversation_id = _extract_conversation_id(ctx.model) + if conversation_id is not None: + await append_turn_to_conversation( + AsyncOgxClientHolder().get_client(), + conversation_id, + user_message, + self.config.invalid_question_response, + ) + else: + logger.warning( + "Unable to determine conversation ID from model settings; " + "skipping v1/conversation persistence for rejected question." + ) + return AgentRunResult( output=self.config.invalid_question_response, _state=state ) diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py index cae0d2916..9658e917a 100644 --- a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py @@ -22,6 +22,7 @@ SUBJECT_ALLOWED, SUBJECT_REJECTED, QuestionValidity, + _extract_conversation_id, _extract_message_str_from_user_content, ) @@ -65,6 +66,47 @@ def test_sequence_with_non_text_content(self) -> None: assert result == "keep" +class TestExtractConversationId: + """Tests for _extract_conversation_id helper.""" + + def test_extracts_conversation_id(self, mocker: MockerFixture) -> None: + """Test extraction when extra_body.conversation is set.""" + model = mocker.Mock() + model.settings = {"extra_body": {"conversation": "conv_123"}} + + assert _extract_conversation_id(model) == "conv_123" + + def test_returns_none_when_settings_missing(self, mocker: MockerFixture) -> None: + """Test that None settings yields None.""" + model = mocker.Mock() + model.settings = None + + assert _extract_conversation_id(model) is None + + def test_returns_none_when_extra_body_missing(self, mocker: MockerFixture) -> None: + """Test that missing extra_body yields None.""" + model = mocker.Mock() + model.settings = {} + + assert _extract_conversation_id(model) is None + + def test_returns_none_when_extra_body_not_dict(self, mocker: MockerFixture) -> None: + """Test that a non-dict extra_body yields None instead of raising.""" + model = mocker.Mock() + model.settings = {"extra_body": "not-a-dict"} + + assert _extract_conversation_id(model) is None + + def test_returns_none_when_conversation_not_string( + self, mocker: MockerFixture + ) -> None: + """Test that a non-string conversation value yields None.""" + model = mocker.Mock() + model.settings = {"extra_body": {"conversation": 123}} + + assert _extract_conversation_id(model) is None + + class TestQuestionValidityConfigInit: """Tests for QuestionValidityConfig initialization.""" @@ -216,12 +258,21 @@ def _mock_create_model(self, mocker: MockerFixture) -> None: mocker.patch(f"{_MODULE}.AsyncOgxClientHolder") mocker.patch(f"{_MODULE}.OgxResponsesModel.from_ogx_client") + @pytest.fixture(name="mock_append_turn", autouse=True) + def mock_append_turn_fixture(self, mocker: MockerFixture) -> MockType: + """Mock the conversation-persistence call used on rejection.""" + return mocker.patch( + f"{_MODULE}.append_turn_to_conversation", new_callable=mocker.AsyncMock + ) + @pytest.fixture(name="mock_ctx") def mock_ctx_fixture(self, mocker: MockerFixture) -> RunContext: - """Create a mock RunContext.""" + """Create a mock RunContext bound to a model with a conversation ID.""" ctx = mocker.Mock(spec=RunContext) ctx.prompt = "How do I create a pod?" ctx.usage = RunUsage() + ctx.model = mocker.Mock() + ctx.model.settings = {"extra_body": {"conversation": "conv_test"}} return ctx @pytest.fixture(name="mock_handler") @@ -280,6 +331,89 @@ async def test_rejected_question_returns_rejection( assert isinstance(result, AgentRunResult) assert result.output == DEFAULT_INVALID_QUESTION_RESPONSE + @pytest.mark.asyncio + async def test_rejected_question_persists_turn_to_conversation( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that a rejection appends the user question and refusal to the conversation.""" + mock_client = mocker.Mock() + mocker.patch( + f"{_MODULE}.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + mock_response = ModelResponse( + parts=[TextPart(content=SUBJECT_REJECTED)], + usage=RequestUsage(input_tokens=10, output_tokens=1), + ) + mocker.patch( + "pydantic_ai_lightspeed.capabilities.question_validity._capability.model_request", + return_value=mock_response, + ) + + config = QuestionValidityConfig(model_id="test") + qv = QuestionValidity(config=config) + await qv.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_awaited_once_with( + mock_client, + "conv_test", + "How do I create a pod?", + DEFAULT_INVALID_QUESTION_RESPONSE, + ) + + @pytest.mark.asyncio + async def test_rejection_skips_persistence_when_conversation_id_missing( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that persistence is skipped (not crashed) without a conversation ID.""" + mock_ctx.model = mocker.Mock(settings={}) + mock_response = ModelResponse( + parts=[TextPart(content=SUBJECT_REJECTED)], + usage=RequestUsage(), + ) + mocker.patch( + "pydantic_ai_lightspeed.capabilities.question_validity._capability.model_request", + return_value=mock_response, + ) + + config = QuestionValidityConfig(model_id="test") + qv = QuestionValidity(config=config) + result = await qv.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_not_awaited() + assert result.output == DEFAULT_INVALID_QUESTION_RESPONSE + + @pytest.mark.asyncio + async def test_allowed_question_does_not_persist_turn( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that an allowed question does not touch the conversation.""" + mock_response = ModelResponse( + parts=[TextPart(content=SUBJECT_ALLOWED)], + usage=RequestUsage(input_tokens=10, output_tokens=1), + ) + mocker.patch( + "pydantic_ai_lightspeed.capabilities.question_validity._capability.model_request", + return_value=mock_response, + ) + + config = QuestionValidityConfig(model_id="test") + qv = QuestionValidity(config=config) + await qv.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_not_awaited() + @pytest.mark.asyncio async def test_unexpected_response_treated_as_rejected( self, @@ -469,6 +603,8 @@ async def test_wrap_run_with_none_prompt( ctx = mocker.Mock(spec=RunContext) ctx.prompt = None ctx.usage = RunUsage() + ctx.model = mocker.Mock() + ctx.model.settings = {"extra_body": {"conversation": "conv_test"}} mock_response = ModelResponse( parts=[TextPart(content=SUBJECT_REJECTED)],