From 18500695622cc3d576f97f93ca3d8d23fc41b492 Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Mon, 17 Aug 2026 11:49:54 +0200 Subject: [PATCH] LCORE-3582: fix compacted-mode 500s in the agent pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Once a conversation compacted, every subsequent request on it failed with HTTP 500 on both /v1/query and /v1/streaming_query, permanently bricking the conversation. Root cause: in compacted mode (LCORE-1572) CompactionResult.params.input is an explicit item list (summaries + recent verbatim turns + new query) with the conversation parameter omitted, but the pydantic-ai agent pipeline did prompt = cast(str, responses_params.input) and handed the list to agent.run(), which dies client-side before any request reaches Llama Stack. The A2A executor had a quieter variant of the same gap: it passes the raw user text as the prompt, so compacted A2A turns silently lost all conversation context. The fix makes the explicit input reach the wire through the existing params-to-model-settings seam: - _model_settings_from_responses_params (pydantic_ai_lightspeed/ llamastack/_model.py): when omit_conversation is set and input is an item list, the dumped list is added to extra_body. The OpenAI SDK merges extra_body into the request body with precedence (_merge_mappings: "the second mapping takes precedence"), so the explicit list replaces the prompt-derived input — the request body matches what the non-agent /v1/responses path sends. This also fixes the A2A context loss with no a2a.py changes, since build_agent already routes params through this seam. - OgxResponsesModel gains _prepare_compacted_input, applied in both request() and request_stream() after the conversation-continuation trim: once a ModelResponse exists in the message history (a client-side tool-loop continuation), the input override is dropped so pydantic-ai's mapped messages — which carry the tool results — win. - New agent_prompt_text() helper (utils/conversation_compaction.py) replaces the cast(str, ...) at all four call sites (non-streaming and streaming, plain and multimodal): returns input unchanged when it is a string, else the text of the trailing message item of the explicit list. The prompt still drives capabilities and multimodal input construction; the wire input comes from the override. - The blocked-moderation path in retrieve_agent_response now skips append_turn_items_to_conversation when omit_conversation is set, mirroring the streaming path — appending the full explicit list would duplicate summaries and history into the conversation. - map_agent_inference_error now logs the original exception at error level with the traceback before mapping. Previously the mapped HTTPException discarded it and callers raised without logging, so these failures produced a generic 500 with nothing in the logs — which is what made this bug expensive to diagnose. Verified live against llama-stack 0.6.0: on a compacted conversation, turn after turn returns 200 on both endpoints (previously 500), the streaming path emits compaction/token/end events, and the model correctly answers questions about pre-compaction turns, proving the summary context reaches the model. Known limitation: image attachments are not folded into the overridden explicit input, so a compacted turn with images sends the text-only list (compaction+images previously hard -failed; noted in LCORE-3582). Unit tests cover the extra_body override (present in compacted mode, absent otherwise), the tool-loop guard, agent_prompt_text, the prompt threading through both retrieve paths, the moderation-append guard, and the error logging. --- .../llamastack/_model.py | 40 ++++++++ src/utils/agents/error_handler.py | 3 + src/utils/agents/query.py | 20 ++-- src/utils/agents/streaming.py | 7 +- src/utils/conversation_compaction.py | 28 ++++++ .../llamastack/test_model.py | 81 ++++++++++++++++ tests/unit/utils/agents/test_query.py | 92 +++++++++++++++++++ tests/unit/utils/agents/test_streaming.py | 52 +++++++++++ .../utils/test_conversation_compaction.py | 21 +++++ 9 files changed, 332 insertions(+), 12 deletions(-) diff --git a/src/pydantic_ai_lightspeed/llamastack/_model.py b/src/pydantic_ai_lightspeed/llamastack/_model.py index 1782a15c3..496bce278 100644 --- a/src/pydantic_ai_lightspeed/llamastack/_model.py +++ b/src/pydantic_ai_lightspeed/llamastack/_model.py @@ -74,6 +74,16 @@ def _model_settings_from_responses_params( """Map ``ResponsesApiParams`` into Pydantic AI OpenAI Responses model settings.""" payload = responses_params.model_dump(exclude_none=True) extra_body = {k: v for k, v in payload.items() if k in _LLS_RESPONSES_EXTRA_FIELDS} + if responses_params.omit_conversation and not isinstance( + responses_params.input, str + ): + # Compacted mode (LCORE-3582): the request must carry the explicit item + # list (summaries + recent turns + new query), but pydantic-ai builds + # the wire ``input`` from the prompt alone. Overriding via extra_body + # replaces it with the explicit list, exactly as the non-agent + # /v1/responses path sends it. Dropped again on tool-loop + # continuations — see ``_prepare_compacted_input``. + extra_body["input"] = payload["input"] settings_dict: dict[str, Any] = {} if extra_body: settings_dict["extra_body"] = extra_body @@ -291,6 +301,7 @@ async def request( # pylint: disable=unused-argument messages, model_settings = self._prepare_conversation_continuation( messages, model_settings ) + model_settings = self._prepare_compacted_input(messages, model_settings) return await super().request(messages, model_settings, model_request_parameters) def _prepare_conversation_continuation( @@ -334,6 +345,34 @@ def _prepare_conversation_continuation( new_settings.pop("openai_previous_response_id", None) return trimmed_messages, cast(ModelSettings, new_settings) + def _prepare_compacted_input( + self, + messages: list[ModelMessage], + model_settings: Optional[ModelSettings], + ) -> Optional[ModelSettings]: + """Drop the compacted ``input`` override on tool-loop continuations. + + In compacted mode (LCORE-3582) the request body ``input`` is overridden + via ``extra_body`` with the explicit item list. That override is only + valid for the first request of an agent run: on client-side tool-loop + iterations pydantic-ai's mapped messages carry the tool results and + must win, so the override is removed once a ``ModelResponse`` exists in + the message history. + """ + if not model_settings or not isinstance(model_settings, dict): + return model_settings + extra_body = model_settings.get("extra_body") + if not isinstance(extra_body, dict) or "input" not in extra_body: + return model_settings + if not any(isinstance(message, ModelResponse) for message in messages): + return model_settings + + new_extra_body = dict(extra_body) + new_extra_body.pop("input") + new_settings = dict(model_settings) + new_settings["extra_body"] = new_extra_body + return cast(ModelSettings, new_settings) + @asynccontextmanager async def request_stream( # pylint: disable=unused-argument self, @@ -360,6 +399,7 @@ async def request_stream( # pylint: disable=unused-argument messages, model_settings = self._prepare_conversation_continuation( messages, model_settings ) + model_settings = self._prepare_compacted_input(messages, model_settings) model_settings_cast = cast(OpenAIResponsesModelSettings, model_settings or {}) response = await self._responses_create( diff --git a/src/utils/agents/error_handler.py b/src/utils/agents/error_handler.py index 8480ebb30..f99e0e6fe 100644 --- a/src/utils/agents/error_handler.py +++ b/src/utils/agents/error_handler.py @@ -48,6 +48,9 @@ def map_agent_inference_error( RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is not a recognized context-length failure. """ + # The mapped HTTPException loses the original exception, and callers raise + # it without logging — log here so failures are diagnosable (LCORE-3582). + logger.error("Agent inference failed: %s", exc, exc_info=exc) match exc: case AgentRunError() as agent_exc: return map_pydantic_agent_run_error(agent_exc, model_id) diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index dabc956cf..7dd2d9e43 100644 --- a/src/utils/agents/query.py +++ b/src/utils/agents/query.py @@ -3,7 +3,7 @@ from __future__ import annotations from enum import Enum -from typing import Optional, cast +from typing import Optional from fastapi import HTTPException from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient @@ -36,6 +36,7 @@ process_native_tool_call, process_native_tool_result, ) +from utils.conversation_compaction import agent_prompt_text from utils.conversations import append_turn_items_to_conversation from utils.otel_tracing import ( SpanAttributes, @@ -276,12 +277,13 @@ async def retrieve_agent_response( ) if moderation_result.decision == "blocked": - await append_turn_items_to_conversation( - client, - responses_params.conversation, - responses_params.input, - [moderation_result.refusal_response], - ) + if not responses_params.omit_conversation: + await append_turn_items_to_conversation( + client, + responses_params.conversation, + responses_params.input, + [moderation_result.refusal_response], + ) return TurnSummary( id=moderation_result.moderation_id, llm_response=moderation_result.message, @@ -301,11 +303,11 @@ async def retrieve_agent_response( logger.debug("Starting agent non-streaming response processing") if image_attachments: prompt = build_multimodal_input( - cast(str, responses_params.input), + agent_prompt_text(responses_params), image_attachments, ) else: - prompt = cast(str, responses_params.input) + prompt = agent_prompt_text(responses_params) run_result = await agent.run(prompt) except ( AgentRunError, diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 55666447d..90f320d96 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -8,7 +8,7 @@ import datetime from collections.abc import AsyncIterator from functools import singledispatch -from typing import Any, Final, Optional, cast +from typing import Any, Final, Optional from fastapi import HTTPException from ogx_client import APIConnectionError, APIStatusError @@ -59,6 +59,7 @@ process_native_tool_call, process_native_tool_result, ) +from utils.conversation_compaction import agent_prompt_text from utils.conversations import append_turn_items_to_conversation from utils.pydantic_ai_helpers import build_agent from utils.query import ( @@ -334,11 +335,11 @@ async def agent_response_generator( ) if image_attachments: prompt = build_multimodal_input( - cast(str, responses_params.input), + agent_prompt_text(responses_params), image_attachments, ) else: - prompt = cast(str, responses_params.input) + prompt = agent_prompt_text(responses_params) logger.debug("Starting agent streaming response processing") async with agent.run_stream_events(prompt) as stream: diff --git a/src/utils/conversation_compaction.py b/src/utils/conversation_compaction.py index e45cefda4..62c5b9f07 100644 --- a/src/utils/conversation_compaction.py +++ b/src/utils/conversation_compaction.py @@ -243,6 +243,34 @@ def _verbatim_input_message(item: Any) -> Optional[OpenAIResponseMessage]: return OpenAIResponseMessage(role=cast(Any, role), content=text) +def agent_prompt_text(params: ResponsesApiParams) -> str: + """Return the textual user prompt for a pydantic-ai agent run. + + In compacted mode ``params.input`` is the explicit item list built by + :func:`_build_explicit_input` (summaries + recent turns + new query), so + the new user query is the trailing message item. The agent pipeline still + needs a plain string prompt (capabilities and multimodal input operate on + it); the full explicit list reaches the request body separately via the + ``extra_body`` input override (LCORE-3582). + + Args: + params: Prepared (possibly compaction-rewritten) request parameters. + + Returns: + ``params.input`` unchanged when it is a string; otherwise the text of + the last message item in the explicit list, or ``""`` when there is + none. + """ + if isinstance(params.input, str): + return params.input + for item in reversed(list(params.input)): + if is_message_item(item): + text = extract_message_text(item) + if text: + return text + return "" + + def _query_input_message(original_input: ResponseInput) -> list[Any]: """Render the new user query as explicit input items. diff --git a/tests/unit/pydantic_ai_lightspeed/llamastack/test_model.py b/tests/unit/pydantic_ai_lightspeed/llamastack/test_model.py index 9a059cd72..3773c37d9 100644 --- a/tests/unit/pydantic_ai_lightspeed/llamastack/test_model.py +++ b/tests/unit/pydantic_ai_lightspeed/llamastack/test_model.py @@ -6,6 +6,7 @@ from typing import Any import pytest +from ogx_api.openai_responses import OpenAIResponseMessage from openai.types import responses from pydantic_ai import ModelMessage, UnexpectedModelBehavior from pydantic_ai.messages import ModelResponse @@ -120,6 +121,86 @@ def test_none_fields_excluded(self) -> None: assert "extra_headers" not in settings assert "openai_previous_response_id" not in settings + def test_compacted_input_overrides_via_extra_body(self) -> None: + """Test compacted params carry the explicit input list in extra_body.""" + items = [ + OpenAIResponseMessage( + role="user", content="Summary of earlier conversation:\nS1" + ), + OpenAIResponseMessage(role="user", content="new question"), + ] + params = _make_params(input=items, omit_conversation=True) + settings = _model_settings_from_responses_params(params) + extra_body = settings["extra_body"] + assert "conversation" not in extra_body + assert extra_body["input"] == [ + { + "role": "user", + "content": "Summary of earlier conversation:\nS1", + "type": "message", + }, + {"role": "user", "content": "new question", "type": "message"}, + ] + + def test_string_input_never_lands_in_extra_body(self) -> None: + """Test that a plain string input is not duplicated into extra_body.""" + params = _make_params(input="hello", omit_conversation=True) + settings = _model_settings_from_responses_params(params) + assert "input" not in settings.get("extra_body", {}) + + def test_non_compacted_list_input_not_in_extra_body(self) -> None: + """Test that without omit_conversation the input stays out of extra_body.""" + items = [OpenAIResponseMessage(role="user", content="q")] + params = _make_params(input=items, omit_conversation=False) + settings = _model_settings_from_responses_params(params) + assert "input" not in settings.get("extra_body", {}) + + +class TestPrepareCompactedInput: + """Tests for the compacted-input tool-loop guard.""" + + @pytest.fixture(name="model") + def model_fixture(self, mocker: MockerFixture) -> OgxResponsesModel: + """Create a OgxResponsesModel with mocked __init__.""" + mocker.patch.object(OgxResponsesModel, "__init__", return_value=None) + return OgxResponsesModel("test-model") + + def test_no_input_override_returns_unchanged( + self, model: OgxResponsesModel, mocker: MockerFixture + ) -> None: + """Test settings without an input override pass through untouched.""" + messages = [mocker.Mock()] + settings: ModelSettings = {"extra_body": {"max_infer_iters": 5}} + assert model._prepare_compacted_input(messages, settings) is settings + + def test_first_request_keeps_input_override( + self, model: OgxResponsesModel, mocker: MockerFixture + ) -> None: + """Test the override survives when no ModelResponse is in messages.""" + messages = [mocker.Mock(spec=[])] + settings: ModelSettings = { + "extra_body": {"input": [{"role": "user", "content": "q"}]} + } + assert model._prepare_compacted_input(messages, settings) is settings + + def test_tool_loop_continuation_drops_input_override( + self, model: OgxResponsesModel + ) -> None: + """Test the override is dropped once a ModelResponse exists.""" + messages: list[ModelMessage] = [ModelResponse(parts=[])] + settings: ModelSettings = { + "extra_body": { + "input": [{"role": "user", "content": "q"}], + "max_infer_iters": 5, + } + } + result = model._prepare_compacted_input(messages, settings) + assert result is not settings + assert "input" not in result["extra_body"] + assert result["extra_body"]["max_infer_iters"] == 5 + # original settings untouched + assert "input" in settings["extra_body"] + class TestFromOgxClient: """Tests for OgxResponsesModel.from_ogx_client factory.""" diff --git a/tests/unit/utils/agents/test_query.py b/tests/unit/utils/agents/test_query.py index 97aa40649..775a2fe5b 100644 --- a/tests/unit/utils/agents/test_query.py +++ b/tests/unit/utils/agents/test_query.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException +from ogx_api.openai_responses import OpenAIResponseMessage from ogx_client import APIConnectionError, APIStatusError from pydantic_ai.messages import ( FinishReason, @@ -427,6 +428,97 @@ async def test_success_returns_turn_summary( assert summary.llm_response == "Hello!" assert summary.id == "resp-success" + @pytest.mark.asyncio + async def test_compacted_input_runs_agent_with_prompt_text( + self, + mocker: MockerFixture, + make_agent_run_result: Callable[..., Any], + make_responses_params: Callable[..., ResponsesApiParams], + patch_recording_metrics: None, + ) -> None: + """Test compacted explicit input is reduced to the query text for agent.run.""" + explicit = [ + OpenAIResponseMessage( + role="user", content="Summary of earlier conversation:\nS1" + ), + OpenAIResponseMessage(role="user", content="new question"), + ] + params = make_responses_params(input_text="ignored").model_copy( + update={"input": explicit, "omit_conversation": True} + ) + run_result = make_agent_run_result(content="Answer") + mock_agent = mocker.AsyncMock() + mock_agent.run = mocker.AsyncMock(return_value=run_result) + mocker.patch("utils.agents.query.build_agent", return_value=mock_agent) + + summary = await retrieve_agent_response( + client=mocker.AsyncMock(), + responses_params=params, + moderation_result=ShieldModerationPassed(), + endpoint_path=ENDPOINT_PATH_QUERY, + ) + + mock_agent.run.assert_awaited_once_with("new question") + assert summary.llm_response == "Answer" + + @pytest.mark.asyncio + async def test_blocked_moderation_compacted_skips_append( + self, + mocker: MockerFixture, + make_responses_params: Callable[..., ResponsesApiParams], + blocked_moderation: ShieldModerationBlocked, + ) -> None: + """Test blocked moderation does not append explicit input in compacted mode.""" + params = make_responses_params().model_copy( + update={ + "input": [OpenAIResponseMessage(role="user", content="q")], + "omit_conversation": True, + } + ) + mock_append = mocker.patch( + "utils.agents.query.append_turn_items_to_conversation", + new=mocker.AsyncMock(), + ) + + summary = await retrieve_agent_response( + client=mocker.AsyncMock(), + responses_params=params, + moderation_result=blocked_moderation, + endpoint_path=ENDPOINT_PATH_QUERY, + ) + + mock_append.assert_not_awaited() + assert summary.llm_response == "Content blocked by shield." + + @pytest.mark.asyncio + async def test_inference_error_is_logged( + self, + mocker: MockerFixture, + make_responses_params: Callable[..., ResponsesApiParams], + patch_recording_metrics: None, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Test mapped agent inference errors are logged before raising.""" + mock_agent = mocker.AsyncMock() + mock_agent.run = mocker.AsyncMock( + side_effect=APIConnectionError(request=mocker.Mock()) + ) + mocker.patch("utils.agents.query.build_agent", return_value=mock_agent) + + with caplog.at_level("ERROR"): + with pytest.raises(HTTPException): + await retrieve_agent_response( + client=mocker.AsyncMock(), + responses_params=make_responses_params(), + moderation_result=ShieldModerationPassed(), + endpoint_path=ENDPOINT_PATH_QUERY, + ) + + assert any( + record.levelname == "ERROR" and "Agent inference failed" in record.message + for record in caplog.records + ) + @pytest.mark.asyncio async def test_success_with_image_attachments_sends_multimodal_prompt( self, diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index f452f8106..c9c83a809 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -10,6 +10,7 @@ import pytest from fastapi import HTTPException +from ogx_api.openai_responses import OpenAIResponseMessage from ogx_client import APIStatusError from pydantic_ai import AgentRunResultEvent from pydantic_ai.exceptions import AgentRunError @@ -886,6 +887,57 @@ async def test_streams_token_events_and_updates_summary( assert turn_summary.token_usage.input_tokens == 4 assert turn_summary.token_usage.output_tokens == 2 + @pytest.mark.asyncio + async def test_compacted_input_streams_with_prompt_text( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + make_responses_params: Callable[..., ResponsesApiParams], + make_agent_run_result: Callable[..., Any], + patch_recording_metrics: None, + ) -> None: + """Test compacted explicit input is reduced to the query text for streaming.""" + context = make_generator_context() + turn_summary = TurnSummary() + run_result = make_agent_run_result(content="Answer", response_id="resp-c1") + events = [ + PartStartEvent(index=0, part=TextPart(content="Answer")), + AgentRunResultEvent(result=run_result), + ] + mock_agent = mocker.Mock() + mock_agent.run_stream_events.return_value = _mock_run_stream(events) + mocker.patch( + "utils.agents.streaming.get_agent_finish_reason", + return_value=AgentFinishReason.SUCCESS, + ) + mocker.patch( + "utils.agents.streaming.deduplicate_referenced_documents", + side_effect=lambda docs: docs, + ) + + explicit = [ + OpenAIResponseMessage( + role="user", content="Summary of earlier conversation:\nS1" + ), + OpenAIResponseMessage(role="user", content="new question"), + ] + params = make_responses_params(input_text="ignored").model_copy( + update={"input": explicit, "omit_conversation": True} + ) + + _ = [ + event + async for event in agent_response_generator( + mock_agent, + params, + context, + turn_summary, + ENDPOINT_PATH_STREAMING_QUERY, + ) + ] + + assert mock_agent.run_stream_events.call_args[0][0] == "new question" + @pytest.mark.asyncio async def test_streams_with_image_attachments_passes_multimodal_prompt( self, diff --git a/tests/unit/utils/test_conversation_compaction.py b/tests/unit/utils/test_conversation_compaction.py index 29678d691..1a35cdfd2 100644 --- a/tests/unit/utils/test_conversation_compaction.py +++ b/tests/unit/utils/test_conversation_compaction.py @@ -121,6 +121,27 @@ def test_should_compact() -> None: ) +def test_agent_prompt_text_string_input() -> None: + """A plain string input is returned unchanged.""" + assert cc.agent_prompt_text(_params("what is a pod?")) == "what is a pod?" + + +def test_agent_prompt_text_explicit_list_returns_last_message_text() -> None: + """For compacted explicit input, the trailing user query text is returned.""" + params = _params() + explicit = cc._build_explicit_input( + ["earlier summary"], [_msg("assistant", "prior answer")], "new question" + ) + compacted = params.model_copy(update={"input": explicit, "omit_conversation": True}) + assert cc.agent_prompt_text(compacted) == "new question" + + +def test_agent_prompt_text_empty_list_returns_empty() -> None: + """An empty explicit list yields an empty prompt rather than crashing.""" + params = _params().model_copy(update={"input": [], "omit_conversation": True}) + assert cc.agent_prompt_text(params) == "" + + # --- apply_compaction ---