diff --git a/src/pipecat/services/google/llm_openai.py b/src/pipecat/services/google/llm_openai.py index c781fe0d3..5d396b583 100644 --- a/src/pipecat/services/google/llm_openai.py +++ b/src/pipecat/services/google/llm_openai.py @@ -91,52 +91,55 @@ class GoogleLLMOpenAIBetaService(OpenAILLMService): ChatCompletionChunk ] = await self._stream_chat_completions_specific_context(context) - async for chunk in chunk_stream: - if chunk.usage: - tokens = LLMTokenUsage( - prompt_tokens=chunk.usage.prompt_tokens or 0, - completion_tokens=chunk.usage.completion_tokens or 0, - total_tokens=chunk.usage.total_tokens or 0, - ) - await self.start_llm_usage_metrics(tokens) + # Use context manager to ensure stream is closed on cancellation/exception. + # Without this, CancelledError during iteration leaves the underlying socket open. + async with chunk_stream: + async for chunk in chunk_stream: + if chunk.usage: + tokens = LLMTokenUsage( + prompt_tokens=chunk.usage.prompt_tokens or 0, + completion_tokens=chunk.usage.completion_tokens or 0, + total_tokens=chunk.usage.total_tokens or 0, + ) + await self.start_llm_usage_metrics(tokens) - if chunk.choices is None or len(chunk.choices) == 0: - continue + if chunk.choices is None or len(chunk.choices) == 0: + continue - await self.stop_ttfb_metrics() + await self.stop_ttfb_metrics() - if not chunk.choices[0].delta: - continue + if not chunk.choices[0].delta: + continue - if chunk.choices[0].delta.tool_calls: - # We're streaming the LLM response to enable the fastest response times. - # For text, we just yield each chunk as we receive it and count on consumers - # to do whatever coalescing they need (eg. to pass full sentences to TTS) - # - # If the LLM is a function call, we'll do some coalescing here. - # If the response contains a function name, we'll yield a frame to tell consumers - # that they can start preparing to call the function with that name. - # We accumulate all the arguments for the rest of the streamed response, then when - # the response is done, we package up all the arguments and the function name and - # yield a frame containing the function name and the arguments. - logger.debug(f"Tool call: {chunk.choices[0].delta.tool_calls}") - tool_call = chunk.choices[0].delta.tool_calls[0] - if tool_call.index != func_idx: - functions_list.append(function_name) - arguments_list.append(arguments) - tool_id_list.append(tool_call_id) - function_name = "" - arguments = "" - tool_call_id = "" - func_idx += 1 - if tool_call.function and tool_call.function.name: - function_name += tool_call.function.name - tool_call_id = tool_call.id - if tool_call.function and tool_call.function.arguments: - # Keep iterating through the response to collect all the argument fragments - arguments += tool_call.function.arguments - elif chunk.choices[0].delta.content: - await self.push_frame(LLMTextFrame(chunk.choices[0].delta.content)) + if chunk.choices[0].delta.tool_calls: + # We're streaming the LLM response to enable the fastest response times. + # For text, we just yield each chunk as we receive it and count on consumers + # to do whatever coalescing they need (eg. to pass full sentences to TTS) + # + # If the LLM is a function call, we'll do some coalescing here. + # If the response contains a function name, we'll yield a frame to tell consumers + # that they can start preparing to call the function with that name. + # We accumulate all the arguments for the rest of the streamed response, then when + # the response is done, we package up all the arguments and the function name and + # yield a frame containing the function name and the arguments. + logger.debug(f"Tool call: {chunk.choices[0].delta.tool_calls}") + tool_call = chunk.choices[0].delta.tool_calls[0] + if tool_call.index != func_idx: + functions_list.append(function_name) + arguments_list.append(arguments) + tool_id_list.append(tool_call_id) + function_name = "" + arguments = "" + tool_call_id = "" + func_idx += 1 + if tool_call.function and tool_call.function.name: + function_name += tool_call.function.name + tool_call_id = tool_call.id + if tool_call.function and tool_call.function.arguments: + # Keep iterating through the response to collect all the argument fragments + arguments += tool_call.function.arguments + elif chunk.choices[0].delta.content: + await self.push_frame(LLMTextFrame(chunk.choices[0].delta.content)) # if we got a function name and arguments, check to see if it's a function with # a registered handler. If so, run the registered callback, save the result to diff --git a/src/pipecat/services/sambanova/llm.py b/src/pipecat/services/sambanova/llm.py index d50978d72..047ce0e6c 100644 --- a/src/pipecat/services/sambanova/llm.py +++ b/src/pipecat/services/sambanova/llm.py @@ -131,59 +131,62 @@ class SambaNovaLLMService(OpenAILLMService): # type: ignore else self._stream_chat_completions_universal_context(context) ) - async for chunk in chunk_stream: - if chunk.usage: - tokens = LLMTokenUsage( - prompt_tokens=chunk.usage.prompt_tokens, - completion_tokens=chunk.usage.completion_tokens, - total_tokens=chunk.usage.total_tokens, - ) - await self.start_llm_usage_metrics(tokens) + # Use context manager to ensure stream is closed on cancellation/exception. + # Without this, CancelledError during iteration leaves the underlying socket open. + async with chunk_stream: + async for chunk in chunk_stream: + if chunk.usage: + tokens = LLMTokenUsage( + prompt_tokens=chunk.usage.prompt_tokens, + completion_tokens=chunk.usage.completion_tokens, + total_tokens=chunk.usage.total_tokens, + ) + await self.start_llm_usage_metrics(tokens) - if chunk.choices is None or len(chunk.choices) == 0: - continue + if chunk.choices is None or len(chunk.choices) == 0: + continue - await self.stop_ttfb_metrics() + await self.stop_ttfb_metrics() - if not chunk.choices[0].delta: - continue + if not chunk.choices[0].delta: + continue - if chunk.choices[0].delta.tool_calls: - # We're streaming the LLM response to enable the fastest response times. - # For text, we just yield each chunk as we receive it and count on consumers - # to do whatever coalescing they need (eg. to pass full sentences to TTS) - # - # If the LLM is a function call, we'll do some coalescing here. - # If the response contains a function name, we'll yield a frame to tell consumers - # that they can start preparing to call the function with that name. - # We accumulate all the arguments for the rest of the streamed response, then when - # the response is done, we package up all the arguments and the function name and - # yield a frame containing the function name and the arguments. + if chunk.choices[0].delta.tool_calls: + # We're streaming the LLM response to enable the fastest response times. + # For text, we just yield each chunk as we receive it and count on consumers + # to do whatever coalescing they need (eg. to pass full sentences to TTS) + # + # If the LLM is a function call, we'll do some coalescing here. + # If the response contains a function name, we'll yield a frame to tell consumers + # that they can start preparing to call the function with that name. + # We accumulate all the arguments for the rest of the streamed response, then when + # the response is done, we package up all the arguments and the function name and + # yield a frame containing the function name and the arguments. - tool_call = chunk.choices[0].delta.tool_calls[0] - if tool_call.index != func_idx: - functions_list.append(function_name) - arguments_list.append(arguments) - tool_id_list.append(tool_call_id) - function_name = "" - arguments = "" - tool_call_id = "" - func_idx += 1 - if tool_call.function and tool_call.function.name: - function_name += tool_call.function.name - tool_call_id = tool_call.id # type: ignore - if tool_call.function and tool_call.function.arguments: - # Keep iterating through the response to collect all the argument fragments - arguments += tool_call.function.arguments - elif chunk.choices[0].delta.content: - await self.push_frame(LLMTextFrame(chunk.choices[0].delta.content)) + tool_call = chunk.choices[0].delta.tool_calls[0] + if tool_call.index != func_idx: + functions_list.append(function_name) + arguments_list.append(arguments) + tool_id_list.append(tool_call_id) + function_name = "" + arguments = "" + tool_call_id = "" + func_idx += 1 + if tool_call.function and tool_call.function.name: + function_name += tool_call.function.name + tool_call_id = tool_call.id # type: ignore + if tool_call.function and tool_call.function.arguments: + # Keep iterating through the response to collect all the argument fragments + arguments += tool_call.function.arguments + elif chunk.choices[0].delta.content: + await self.push_frame(LLMTextFrame(chunk.choices[0].delta.content)) - # When gpt-4o-audio / gpt-4o-mini-audio is used for llm or stt+llm - # we need to get LLMTextFrame for the transcript - elif hasattr(chunk.choices[0].delta, "audio") and chunk.choices[0].delta.audio.get( - "transcript" - ): - await self.push_frame(LLMTextFrame(chunk.choices[0].delta.audio["transcript"])) + # When gpt-4o-audio / gpt-4o-mini-audio is used for llm or stt+llm + # we need to get LLMTextFrame for the transcript + elif hasattr(chunk.choices[0].delta, "audio") and chunk.choices[0].delta.audio.get( + "transcript" + ): + await self.push_frame(LLMTextFrame(chunk.choices[0].delta.audio["transcript"])) # if we got a function name and arguments, check to see if it's a function with # a registered handler. If so, run the registered callback, save the result to diff --git a/tests/test_google_llm_openai.py b/tests/test_google_llm_openai.py new file mode 100644 index 000000000..2940bf7d7 --- /dev/null +++ b/tests/test_google_llm_openai.py @@ -0,0 +1,81 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Unit tests for Google LLM OpenAI Beta service.""" + +import asyncio +import warnings +from unittest.mock import AsyncMock, patch + +import pytest + +from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext + +try: + from pipecat.services.google.llm_openai import GoogleLLMOpenAIBetaService + + google_available = True +except Exception: + google_available = False + + +@pytest.mark.asyncio +@pytest.mark.skipif(not google_available, reason="Google dependencies not installed") +async def test_google_llm_openai_stream_closed_on_cancellation(): + """Test that the stream is closed when CancelledError occurs during iteration. + + This prevents socket leaks when the pipeline is interrupted (e.g., user interruption). + See issue #3639. + """ + with patch.object(GoogleLLMOpenAIBetaService, "create_client"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + service = GoogleLLMOpenAIBetaService(api_key="test-key", model="test-model") + service._client = AsyncMock() + + stream_closed = False + + class MockAsyncStream: + """Mock AsyncStream that tracks close() calls and raises CancelledError.""" + + def __init__(self): + self.iteration_count = 0 + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + nonlocal stream_closed + stream_closed = True + return False + + def __aiter__(self): + return self + + async def __anext__(self): + self.iteration_count += 1 + if self.iteration_count > 1: + raise asyncio.CancelledError() + mock_chunk = AsyncMock() + mock_chunk.usage = None + mock_chunk.choices = [] + return mock_chunk + + mock_stream = MockAsyncStream() + + service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) + service.start_ttfb_metrics = AsyncMock() + service.stop_ttfb_metrics = AsyncMock() + service.start_llm_usage_metrics = AsyncMock() + + context = OpenAILLMContext( + messages=[{"role": "user", "content": "Hello"}], + ) + + with pytest.raises(asyncio.CancelledError): + await service._process_context(context) + + assert stream_closed, "Stream should be closed even when CancelledError occurs" diff --git a/tests/test_sambanova_llm.py b/tests/test_sambanova_llm.py new file mode 100644 index 000000000..6632951fc --- /dev/null +++ b/tests/test_sambanova_llm.py @@ -0,0 +1,72 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Unit tests for SambaNova LLM service.""" + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +from pipecat.processors.aggregators.llm_context import LLMContext +from pipecat.services.sambanova.llm import SambaNovaLLMService + + +@pytest.mark.asyncio +async def test_sambanova_llm_stream_closed_on_cancellation(): + """Test that the stream is closed when CancelledError occurs during iteration. + + This prevents socket leaks when the pipeline is interrupted (e.g., user interruption). + See issue #3639. + """ + with patch.object(SambaNovaLLMService, "create_client"): + service = SambaNovaLLMService(api_key="test-key", model="test-model") + service._client = AsyncMock() + + stream_closed = False + + class MockAsyncStream: + """Mock AsyncStream that tracks close() calls and raises CancelledError.""" + + def __init__(self): + self.iteration_count = 0 + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + nonlocal stream_closed + stream_closed = True + return False + + def __aiter__(self): + return self + + async def __anext__(self): + self.iteration_count += 1 + if self.iteration_count > 1: + raise asyncio.CancelledError() + mock_chunk = AsyncMock() + mock_chunk.usage = None + mock_chunk.choices = [] + return mock_chunk + + mock_stream = MockAsyncStream() + + service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) + service._stream_chat_completions_universal_context = AsyncMock(return_value=mock_stream) + service.start_ttfb_metrics = AsyncMock() + service.stop_ttfb_metrics = AsyncMock() + service.start_llm_usage_metrics = AsyncMock() + + context = LLMContext( + messages=[{"role": "user", "content": "Hello"}], + ) + + with pytest.raises(asyncio.CancelledError): + await service._process_context(context) + + assert stream_closed, "Stream should be closed even when CancelledError occurs"