diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 84bb60ebd..0e68bc965 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -19,8 +19,8 @@ overload, ) -from typing_extensions import ParamSpec from opentelemetry.util._decorator import _AgnosticContextManager +from typing_extensions import ParamSpec from langfuse._client.environment_variables import ( LANGFUSE_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED, @@ -267,6 +267,15 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any: result = await func(*args, **kwargs) if capture_output is True: + if inspect.isgenerator(result): + is_return_type_generator = True + + return self._wrap_sync_generator_result( + langfuse_span_or_generation, + result, + transform_to_string, + ) + if inspect.isasyncgen(result): is_return_type_generator = True @@ -369,6 +378,15 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: transform_to_string, ) + if inspect.isasyncgen(result): + is_return_type_generator = True + + return self._wrap_async_generator_result( + langfuse_span_or_generation, + result, + transform_to_string, + ) + langfuse_span_or_generation.update(output=result) return result diff --git a/tests/test_decorators.py b/tests/test_decorators.py index a08362f7b..47bc3b015 100644 --- a/tests/test_decorators.py +++ b/tests/test_decorators.py @@ -798,7 +798,7 @@ async def main_async(**kwargs): assert main_observation.output == mock_output assert nested_observation.name == "async_generator_function" - assert nested_observation.output == "" + assert nested_observation.output == "Hello--, async --World!" @pytest.mark.asyncio