diff --git a/src/pydantic_ai_lightspeed/llamastack/_model.py b/src/pydantic_ai_lightspeed/llamastack/_model.py index 3df0b9640..2338d241e 100644 --- a/src/pydantic_ai_lightspeed/llamastack/_model.py +++ b/src/pydantic_ai_lightspeed/llamastack/_model.py @@ -26,7 +26,7 @@ from pydantic_ai import UnexpectedModelBehavior from pydantic_ai._run_context import RunContext from pydantic_ai._utils import PeekableAsyncStream, Unset, number_to_datetime -from pydantic_ai.messages import ModelMessage +from pydantic_ai.messages import ModelMessage, ModelResponse from pydantic_ai.models import ( ModelRequestParameters, StreamedResponse, @@ -181,15 +181,79 @@ class LlamaStackResponsesModel(OpenAIResponsesModel): before the corresponding ``McpCall`` or ``ResponseFunctionToolCall`` item. """ + async def request( # pylint: disable=unused-argument + self, + messages: list[ModelMessage], + model_settings: ModelSettings | None, + model_request_parameters: ModelRequestParameters, + run_context: RunContext[Any] | None = None, + ) -> Any: + """Non-streaming request with Llama Stack conversation continuation fix. + + Llama Stack rejects requests containing both ``conversation`` and + ``previous_response_id``. On continuation turns (where a prior + ``ModelResponse`` exists), we trim messages to only the new input and + disable ``previous_response_id`` so that only ``conversation`` is sent. + This ensures all responses are persisted to the conversation. + """ + messages, model_settings = self._prepare_conversation_continuation( + messages, model_settings + ) + return await super().request(messages, model_settings, model_request_parameters) + + def _prepare_conversation_continuation( + self, + messages: list[ModelMessage], + model_settings: ModelSettings | None, + ) -> tuple[list[ModelMessage], ModelSettings | None]: + """Trim messages and disable previous_response_id for conversation continuations. + + Llama Stack rejects requests with both ``previous_response_id`` and + ``conversation``. When ``conversation`` is in ``extra_body`` and there's + already a ModelResponse in the history (a continuation turn), we: + + 1. Trim messages to only those AFTER the last ModelResponse (new input only) + 2. Disable ``openai_previous_response_id`` so pydantic-ai won't resolve one + + This means Llama Stack receives ``conversation`` (for persistence) plus only + the new input items. Llama Stack reconstructs prior history from the + conversation and appends the new input correctly. + """ + if not model_settings or not isinstance(model_settings, dict): + return messages, model_settings + + extra_body = model_settings.get("extra_body") + if not isinstance(extra_body, dict) or "conversation" not in extra_body: + return messages, model_settings + + last_response_idx = None + for i in range(len(messages) - 1, -1, -1): + msg = messages[i] + if isinstance(msg, ModelResponse) and msg.provider_response_id: + last_response_idx = i + break + + if last_response_idx is None: + return messages, model_settings + + trimmed_messages = messages[last_response_idx + 1 :] + + new_settings = dict(model_settings) + new_settings.pop("openai_previous_response_id", None) + return trimmed_messages, cast(ModelSettings, new_settings) + @asynccontextmanager - async def request_stream( + async def request_stream( # pylint: disable=unused-argument self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[Any] | None = None, ) -> AsyncIterator[StreamedResponse]: - """Request a streaming response, filtering Llama Stack-specific event quirks. + """Request a streaming response with Llama Stack compatibility fixes. + + Applies the same conversation continuation handling as :meth:`request` + before calling the Responses API, then filters streaming tool-call events. Args: messages: Model messages for the request. @@ -201,10 +265,10 @@ async def request_stream( A StreamedResponse with the filtered event stream. """ check_allow_model_requests() - model_settings, model_request_parameters = self.prepare_request( - model_settings, - model_request_parameters, + messages, model_settings = self._prepare_conversation_continuation( + messages, model_settings ) + model_settings_cast = cast(OpenAIResponsesModelSettings, model_settings or {}) response = await self._responses_create( messages, True, model_settings_cast, model_request_parameters diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 138852bc2..9b0c41226 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -221,12 +221,29 @@ async def generate_agent_response( context.query_request.conversation_id is None and bool(context.query_request.generate_topic_summary) ) - topic_summary = await maybe_get_topic_summary( - generate_topic_summary=should_generate_topic_summary, - input_text=context.query_request.query, - client=context.client, - model_id=responses_params.model, - ) + try: + topic_summary = await maybe_get_topic_summary( + generate_topic_summary=should_generate_topic_summary, + input_text=context.query_request.query, + client=context.client, + model_id=responses_params.model, + ) + except HTTPException as exc: + logger.warning( + "Topic summary failed for request %s: %s", + context.request_id, + exc.detail, + ) + detail: dict[str, str] = exc.detail if isinstance(exc.detail, dict) else {} + yield serialize_event( + ErrorStreamPayload.create( + status_code=exc.status_code, + response=detail.get("response", "Internal server error"), + cause=detail.get("cause", str(exc.detail)), + ), + media_type, + ) + return logger.info("Consuming tokens") consume_query_tokens( user_id=context.user_id, diff --git a/src/utils/pydantic_ai.py b/src/utils/pydantic_ai.py index f4a1cf18c..ef2e39d54 100644 --- a/src/utils/pydantic_ai.py +++ b/src/utils/pydantic_ai.py @@ -21,7 +21,6 @@ { "conversation", "max_infer_iters", - "tools", "tool_choice", "include", "text", @@ -68,6 +67,8 @@ def _model_settings_from_responses_params( if responses_params.extra_headers: settings_dict["extra_headers"] = dict(responses_params.extra_headers) settings_dict["openai_store"] = responses_params.store + if responses_params.tools is not None: + settings_dict["openai_native_tools"] = responses_params.tools if responses_params.previous_response_id is not None: settings_dict["openai_previous_response_id"] = ( responses_params.previous_response_id diff --git a/src/utils/query.py b/src/utils/query.py index 51e246d91..4ccb3abeb 100644 --- a/src/utils/query.py +++ b/src/utils/query.py @@ -574,10 +574,9 @@ def handle_known_apistatus_errors( Returns: AbstractErrorResponse: The error response model. """ - if error.status_code == 400: - error_message = getattr(error, "message", str(error)) - if is_context_length_error(error_message): - return PromptTooLongResponse(model=model_id) - elif error.status_code == 429: + error_message = getattr(error, "message", str(error)) + if is_context_length_error(error_message): + return PromptTooLongResponse(model=model_id) + if error.status_code == 429: return QuotaExceededResponse.model(model_id) return InternalServerErrorResponse.generic() diff --git a/tests/unit/utils/test_pydantic_ai.py b/tests/unit/utils/test_pydantic_ai.py index 2a9a59d52..219b3fc57 100644 --- a/tests/unit/utils/test_pydantic_ai.py +++ b/tests/unit/utils/test_pydantic_ai.py @@ -82,6 +82,7 @@ def minimal_params_fixture(self, mocker: MockerFixture) -> object: params.parallel_tool_calls = None params.extra_headers = None params.store = False + params.tools = None params.previous_response_id = None return params @@ -138,7 +139,6 @@ def test_extra_body_from_lls_fields(self, mocker: MockerFixture) -> None: "model": "test/model", "conversation": "conv-123", "max_infer_iters": 5, - "tools": [{"type": "function"}], "tool_choice": "auto", } params.max_output_tokens = None @@ -147,14 +147,15 @@ def test_extra_body_from_lls_fields(self, mocker: MockerFixture) -> None: params.extra_headers = None params.store = False params.previous_response_id = None + params.tools = [{"type": "function"}] settings = _model_settings_from_responses_params(params) assert "extra_body" in settings assert settings["extra_body"]["conversation"] == "conv-123" assert settings["extra_body"]["max_infer_iters"] == 5 - assert settings["extra_body"]["tools"] == [{"type": "function"}] assert settings["extra_body"]["tool_choice"] == "auto" + assert settings["openai_native_tools"] == [{"type": "function"}] def test_extra_body_only_includes_known_fields(self, mocker: MockerFixture) -> None: """Test that extra_body only includes fields in _LLS_RESPONSES_EXTRA_FIELDS.""" @@ -189,7 +190,6 @@ def test_contains_expected_fields(self) -> None: expected = { "conversation", "max_infer_iters", - "tools", "tool_choice", "include", "text",