diff --git a/python/semantic_kernel/functions/kernel_function.py b/python/semantic_kernel/functions/kernel_function.py index fee32dd437c1..b8f9a3a5fdbe 100644 --- a/python/semantic_kernel/functions/kernel_function.py +++ b/python/semantic_kernel/functions/kernel_function.py @@ -5,6 +5,7 @@ import time from abc import abstractmethod from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from contextlib import aclosing, closing from copy import copy, deepcopy from inspect import isasyncgen, isgenerator from typing import TYPE_CHECKING, Any @@ -352,13 +353,15 @@ async def invoke_stream( function_results: list[Any] = [] if function_context.result is not None: if isasyncgen(function_context.result.value): - async for partial in function_context.result.value: - function_results.append(partial) - yield partial + async with aclosing(function_context.result.value) as async_stream: + async for partial in async_stream: + function_results.append(partial) + yield partial elif isgenerator(function_context.result.value): - for partial in function_context.result.value: - function_results.append(partial) - yield partial + with closing(function_context.result.value) as stream: + for partial in stream: + function_results.append(partial) + yield partial else: function_results.append(function_context.result.value) yield function_context.result diff --git a/python/semantic_kernel/kernel.py b/python/semantic_kernel/kernel.py index 780a3ce5e294..e9e597e70652 100644 --- a/python/semantic_kernel/kernel.py +++ b/python/semantic_kernel/kernel.py @@ -2,7 +2,7 @@ import logging from collections.abc import AsyncGenerator, AsyncIterable, Callable -from contextlib import AbstractAsyncContextManager +from contextlib import AbstractAsyncContextManager, aclosing from copy import copy, deepcopy from typing import TYPE_CHECKING, Any, Literal, TypeVar @@ -142,15 +142,16 @@ async def invoke_stream( function_result: list[list["StreamingContentMixin"] | Any] = [] - async for stream_message in function.invoke_stream(self, arguments): - if isinstance(stream_message, FunctionResult) and ( - exception := stream_message.metadata.get(METADATA_EXCEPTION_KEY, None) - ): - raise KernelInvokeException( - f"Error occurred while invoking function: '{function.fully_qualified_name}'" - ) from exception - function_result.append(stream_message) - yield stream_message + async with aclosing(function.invoke_stream(self, arguments)) as stream: + async for stream_message in stream: + if isinstance(stream_message, FunctionResult) and ( + exception := stream_message.metadata.get(METADATA_EXCEPTION_KEY, None) + ): + raise KernelInvokeException( + f"Error occurred while invoking function: '{function.fully_qualified_name}'" + ) from exception + function_result.append(stream_message) + yield stream_message if return_function_results: output_function_result: list["StreamingContentMixin"] = [] diff --git a/python/tests/unit/functions/test_kernel_function_from_method.py b/python/tests/unit/functions/test_kernel_function_from_method.py index 2051ea0dc623..03b796d4dfef 100644 --- a/python/tests/unit/functions/test_kernel_function_from_method.py +++ b/python/tests/unit/functions/test_kernel_function_from_method.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import logging +from asyncio import CancelledError from collections.abc import AsyncGenerator, Iterable from contextlib import nullcontext from dataclasses import dataclass @@ -26,6 +27,85 @@ kernel_function_module = import_module("semantic_kernel.functions.kernel_function") +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("via_kernel", [False, True]) +@pytest.mark.parametrize("termination", ["close", "cancel", "exhaust", "error"]) +async def test_invoke_stream_closes_native_generator( + kernel: Kernel, asynchronous: bool, via_kernel: bool, termination: str +): + closed = [] + failure = ValueError("native stream failed") + + @kernel_function + def sync_stream() -> Iterable[str]: + try: + yield "first" + if termination == "error": + raise failure + yield "second" + finally: + closed.append(True) + + @kernel_function + async def async_stream() -> AsyncGenerator[str, None]: + try: + yield "first" + if termination == "error": + raise failure + yield "second" + finally: + closed.append(True) + + # Retain the source so garbage collection cannot stand in for explicit cleanup. + method = async_stream if asynchronous else sync_stream + source = method() + function = KernelFunctionFromMethod(method=method, stream_method=lambda: source) + stream = kernel.invoke_stream(function) if via_kernel else function.invoke_stream(kernel) + try: + assert await anext(stream) == "first" + assert not closed + if termination == "close": + await stream.aclose() + elif termination == "cancel": + with pytest.raises(CancelledError): + await stream.athrow(CancelledError()) + elif termination == "error": + with pytest.raises(ValueError) as exc_info: + await anext(stream) + assert exc_info.value is failure + else: + assert [part async for part in stream] == ["second"] + assert closed == [True] + finally: + await stream.aclose() + if asynchronous: + await source.aclose() + else: + source.close() + + +@pytest.mark.parametrize("via_kernel", [False, True]) +async def test_invoke_stream_does_not_close_ordinary_result(kernel: Kernel, via_kernel: bool): + closed = [] + + class ResultList(list): + def close(self): + closed.append(True) + + @kernel_function + def native() -> str: + return "non-streaming result" + + result = ResultList(["ordinary result"]) + function = KernelFunctionFromMethod(method=native, stream_method=lambda: result) + stream = kernel.invoke_stream(function) if via_kernel else function.invoke_stream(kernel) + parts = [part async for part in stream] + assert len(parts) == 1 + assert parts[0].value is result + assert result == ["ordinary result"] + assert not closed + + class CustomType(KernelBaseModel): id: str name: str