From 3e0520704193bac932f26d92aefbf1e2a1c97ce8 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Tue, 11 Aug 2026 09:03:59 +0200 Subject: [PATCH] Added custom serializer Co-authored-by: Cursor --- src/app/endpoints/prompts.py | 9 +- src/app/endpoints/providers.py | 3 +- src/app/endpoints/responses.py | 11 +- src/utils/ogx_serialization.py | 19 ++ tests/integration/conftest.py | 139 ++++++----- .../test_responses_byok_integration.py | 6 +- .../endpoints/test_responses_integration.py | 20 +- tests/unit/app/endpoints/test_responses.py | 75 ++++-- .../app/endpoints/test_responses_splunk.py | 6 +- tests/unit/utils/test_ogx_serialization.py | 221 ++++++++++++++++++ 10 files changed, 403 insertions(+), 106 deletions(-) create mode 100644 src/utils/ogx_serialization.py create mode 100644 tests/unit/utils/test_ogx_serialization.py diff --git a/src/app/endpoints/prompts.py b/src/app/endpoints/prompts.py index e672ab0fb..17525da6a 100644 --- a/src/app/endpoints/prompts.py +++ b/src/app/endpoints/prompts.py @@ -28,6 +28,7 @@ ) from models.config import Action from utils.endpoints import check_configuration_loaded +from utils.ogx_serialization import dump_ogx_model from utils.query import handle_known_apistatus_errors from utils.suid import check_suid_prompt @@ -136,7 +137,7 @@ async def create_prompt_handler( client = AsyncOgxClientHolder().get_client() payload = body.model_dump(exclude_none=True) created = await client.prompts.create(**payload) - return PromptResourceResponse.model_validate(created.model_dump()) + return PromptResourceResponse.model_validate(dump_ogx_model(created)) except ApiException as e: if not e.status: logger.error("Unable to connect to Llama Stack: %s", e) @@ -185,7 +186,7 @@ async def list_prompts_handler( try: client = AsyncOgxClientHolder().get_client() items = await client.prompts.list() - data = [PromptResourceResponse.model_validate(p.model_dump()) for p in items] + data = [PromptResourceResponse.model_validate(dump_ogx_model(p)) for p in items] return PromptsListResponse(data=data) except ApiException as e: if not e.status: @@ -246,7 +247,7 @@ async def get_prompt_handler( try: client = AsyncOgxClientHolder().get_client() retrieved = await client.prompts.retrieve(prompt_id, version=version) - return PromptResourceResponse.model_validate(retrieved.model_dump()) + return PromptResourceResponse.model_validate(dump_ogx_model(retrieved)) except (BadRequestError, ValueError) as e: logger.error("Prompt not found: %s", e) response = NotFoundResponse(resource="prompt", resource_id=prompt_id) @@ -316,7 +317,7 @@ async def update_prompt_handler( client = AsyncOgxClientHolder().get_client() payload = body.model_dump(exclude_none=True, exclude_unset=True) updated = await client.prompts.update(prompt_id, **payload) - return PromptResourceResponse.model_validate(updated.model_dump()) + return PromptResourceResponse.model_validate(dump_ogx_model(updated)) except (BadRequestError, ValueError) as e: logger.error("Prompt update failed: %s", e) response = NotFoundResponse(resource="prompt", resource_id=prompt_id) diff --git a/src/app/endpoints/providers.py b/src/app/endpoints/providers.py index 28c06a9ac..a1be36b4e 100644 --- a/src/app/endpoints/providers.py +++ b/src/app/endpoints/providers.py @@ -27,6 +27,7 @@ ) from models.config import Action from utils.endpoints import check_configuration_loaded +from utils.ogx_serialization import dump_ogx_model logger = get_logger(__name__) router = APIRouter(tags=["providers"]) @@ -164,7 +165,7 @@ async def get_provider_endpoint_handler( try: client = AsyncOgxClientHolder().get_client() provider = await client.providers.retrieve(provider_id) - return ProviderResponse(**provider.model_dump()) + return ProviderResponse(**dump_ogx_model(provider)) except (BadRequestError, ValueError) as e: # Server mode raises BadRequestError; library mode raises ValueError. diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index e5e2b5872..630e065ef 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -69,6 +69,7 @@ ) from utils.mcp_headers import mcp_headers_dependency from utils.mcp_oauth_probe import check_mcp_auth +from utils.ogx_serialization import dump_ogx_model from utils.prompts import get_system_prompt from utils.query import ( consume_query_tokens, @@ -711,7 +712,7 @@ def _sanitize_response_dict( response object before it is forwarded to the client. Args: - response_dict: Mutable dict produced by ``model_dump`` on a response + response_dict: Mutable dict produced by ``dump_ogx_model`` on a response object. Modified in-place. configured_mcp_labels: Set of ``server_label`` values that identify server-deployed MCP servers. @@ -886,7 +887,7 @@ async def response_generator( ): continue - chunk_dict = chunk.model_dump(exclude_none=True, by_alias=True) + chunk_dict = dump_ogx_model(chunk) # Create own sequence number for chunks to maintain order chunk_dict["sequence_number"] = sequence_number @@ -1164,7 +1165,11 @@ async def handle_non_streaming_response( output_text, ) configured_mcp_labels = {s.name for s in configuration.mcp_servers} - response_dict = api_response.model_dump(exclude_none=True) + response_dict = ( + api_response.model_dump(exclude_none=True) + if context.moderation_result.decision == "blocked" + else dump_ogx_model(api_response) + ) _sanitize_response_dict( response_dict, configured_mcp_labels, diff --git a/src/utils/ogx_serialization.py b/src/utils/ogx_serialization.py new file mode 100644 index 000000000..4912b150f --- /dev/null +++ b/src/utils/ogx_serialization.py @@ -0,0 +1,19 @@ +"""Serialization helpers for ``ogx_client`` models.""" + +from typing import Any + +from ogx_client.api_client import ApiClient + +_OGX_SERIALIZER = ApiClient() + + +def dump_ogx_model(obj: Any) -> Any: + """Dump an ogx_client model to a JSON-safe structure. + + Args: + obj: An ogx_client model instance. + + Returns: + A JSON-serializable dict, list, or primitive produced from obj. + """ + return _OGX_SERIALIZER.sanitize_for_serialization(obj) # type: ignore[no-untyped-call] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 9739a6fdf..5427dca4f 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -8,12 +8,12 @@ import pytest from fastapi import Request, Response from fastapi.testclient import TestClient -from ogx_api.openai_responses import OpenAIResponseObject from ogx_client.models.list_models_v1_models_get200_response import ( ListModelsV1ModelsGet200Response, ) from ogx_client.models.open_ai_list_models_response import OpenAIListModelsResponse from ogx_client.models.open_ai_model import OpenAIModel +from ogx_client.models.open_ai_response_object import OpenAIResponseObject from ogx_client.models.version_info import VersionInfo from pydantic_ai import AgentRunResultEvent from pydantic_ai.messages import ( @@ -112,52 +112,99 @@ def make_openai_model( ) +def make_openai_response_object( # pylint: disable=too-many-arguments,too-many-positional-arguments + *, + response_id: str = "response-123", + content: str = "This is a test response about Ansible.", + model: str = TEST_MODEL, + tool_calls: Optional[list[dict[str, Any]]] = None, + refusal: Optional[str] = None, + input_tokens: int = 10, + output_tokens: int = 5, +) -> OpenAIResponseObject: + """Build a real ``ogx_client`` OpenAI Responses API object for tests. + + Parameters: + response_id: Response identifier returned by the mocked API. + content: Assistant message text for the default output item. + model: Model identifier on the response object. + tool_calls: Optional function-call output items to append. + refusal: Optional refusal text; when set, emits a refusal content part. + input_tokens: Input token count for usage metadata. + output_tokens: Output token count for usage metadata. + + Returns: + ``OpenAIResponseObject`` instance suitable for ``dump_ogx_model()``. + """ + output: list[dict[str, Any]] = list(tool_calls or []) + message_content: list[dict[str, Any]] = ( + [{"type": "refusal", "refusal": refusal}] + if refusal + else [{"type": "output_text", "text": content, "annotations": []}] + ) + output.append( + { + "type": "message", + "id": "msg-1", + "role": "assistant", + "status": "completed", + "content": message_content, + } + ) + + payload: dict[str, Any] = { + "id": response_id, + "object": "response", + "created_at": 1_700_000_000, + "status": "completed", + "model": model, + "store": False, + "output": output, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + response = OpenAIResponseObject.from_dict(payload) + assert response is not None + return response + + def create_mock_llm_response( # pylint: disable=too-many-arguments,too-many-positional-arguments mocker: MockerFixture, content: str = "This is a test response about Ansible.", - tool_calls: Optional[list[Any]] = None, + tool_calls: Optional[list[dict[str, Any]]] = None, refusal: Optional[str] = None, input_tokens: int = 10, output_tokens: int = 5, -) -> Any: - """Create a customizable mock LLM response. +) -> OpenAIResponseObject: + """Create a customizable LLM response for integration test mocks. - Helper function to create mock LLM responses with configurable content, - tool calls, refusals, and token counts. Useful for tests that need to - customize the response behavior. + Helper function to create ``ogx_client`` response objects with configurable + content, tool calls, refusals, and token counts. Args: - mocker: pytest-mock fixture - content: Response content text - tool_calls: Optional list of tool calls - refusal: Optional refusal message (for shield violations) - input_tokens: Input token count for usage - output_tokens: Output token count for usage + mocker: pytest-mock fixture (kept for backward compatibility). + content: Response content text. + tool_calls: Optional function-call output items. + refusal: Optional refusal message (for shield violations). + input_tokens: Input token count for usage. + output_tokens: Output token count for usage. Returns: - Mock LLM response object with the specified configuration. + ``OpenAIResponseObject`` with the specified configuration. """ - mock_response = mocker.MagicMock(spec=OpenAIResponseObject) - mock_response.id = "response-123" - - # Create output message - mock_output_item = mocker.MagicMock() - mock_output_item.type = "message" - mock_output_item.role = "assistant" - mock_output_item.content = content - mock_output_item.refusal = refusal - - mock_response.output = [mock_output_item] - mock_response.stop_reason = "end_turn" if not refusal else "stop" - mock_response.tool_calls = tool_calls or [] - - # Mock usage - mock_usage = mocker.MagicMock() - mock_usage.input_tokens = input_tokens - mock_usage.output_tokens = output_tokens - mock_response.usage = mock_usage - - return mock_response + _ = mocker + return make_openai_response_object( + content=content, + tool_calls=tool_calls, + refusal=refusal, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) def create_mock_vector_store_response( @@ -795,25 +842,9 @@ def mock_ogx_client_fixture( mock_client = mocker.AsyncMock() # Mock responses.create with default assistant response - mock_response = mocker.MagicMock(spec=OpenAIResponseObject) - mock_response.id = "response-123" - - mock_output_item = mocker.MagicMock() - mock_output_item.type = "message" - mock_output_item.role = "assistant" - mock_output_item.content = "This is a test response about Ansible." - mock_output_item.refusal = None - - mock_response.output = [mock_output_item] - mock_response.stop_reason = "end_turn" - mock_response.tool_calls = [] - - mock_usage = mocker.MagicMock() - mock_usage.input_tokens = 10 - mock_usage.output_tokens = 5 - mock_response.usage = mock_usage - - mock_client.responses.create.return_value = mock_response + mock_client.responses.create = mocker.AsyncMock( + return_value=make_openai_response_object() + ) # Mock openai.list mock_client.openai.list.return_value = make_openai_models_list_response( diff --git a/tests/integration/endpoints/test_responses_byok_integration.py b/tests/integration/endpoints/test_responses_byok_integration.py index 72a7c29fe..2f9b9429b 100644 --- a/tests/integration/endpoints/test_responses_byok_integration.py +++ b/tests/integration/endpoints/test_responses_byok_integration.py @@ -4,6 +4,7 @@ import pytest from fastapi import Request +from ogx_client.models.open_ai_response_object import OpenAIResponseObject from pytest_mock import MockerFixture import constants @@ -36,6 +37,7 @@ "created_at": 1700000000, "status": "completed", "model": "test-provider/test-model", + "store": False, "output": [ { "type": "message", @@ -68,8 +70,8 @@ def _build_responses_mock_client(mocker: MockerFixture) -> Any: """Build a mock client suitable for the /responses endpoint.""" mock_client = _build_base_mock_client(mocker) - mock_client.responses.create.return_value.model_dump.return_value = ( - _RESPONSE_DUMP.copy() + mock_client.responses.create = mocker.AsyncMock( + return_value=OpenAIResponseObject.from_dict(_RESPONSE_DUMP) ) return mock_client diff --git a/tests/integration/endpoints/test_responses_integration.py b/tests/integration/endpoints/test_responses_integration.py index 652f8ed22..960c0f829 100644 --- a/tests/integration/endpoints/test_responses_integration.py +++ b/tests/integration/endpoints/test_responses_integration.py @@ -11,6 +11,7 @@ import pytest from fastapi import Request from fastapi.responses import StreamingResponse +from ogx_client.models.open_ai_response_object import OpenAIResponseObject from pytest_mock import MockerFixture from sqlalchemy.orm import Session @@ -43,6 +44,7 @@ "created_at": 1700000000, "status": "completed", "model": "test-provider/test-model", + "store": False, "output": [ { "type": "message", @@ -77,21 +79,9 @@ def _build_mock_client(mocker: MockerFixture) -> Any: """ mock_client = mocker.AsyncMock() - mock_response = mocker.MagicMock() - mock_response.id = "resp_integ_test" - mock_output = mocker.MagicMock() - mock_output.type = "message" - mock_output.role = "assistant" - mock_output.content = "Ansible is an automation tool." - mock_output.refusal = None - mock_response.output = [mock_output] - mock_response.usage = mocker.MagicMock() - mock_response.usage.input_tokens = 10 - mock_response.usage.output_tokens = 5 - mock_response.status = "completed" - mock_response.model = "test-provider/test-model" - mock_response.model_dump.return_value = _RESPONSE_DUMP.copy() - mock_client.responses.create = mocker.AsyncMock(return_value=mock_response) + mock_client.responses.create = mocker.AsyncMock( + return_value=OpenAIResponseObject.from_dict(_RESPONSE_DUMP) + ) mock_client.openai.list.return_value = make_openai_models_list_response( make_openai_model() diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index b64131f16..3fa803586 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -822,7 +822,7 @@ async def test_handle_non_streaming_success_returns_response( mock_api_response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - mock_api_response.model_dump.return_value = { + serialized_response = { "id": "resp_1", "object": "response", "created_at": 0, @@ -838,6 +838,7 @@ async def test_handle_non_streaming_success_returns_response( }, } mock_client.responses.create = mocker.AsyncMock(return_value=mock_api_response) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_response) _patch_handle_non_streaming_common(mocker, minimal_config) mocker.patch( @@ -903,7 +904,7 @@ async def test_handle_non_streaming_with_previous_response_id_appends_turn( mock_api_response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - mock_api_response.model_dump.return_value = { + serialized_response = { "id": "resp_1", "object": "response", "created_at": 0, @@ -919,6 +920,7 @@ async def test_handle_non_streaming_with_previous_response_id_appends_turn( }, } mock_client.responses.create = mocker.AsyncMock(return_value=mock_api_response) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_response) _patch_handle_non_streaming_common(mocker, minimal_config) mocker.patch( @@ -1222,7 +1224,7 @@ async def test_handle_streaming_success_returns_sse_consumes_response_generator( mock_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - mock_chunk.model_dump.return_value = { + serialized_chunk = { "type": "response.completed", "response": {"id": "r1", "usage": {"input_tokens": 1}}, } @@ -1231,6 +1233,7 @@ async def mock_stream() -> Any: yield mock_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_chunk) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -1295,10 +1298,6 @@ async def test_handle_streaming_in_progress_chunk_sets_quotas_and_output_text( in_progress_chunk = mocker.Mock() in_progress_chunk.type = "response.in_progress" - in_progress_chunk.model_dump.return_value = { - "type": "response.in_progress", - "response": {"id": "r0"}, - } completed_chunk = mocker.Mock() completed_chunk.type = "response.completed" @@ -1308,7 +1307,11 @@ async def test_handle_streaming_in_progress_chunk_sets_quotas_and_output_text( completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { + serialized_in_progress_chunk = { + "type": "response.in_progress", + "response": {"id": "r0"}, + } + serialized_completed_chunk = { "type": "response.completed", "response": {"id": "r1", "usage": {"input_tokens": 1}}, } @@ -1318,6 +1321,10 @@ async def mock_stream() -> Any: yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", + side_effect=[serialized_in_progress_chunk, serialized_completed_chunk], + ) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -1389,7 +1396,7 @@ async def test_handle_streaming_builds_tool_call_summary_from_output( completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { + serialized_completed_chunk = { "type": "response.completed", "response": {"id": "r1", "usage": {"input_tokens": 1}}, } @@ -1398,6 +1405,9 @@ async def mock_stream() -> Any: yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", return_value=serialized_completed_chunk + ) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -1475,7 +1485,7 @@ async def test_handle_streaming_with_previous_response_id_appends_turn( completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { + serialized_completed_chunk = { "type": "response.completed", "response": {"id": "r1", "usage": {"input_tokens": 1}}, } @@ -1484,6 +1494,9 @@ async def mock_stream() -> Any: yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", return_value=serialized_completed_chunk + ) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -2276,7 +2289,7 @@ async def test_non_streaming_sanitizes_mcp_output_and_model( mock_api_response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - mock_api_response.model_dump.return_value = { + serialized_response = { "id": "resp_1", "object": "response", "created_at": 0, @@ -2304,6 +2317,7 @@ async def test_non_streaming_sanitizes_mcp_output_and_model( }, } mock_client.responses.create = mocker.AsyncMock(return_value=mock_api_response) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_response) mocker.patch(f"{MODULE}.configuration", mock_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -2374,7 +2388,7 @@ def _make_streaming_completed_chunk(self, mocker: MockerFixture) -> Any: completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { + completed_chunk.serialized = { "type": "response.completed", "response": { "id": "r1", @@ -2433,6 +2447,10 @@ async def mock_stream() -> Any: yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", + return_value=completed_chunk.serialized, + ) mocker.patch(f"{MODULE}.configuration", mock_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -2536,9 +2554,9 @@ async def test_mcp_events_filtered_without_merge_server_tools_header( # pylint: mcp_added_chunk.type = "response.output_item.added" mcp_added_chunk.item = mcp_item mcp_added_chunk.output_index = 0 - mcp_added_chunk.model_dump.return_value = { - "type": "response.output_item.added", - "output_index": 0, + serialized_completed_chunk = { + "type": "response.completed", + "response": {"id": "r1"}, } completed_chunk = mocker.Mock() @@ -2549,16 +2567,16 @@ async def test_mcp_events_filtered_without_merge_server_tools_header( # pylint: completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { - "type": "response.completed", - "response": {"id": "r1"}, - } async def mock_stream() -> Any: yield mcp_added_chunk yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", + return_value=serialized_completed_chunk, + ) mocker.patch(f"{MODULE}.configuration", mock_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -2632,10 +2650,14 @@ async def test_mcp_events_filtered_with_no_mcp_servers_configured( text_added_chunk.type = "response.output_item.added" text_added_chunk.item = text_item text_added_chunk.output_index = 0 - text_added_chunk.model_dump.return_value = { + serialized_text_added_chunk = { "type": "response.output_item.added", "output_index": 0, } + serialized_completed_chunk = { + "type": "response.completed", + "response": {"id": "r1"}, + } completed_chunk = mocker.Mock() completed_chunk.type = "response.completed" @@ -2645,16 +2667,16 @@ async def test_mcp_events_filtered_with_no_mcp_servers_configured( completed_chunk.response.usage = mocker.Mock( input_tokens=1, output_tokens=2, total_tokens=3 ) - completed_chunk.model_dump.return_value = { - "type": "response.completed", - "response": {"id": "r1"}, - } async def mock_stream() -> Any: yield text_added_chunk yield completed_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch( + f"{MODULE}.dump_ogx_model", + side_effect=[serialized_text_added_chunk, serialized_completed_chunk], + ) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) @@ -2721,7 +2743,10 @@ async def test_response_generator_records_failure_when_stream_iteration_raises( ok_chunk = mocker.Mock() ok_chunk.type = "response.output_item.added" - ok_chunk.model_dump.return_value = {"type": "response.output_item.added"} + mocker.patch( + f"{MODULE}.dump_ogx_model", + return_value={"type": "response.output_item.added"}, + ) async def failing_stream() -> AsyncIterator[Any]: """Async generator that optionally yields a chunk then raises.""" diff --git a/tests/unit/app/endpoints/test_responses_splunk.py b/tests/unit/app/endpoints/test_responses_splunk.py index 8ce3ea99c..f67757f68 100644 --- a/tests/unit/app/endpoints/test_responses_splunk.py +++ b/tests/unit/app/endpoints/test_responses_splunk.py @@ -384,7 +384,7 @@ async def test_non_streaming_success( mock_api_response.usage = mocker.Mock( input_tokens=100, output_tokens=50, total_tokens=150 ) - mock_api_response.model_dump.return_value = { + serialized_response = { "id": "resp_1", "object": "response", "created_at": 0, @@ -400,6 +400,7 @@ async def test_non_streaming_success( }, } mock_client.responses.create = mocker.AsyncMock(return_value=mock_api_response) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_response) _patch_handle_non_streaming_common(mocker, minimal_config) mocker.patch( @@ -613,7 +614,7 @@ async def test_streaming_success( mock_chunk.response.usage = mocker.Mock( input_tokens=100, output_tokens=50, total_tokens=150 ) - mock_chunk.model_dump.return_value = { + serialized_chunk = { "type": "response.completed", "response": { "id": "r1", @@ -626,6 +627,7 @@ async def mock_stream() -> Any: yield mock_chunk mock_client.responses.create = mocker.AsyncMock(return_value=mock_stream()) + mocker.patch(f"{MODULE}.dump_ogx_model", return_value=serialized_chunk) mocker.patch(f"{MODULE}.configuration", minimal_config) mocker.patch(f"{MODULE}.get_available_quotas", return_value={}) diff --git a/tests/unit/utils/test_ogx_serialization.py b/tests/unit/utils/test_ogx_serialization.py new file mode 100644 index 000000000..6dcfe851f --- /dev/null +++ b/tests/unit/utils/test_ogx_serialization.py @@ -0,0 +1,221 @@ +"""Unit tests for ogx_client serialization helpers.""" + +import json +from typing import Any + +import pytest +from ogx_client.models.open_ai_response_object import OpenAIResponseObject +from ogx_client.models.open_ai_response_object_stream import ( + OpenAIResponseObjectStream, +) +from ogx_client.models.open_ai_response_object_stream_response_completed import ( + OpenAIResponseObjectStreamResponseCompleted, +) +from ogx_client.models.open_ai_response_object_stream_response_output_item_added import ( + OpenAIResponseObjectStreamResponseOutputItemAdded, +) + +from models.api.responses.successful.responses_openai import ResponsesResponse +from utils.ogx_serialization import dump_ogx_model + + +@pytest.fixture(name="complex_client_response_payload") +def complex_client_response_payload_fixture() -> dict[str, Any]: + """A response payload that exercises nested OneOf fields in ogx_client.""" + return { + "id": "resp_complex", + "object": "response", + "created_at": 1_234_567_890, + "status": "completed", + "model": "provider/model", + "store": False, + "parallel_tool_calls": True, + "tool_choice": "auto", + "input": "multi-step query with tools", + "tools": [{"type": "file_search", "vector_store_ids": ["vs_1", "vs_2"]}], + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "total_tokens": 150, + "input_tokens_details": {"cached_tokens": 10}, + "output_tokens_details": {"reasoning_tokens": 5}, + }, + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Searching docs...", + "annotations": [], + } + ], + "status": "completed", + "id": "msg_1", + }, + { + "type": "file_search_call", + "id": "fs_1", + "status": "completed", + "queries": ["lightspeed", "quota"], + "results": [ + { + "file_id": "file_abc", + "filename": "guide.pdf", + "score": 0.91, + "text": "relevant chunk", + "attributes": {"page": 2, "section": "limits"}, + } + ], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_weather", + "name": "get_weather", + "arguments": '{"city":"NYC"}', + "status": "completed", + }, + { + "type": "mcp_call", + "id": "mcp_1", + "status": "completed", + "server_label": "portal", + "name": "search", + "arguments": "{}", + "output": "portal result", + }, + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Final answer.", + "annotations": [], + } + ], + "status": "completed", + "id": "msg_2", + }, + ], + } + + +def test_dump_ogx_model_complex_client_response( + complex_client_response_payload: dict[str, Any], +) -> None: + """A multi-item client response must dump all OneOf-backed fields correctly.""" + response = OpenAIResponseObject.from_dict(complex_client_response_payload) + + dumped = dump_ogx_model(response) + + json.dumps(dumped) + assert dumped["tool_choice"] == "auto" + assert dumped["tools"] == [ + {"type": "file_search", "vector_store_ids": ["vs_1", "vs_2"]} + ] + assert [item["type"] for item in dumped["output"]] == [ + "message", + "file_search_call", + "function_call", + "mcp_call", + "message", + ] + assert dumped["output"][0]["content"][0]["text"] == "Searching docs..." + assert dumped["output"][1]["results"][0]["attributes"] == { + "page": 2, + "section": "limits", + } + assert dumped["output"][2]["arguments"] == '{"city":"NYC"}' + assert dumped["output"][3]["server_label"] == "portal" + assert dumped["usage"]["input_tokens_details"]["cached_tokens"] == 10 + + validated = ResponsesResponse.model_validate( + { + **dumped, + "safety_identifier": "safety-id", + "available_quotas": {}, + "conversation": "conv-id", + "completed_at": 1, + "output_text": "Final answer.", + } + ) + assert validated.tool_choice == "auto" + assert validated.output_text == "Final answer." + + +def test_dump_ogx_model_model_dump_leaves_empty_oneof_wrappers( + complex_client_response_payload: dict[str, Any], +) -> None: + """Plain model_dump is the failure mode dump_ogx_model exists to fix.""" + response = OpenAIResponseObject.from_dict(complex_client_response_payload) + + broken = response.model_dump(exclude_none=True) + fixed = dump_ogx_model(response) + + assert broken["tool_choice"] == {} + assert fixed["tool_choice"] == "auto" + + added = OpenAIResponseObjectStreamResponseOutputItemAdded.from_dict( + { + "type": "response.output_item.added", + "output_index": 0, + "sequence_number": 1, + "response_id": "resp_complex", + "item": complex_client_response_payload["output"][0], + } + ) + stream_chunk = OpenAIResponseObjectStream(actual_instance=added) + + broken_chunk = stream_chunk.model_dump(exclude_none=True, by_alias=True) + fixed_chunk = dump_ogx_model(stream_chunk) + + assert broken_chunk["item"] == {} + assert fixed_chunk["item"]["content"][0]["text"] == "Searching docs..." + + +def test_dump_ogx_model_complex_streaming_chunks( + complex_client_response_payload: dict[str, Any], +) -> None: + """Streaming wrappers must preserve nested tool-call and response payloads.""" + mcp_item = complex_client_response_payload["output"][3] + added = OpenAIResponseObjectStreamResponseOutputItemAdded.from_dict( + { + "type": "response.output_item.added", + "output_index": 3, + "sequence_number": 4, + "response_id": "resp_complex", + "item": mcp_item, + } + ) + completed = OpenAIResponseObjectStreamResponseCompleted.from_dict( + { + "type": "response.completed", + "sequence_number": 10, + "response": complex_client_response_payload, + } + ) + + added_dump = dump_ogx_model(OpenAIResponseObjectStream(actual_instance=added)) + completed_dump = dump_ogx_model( + OpenAIResponseObjectStream(actual_instance=completed) + ) + + json.dumps(added_dump) + json.dumps(completed_dump) + + assert added_dump["item"] == { + "id": "mcp_1", + "type": "mcp_call", + "arguments": "{}", + "name": "search", + "server_label": "portal", + "output": "portal result", + } + assert completed_dump["response"]["tool_choice"] == "auto" + assert len(completed_dump["response"]["output"]) == 5 + assert completed_dump["response"]["output"][1]["results"][0]["filename"] == ( + "guide.pdf" + )