fix(shields): fixed QuestionValidity shield

jrobertboos · jrobertboos · commit 02afe552c55e · 2026-07-28T18:49:00.000-04:00
diff --git 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
@@ -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)],
