From 3068b80976416f615b2db053df7660d0b4be71e1 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Wed, 29 Jul 2026 11:33:32 +0200 Subject: [PATCH] Resolve circular import between shields and question_validity. Move append_turn_to_conversation into utils.conversations so capabilities no longer depend on utils.shields at import time. Co-authored-by: Cursor --- .../question_validity/_capability.py | 2 +- src/utils/conversations.py | 38 ++++++++++++++++ src/utils/shields.py | 44 +------------------ src/utils/stream_interrupts.py | 6 ++- tests/unit/utils/test_conversations.py | 32 ++++++++++++++ tests/unit/utils/test_shields.py | 32 -------------- 6 files changed, 77 insertions(+), 77 deletions(-) diff --git a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py index 950ed5c74..c8097d1ad 100644 --- a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py +++ b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py @@ -41,7 +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 +from utils.conversations import append_turn_to_conversation logger = get_logger(__name__) diff --git a/src/utils/conversations.py b/src/utils/conversations.py index d2d039beb..9c98e1ed6 100644 --- a/src/utils/conversations.py +++ b/src/utils/conversations.py @@ -553,3 +553,41 @@ async def get_all_conversation_items( except APIStatusError as e: error_response = InternalServerErrorResponse.generic() raise HTTPException(**error_response.model_dump()) from e + + +async def append_turn_to_conversation( + client: AsyncOgxClient, + conversation_id: str, + user_message: str, + assistant_message: str, +) -> None: + """ + Append a user/assistant turn to a conversation. + + Used to record a conversation turn when a shield blocks the request, + storing both the user's original message and the violation response. + + Parameters: + ---------- + client: The Llama Stack client. + conversation_id: The Llama Stack conversation ID. + user_message: The user's input message. + assistant_message: The shield violation response message. + """ + try: + await client.conversations.items.create( + conversation_id, + items=[ + {"type": "message", "role": "user", "content": user_message}, + {"type": "message", "role": "assistant", "content": assistant_message}, + ], + ) + except APIConnectionError as e: + error_response = ServiceUnavailableResponse( + backend_name="OGX", + cause=str(e), + ) + raise HTTPException(**error_response.model_dump()) from e + except APIStatusError as e: + error_response = InternalServerErrorResponse.generic() + raise HTTPException(**error_response.model_dump()) from e diff --git a/src/utils/shields.py b/src/utils/shields.py index 73cece586..dd72da72c 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -1,18 +1,16 @@ -"""Utility helpers for shield override validation and conversation persistence.""" +"""Utility helpers for shield override validation and moderation.""" from typing import Optional from fastapi import HTTPException -from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient +from ogx_client import 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, NotFoundResponse, - ServiceUnavailableResponse, UnprocessableEntityResponse, ) from models.common.moderation import ( @@ -157,44 +155,6 @@ async def run_shield_moderation( return ShieldModerationPassed() -async def append_turn_to_conversation( - client: AsyncOgxClient, - conversation_id: str, - user_message: str, - assistant_message: str, -) -> None: - """ - Append a user/assistant turn to a conversation after shield violation. - - Used to record the conversation turn when a shield blocks the request, - storing both the user's original message and the violation response. - - Parameters: - ---------- - client: The Llama Stack client. - conversation_id: The Llama Stack conversation ID. - user_message: The user's input message. - assistant_message: The shield violation response message. - """ - try: - await client.conversations.items.create( - conversation_id, - items=[ - {"type": "message", "role": "user", "content": user_message}, - {"type": "message", "role": "assistant", "content": assistant_message}, - ], - ) - except APIConnectionError as e: - error_response = ServiceUnavailableResponse( - backend_name="OGX", - cause=str(e), - ) - raise HTTPException(**error_response.model_dump()) from e - except APIStatusError as e: - error_response = InternalServerErrorResponse.generic() - raise HTTPException(**error_response.model_dump()) from e - - def get_shields_for_request( shields: list[ShieldConfiguration], shield_ids: Optional[list[str]] = None, diff --git a/src/utils/stream_interrupts.py b/src/utils/stream_interrupts.py index 4ab9b1959..acdd771b3 100644 --- a/src/utils/stream_interrupts.py +++ b/src/utils/stream_interrupts.py @@ -19,11 +19,13 @@ 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.conversations import append_turn_items_to_conversation +from utils.conversations import ( + append_turn_items_to_conversation, + append_turn_to_conversation, +) from utils.markdown_repair import close_open_markdown from utils.query import store_query_results, update_conversation_topic_summary from utils.responses import get_topic_summary -from utils.shields import append_turn_to_conversation from utils.types import Singleton logger = get_logger(__name__) diff --git a/tests/unit/utils/test_conversations.py b/tests/unit/utils/test_conversations.py index ea14e4855..50c670aee 100644 --- a/tests/unit/utils/test_conversations.py +++ b/tests/unit/utils/test_conversations.py @@ -26,6 +26,7 @@ _extract_text_from_content, _function_call_output_to_str, append_turn_items_to_conversation, + append_turn_to_conversation, build_conversation_turns_from_items, get_all_conversation_items, ) @@ -825,6 +826,37 @@ async def test_appends_user_input_and_llm_output( assert items[1]["content"] == "I cannot help with that" +class TestAppendTurnToConversation: # pylint: disable=too-few-public-methods + """Tests for append_turn_to_conversation function.""" + + @pytest.mark.asyncio + async def test_appends_user_and_assistant_messages( + self, mocker: MockerFixture + ) -> None: + """Test that append_turn_to_conversation creates conversation items correctly.""" + mock_client = mocker.Mock() + mock_client.conversations.items.create = mocker.AsyncMock(return_value=None) + + await append_turn_to_conversation( + mock_client, + conversation_id="conv-123", + user_message="Hello", + assistant_message="I cannot help with that", + ) + + mock_client.conversations.items.create.assert_called_once_with( + "conv-123", + items=[ + {"type": "message", "role": "user", "content": "Hello"}, + { + "type": "message", + "role": "assistant", + "content": "I cannot help with that", + }, + ], + ) + + class TestGetAllConversationItems: """Tests for get_all_conversation_items function.""" diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py index 4a7d850be..266a4d287 100644 --- a/tests/unit/utils/test_shields.py +++ b/tests/unit/utils/test_shields.py @@ -12,7 +12,6 @@ ShieldConfiguration, ) from utils.shields import ( - append_turn_to_conversation, get_shields_for_request, run_shield_moderation_v2, validate_shield_ids_override, @@ -28,37 +27,6 @@ def _shield_config(name: str) -> QuestionValidityShieldConfiguration: ) -class TestAppendTurnToConversation: # pylint: disable=too-few-public-methods - """Tests for append_turn_to_conversation function.""" - - @pytest.mark.asyncio - async def test_appends_user_and_assistant_messages( - self, mocker: MockerFixture - ) -> None: - """Test that append_turn_to_conversation creates conversation items correctly.""" - mock_client = mocker.Mock() - mock_client.conversations.items.create = mocker.AsyncMock(return_value=None) - - await append_turn_to_conversation( - mock_client, - conversation_id="conv-123", - user_message="Hello", - assistant_message="I cannot help with that", - ) - - mock_client.conversations.items.create.assert_called_once_with( - "conv-123", - items=[ - {"type": "message", "role": "user", "content": "Hello"}, - { - "type": "message", - "role": "assistant", - "content": "I cannot help with that", - }, - ], - ) - - class TestValidateShieldIdsOverride: """Tests for validate_shield_ids_override function."""