diff --git a/tests/models/test_openai_mixed_stream.py b/tests/models/test_openai_mixed_stream.py new file mode 100644 index 000000000..42e5e39e5 --- /dev/null +++ b/tests/models/test_openai_mixed_stream.py @@ -0,0 +1,266 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from trpc_agent_sdk.models import LlmRequest, OpenAIModel +from trpc_agent_sdk.models._openai_model import const + + +async def collect(deltas, streaming_names=None, usage_only=True): + chunks = [{"id": "synthetic-response", "choices": [{"delta": delta, "finish_reason": None}]} for delta in deltas] + usage = {"prompt_tokens": 4, "completion_tokens": 6, "total_tokens": 10} + chunks.append({"choices": [{"delta": {}, "finish_reason": "tool_calls"}]}) + if usage_only: + chunks.append({"choices": [], "usage": usage}) + else: + chunks[0]["usage"] = usage + + async def stream(): + for chunk in chunks: + yield Mock(model_dump=lambda value=chunk: value) + + client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=AsyncMock(return_value=stream()))), + close=AsyncMock()) + model = OpenAIModel(model_name="synthetic", api_key="synthetic") + request = LlmRequest() + request.streaming_tool_names = streaming_names + with patch.object(OpenAIModel, "_create_async_client", return_value=client): + return [response async for response in model._generate_stream({}, request)] + + +def tool_delta(index, arguments, first=True): + result = {"index": index, "function": {"arguments": arguments}} + if first: + result.update(id=f"call_{index}", type="function") + result["function"]["name"] = f"tool_{index}" + return result + + +@pytest.mark.parametrize("text", [{}, { + "content": "text" +}, { + "reasoning_content": "synthetic" +}, { + "content": "text", + "reasoning_content": "synthetic" +}, { + "content": "", + "reasoning_content": None +}]) +@pytest.mark.parametrize("fragmented", [False, True]) +@pytest.mark.parametrize("tool_count", [1, 2]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.asyncio +async def test_mixed_stream_preserves_calls_and_text(text, fragmented, tool_count, streaming): + arguments = json.dumps({"value": "中文", "count": 12}, ensure_ascii=False) + fragments = [arguments[:12], arguments[12:]] if fragmented else [arguments] + deltas = [ + dict(text, tool_calls=[tool_delta(index, fragment, first=part == 0) for index in range(tool_count)]) + for part, fragment in enumerate(fragments) + ] + responses = await collect(deltas, {"tool_0"} if streaming else set(), usage_only=fragmented) + final = responses[-1] + assert not final.partial + assert final.error_code is None + assert final.response_id == "synthetic-response" + assert final.usage_metadata.total_token_count == 10 + calls = [part.function_call for part in final.content.parts if part.function_call] + assert [(call.id, call.name, call.args) for call in calls] == [(f"call_{index}", f"tool_{index}", { + "value": "中文", + "count": 12 + }) for index in range(tool_count)] + for thought, field in [(False, "content"), (True, "reasoning_content")]: + expected = (text.get(field) or "") * len(fragments) + assert "".join(part.text for part in final.content.parts + if part.text and bool(part.thought) == thought) == expected + assert "".join(part.text for response in responses[:-1] if response.content for part in response.content.parts + if part.text and bool(part.thought) == thought) == expected + partial_calls = [ + part.function_call for response in responses[:-1] if response.content for part in response.content.parts + if part.function_call + ] + assert [(call.id, call.name, call.args) for call in partial_calls] == ([("call_0", "tool_0", { + const.TOOL_STREAMING_ARGS: fragment + }) for fragment in fragments] if streaming else []) + + +@pytest.mark.asyncio +async def test_separate_text_and_tool_chunks(): + responses = await collect([{"content": "text"}, {"tool_calls": [tool_delta(0, '{"ok": true}')]}]) + parts = responses[-1].content.parts + assert parts[0].text == "text" + assert parts[1].function_call.args == {"ok": True} + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.asyncio +async def test_invalid_json_preserves_error_envelope(mixed): + delta = {"tool_calls": [tool_delta(0, 'not-json'), tool_delta(1, '{"ok": true}')]} + if mixed: + delta["content"] = "text" + responses = await collect([delta]) + final = responses[-1] + calls = [part.function_call for part in final.content.parts if part.function_call] + assert [call.name for call in calls] == ["tool_0", "tool_1"] + assert const.TOOL_CALL_ARGUMENT_ERRORS in calls[0].args + assert calls[1].args == {"ok": True} + assert final.error_code is None + + +@pytest.mark.parametrize("mixed_index", [0, 1, 2]) +@pytest.mark.asyncio +async def test_one_mixed_fragment_cannot_silently_change_arguments(mixed_index): + fragments = ['{"count":1', '2', '}'] + deltas = [{"tool_calls": [tool_delta(0, fragment, first=index == 0)]} for index, fragment in enumerate(fragments)] + deltas[mixed_index]["content"] = "text" + responses = await collect(deltas, {"tool_0"}) + calls = [part.function_call for part in responses[-1].content.parts if part.function_call] + assert len(calls) == 1 + assert calls[0].id == "call_0" + assert calls[0].args == {"count": 12} + streamed = [ + part.function_call.args[const.TOOL_STREAMING_ARGS] for response in responses[:-1] if response.content + for part in response.content.parts if part.function_call + ] + assert streamed == fragments + + +@pytest.mark.asyncio +async def test_mixed_chunk_text_precedes_streaming_tool_event(): + responses = await collect([{ + "content": "Recording now.", + "tool_calls": [tool_delta(0, '{"count":12}')] + }], {"tool_0"}) + assert responses[0].partial + assert responses[0].content.parts[0].text == "Recording now." + assert responses[1].partial + assert responses[1].content.parts[0].function_call.name == "tool_0" + + +@pytest.mark.asyncio +async def test_mixed_chunks_keep_text_through_anthropic_proxy(): + pytest.importorskip("claude_agent_sdk") + from trpc_agent_sdk.server.agents.claude._proxy import AnthropicMessagesRequest, AnthropicProxyApp + + responses = await collect([ + { + "content": "Recording ", + "tool_calls": [tool_delta(0, '{"count":')] + }, + { + "content": "now.", + "tool_calls": [tool_delta(0, '12}', first=False)] + }, + ], {"tool_0"}) + + async def generate(request, stream=True): + for response in responses: + yield response + + model = SimpleNamespace(generate_async=generate) + request = AnthropicMessagesRequest(model="synthetic", messages=[{"role": "user", "content": "record"}], stream=True) + events = [event async for event in AnthropicProxyApp()._handle_streaming(LlmRequest(), request, model)] + payloads = [ + json.loads(line[6:]) for event in events for line in event.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ] + deltas = [payload["delta"] for payload in payloads if payload["type"] == "content_block_delta"] + assert "".join(delta["text"] for delta in deltas if delta["type"] == "text_delta") == "Recording now." + assert "".join(delta["partial_json"] for delta in deltas if delta["type"] == "input_json_delta") == '{"count":12}' + + starts = [payload["index"] for payload in payloads if payload["type"] == "content_block_start"] + stops = [payload["index"] for payload in payloads if payload["type"] == "content_block_stop"] + assert len(starts) == len(set(starts)) + assert sorted(starts) == sorted(stops) + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.asyncio +async def test_streaming_tool_events_keep_response_identity_and_chunk_metadata(mixed): + deltas = [ + { + "tool_calls": [tool_delta(0, '{"count":')] + }, + { + "tool_calls": [tool_delta(0, '12}', first=False)] + }, + ] + if mixed: + for delta in deltas: + delta["content"] = "text" + responses = await collect(deltas, {"tool_0"}) + tool_events = [ + response for response in responses[:-1] if any(part.function_call for part in response.content.parts) + ] + assert len(tool_events) == 2 + for event, delta in zip(tool_events, deltas): + assert event.response_id == responses[-1].response_id == "synthetic-response" + assert event.custom_metadata[const.TOOL_STREAMING] is True + assert event.custom_metadata[const.CHUNK]["choices"][0]["delta"] == delta + + +@pytest.mark.asyncio +async def test_fragmented_invalid_json_in_mixed_chunks_keeps_error_envelope(): + responses = await collect([ + { + "content": "Recording ", + "tool_calls": [tool_delta(0, 'not-')] + }, + { + "content": "now.", + "tool_calls": [tool_delta(0, 'json', first=False)] + }, + ], {"tool_0"}) + assert [bool(response.content.parts[0].text) for response in responses[:-1]] == [True, False, True, False] + call = next(part.function_call for part in responses[-1].content.parts if part.function_call) + assert call.name == "tool_0" + assert const.TOOL_CALL_ARGUMENT_ERRORS in call.args + assert responses[-1].error_code is None + + +@pytest.mark.asyncio +async def test_proxy_closes_open_blocks_when_interleaved_stream_fails(): + pytest.importorskip("claude_agent_sdk") + from trpc_agent_sdk.server.agents.claude._proxy import AnthropicMessagesRequest, AnthropicProxyApp + + responses = await collect([ + { + "content": "Recording ", + "tool_calls": [tool_delta(0, '{"count":')] + }, + { + "content": "now." + }, + ], {"tool_0"}) + + async def generate(request, stream=True): + for response in responses[:-1]: + yield response + raise RuntimeError("test stream interrupted") + + request = AnthropicMessagesRequest(model="synthetic", messages=[{"role": "user", "content": "record"}], stream=True) + model = SimpleNamespace(generate_async=generate) + events = [event async for event in AnthropicProxyApp()._handle_streaming(LlmRequest(), request, model)] + payloads = [ + json.loads(line[6:]) for event in events for line in event.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ] + open_blocks = set() + for payload in payloads: + if payload["type"] == "content_block_start": + assert payload["index"] not in open_blocks + open_blocks.add(payload["index"]) + elif payload["type"] == "content_block_stop": + assert payload["index"] in open_blocks + open_blocks.remove(payload["index"]) + elif payload["type"] in {"message_delta", "message_stop"}: + assert not open_blocks + assert not open_blocks + assert next(p for p in payloads if p["type"] == "message_delta")["delta"]["stop_reason"] == "error" diff --git a/trpc_agent_sdk/models/_openai_model.py b/trpc_agent_sdk/models/_openai_model.py index cbeaf6923..5d0351080 100644 --- a/trpc_agent_sdk/models/_openai_model.py +++ b/trpc_agent_sdk/models/_openai_model.py @@ -1118,7 +1118,7 @@ def _process_responses_usage(self, response_dict: dict) -> Optional[GenerateCont def _process_chunk_without_content( self, chunk_dict: dict, accumulated_tool_calls: list[dict] ) -> tuple[Optional[FinishReason], Optional[GenerateContentResponseUsageMetadata], dict[int, str]]: - """Process a chunk that doesn't contain content. + """Process tool calls and usage, including chunks that also contain text. Args: chunk_dict (`dict`): The chunk dictionary to process @@ -1149,7 +1149,7 @@ def _process_chunk_without_content( # Handle usage usage = self._process_usage(chunk_dict) - # Handle tool calls in chunks without content (this is where streaming tool calls happen) + # Handle tool calls independently of text content. tool_calls_data = delta.get(const.TOOL_CALLS) # Track delta arguments from this chunk @@ -2519,30 +2519,36 @@ async def _generate_stream( # We can log them but don't need to yield them as content continue - # Verify if the chunk contains valid text content - if not self._verify_text_content_in_delta_response(chunk_dict): - # Process chunk without content (this is where tool calls are streamed) - _, usage, delta_arguments = self._process_chunk_without_content(chunk_dict, accumulated_tool_calls) - - if usage: - last_usage = usage - - # If streaming tool call arguments is enabled, yield partial tool call events - # delta_arguments being non-empty means tool calls were processed - if streaming_tool_names and delta_arguments and accumulated_tool_calls: - # Yield streaming tool call event with delta arguments - # Only for tools in streaming_tool_names - streaming_event = self._create_streaming_tool_call_response(accumulated_tool_calls, - delta_arguments, - streaming_tool_names) - if streaming_event: - yield streaming_event + # Tool deltas and text can coexist in the same chunk. + _, usage, delta_arguments = self._process_chunk_without_content(chunk_dict, accumulated_tool_calls) + + if usage: + last_usage = usage + + # Emit text before tool events from the same chunk for streaming consumers. + streaming_event = None + # If streaming tool call arguments is enabled, prepare partial tool call events + # delta_arguments being non-empty means tool calls were processed + if streaming_tool_names and delta_arguments and accumulated_tool_calls: + # Yield streaming tool call event with delta arguments + # Only for tools in streaming_tool_names + streaming_event = self._create_streaming_tool_call_response(accumulated_tool_calls, + delta_arguments, + streaming_tool_names) + if streaming_event: + streaming_event.response_id = response_id + streaming_event.custom_metadata = {const.CHUNK: chunk_dict, const.TOOL_STREAMING: True} + if not self._verify_text_content_in_delta_response(chunk_dict): + if streaming_event: + yield streaming_event continue # Process chunk with valid content choices = chunk_dict.get(const.CHOICES, []) if not choices: + if streaming_event: + yield streaming_event continue # Skip if no choices available choice: dict[str, dict] = choices[0] delta = choice[const.DELTA] @@ -2558,6 +2564,8 @@ async def _generate_stream( partial_text = self._adapter.filter_streaming_text(reasoning_content, reasoning_filter_state) if not partial_text: + if streaming_event: + yield streaming_event continue # Reasoning content is always thinking content @@ -2589,6 +2597,8 @@ async def _generate_stream( content_filter_state = streaming_text_filter_state["content"] partial_text = self._adapter.filter_streaming_text(content, content_filter_state) if not partial_text: + if streaming_event: + yield streaming_event continue if not is_thinking: visible_content += partial_text @@ -2604,11 +2614,8 @@ async def _generate_stream( response_id=response_id, custom_metadata={const.CHUNK: chunk_dict}, ) - - # Handle usage - usage = self._process_usage(chunk_dict) - if usage: - last_usage = usage + if streaming_event: + yield streaming_event if tool_prompt and streaming_text_filter_state is not None: if self._adapter.should_filter_reasoning_text(): diff --git a/trpc_agent_sdk/server/agents/claude/_proxy.py b/trpc_agent_sdk/server/agents/claude/_proxy.py index a916ee019..6b7a1a643 100644 --- a/trpc_agent_sdk/server/agents/claude/_proxy.py +++ b/trpc_agent_sdk/server/agents/claude/_proxy.py @@ -705,6 +705,9 @@ async def _handle_streaming(self, llm_request: LlmRequest, original_request: Ant Yields: Server-sent events in Anthropic format """ + text_block_closed = True + text_block_index = 0 + streaming_tool_calls: Dict[str, Dict[str, Any]] = {} try: # Send message_start event message_id = f"msg_{uuid.uuid4().hex[:24]}" @@ -746,10 +749,10 @@ async def _handle_streaming(self, llm_request: LlmRequest, original_request: Ant # Track state text_block_closed = False + text_block_index = 0 last_usage = None # Track streaming tool calls: {tool_id: {"index": int, "name": str, "started": bool}} - streaming_tool_calls: Dict[str, Dict[str, Any]] = {} - next_tool_index = 0 # Next available tool block index (0 is text block) + next_tool_index = 0 # Shared index allocator for text and tool blocks # Call model with streaming response_generator = model.generate_async(llm_request, stream=True) @@ -778,8 +781,9 @@ def _is_streaming_tool_call(resp) -> bool: # Close text block if not already closed if not text_block_closed: text_block_closed = True + text_stop_event = {"type": "content_block_stop", "index": text_block_index} yield ("event: content_block_stop\n" - f'data: {json.dumps({"type": "content_block_stop", "index": 0})}\n\n') + f"data: {json.dumps(text_stop_event)}\n\n") # Check if we need to start a new tool block if tool_id not in streaming_tool_calls: @@ -822,11 +826,23 @@ def _is_streaming_tool_call(resp) -> bool: # Handle partial responses (streaming text only) elif llm_response.partial and llm_response.content: for part in llm_response.content.parts: - if part.text and not text_block_closed: + if part.text: + if text_block_closed: + # A model may emit more text after starting a tool call. + next_tool_index += 1 + text_block_index = next_tool_index + text_block_closed = False + text_start_event = { + "type": "content_block_start", + "index": text_block_index, + "content_block": {"type": "text", "text": ""}, + } + yield ("event: content_block_start\n" + f"data: {json.dumps(text_start_event)}\n\n") # Stream text deltas text_delta_event = { "type": "content_block_delta", - "index": 0, + "index": text_block_index, "delta": { "type": "text_delta", "text": part.text @@ -841,7 +857,7 @@ def _is_streaming_tool_call(resp) -> bool: if not text_block_closed: text_block_closed = True yield ("event: content_block_stop\n" - f'data: {json.dumps({"type": "content_block_stop", "index": 0})}\n\n') + f'data: {json.dumps({"type": "content_block_stop", "index": text_block_index})}\n\n') # Close any streaming tool blocks that were started for tool_id, tool_info in streaming_tool_calls.items(): @@ -938,6 +954,18 @@ def _is_streaming_tool_call(resp) -> bool: except Exception as ex: # pylint: disable=broad-except self.logger.error("Error in streaming: %s", str(ex)) + # Close started blocks before terminating the interrupted message. + open_indices = [] + if not text_block_closed: + open_indices.append(text_block_index) + text_block_closed = True + for tool_info in streaming_tool_calls.values(): + if tool_info.get("started") and not tool_info.get("closed"): + open_indices.append(tool_info["index"]) + tool_info["closed"] = True + for index in sorted(open_indices): + stop_event = {"type": "content_block_stop", "index": index} + yield f"event: content_block_stop\ndata: {json.dumps(stop_event)}\n\n" # Send error response to client error_data = { "type": "message_delta",