fix(observe): wrap returned async generators (#1259)

Co-authored-by: @kouk

Author: hassiebp

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

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 == "<async_generator>"
+    assert nested_observation.output == "Hello--, async --World!"
 
 
 @pytest.mark.asyncio