fix tests

hassiebp · hassiebp · commit 678a420efbfa · 2025-09-28T12:32:42.000+02:00
diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py
@@ -11,6 +11,7 @@
     Dict,
     Generator,
     Iterable,
+    List,
     Optional,
     Tuple,
     TypeVar,
@@ -468,64 +469,11 @@ def _wrap_sync_generator_result(
         generator: Generator,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
     ) -> Any:
-        # Capture the current context while the span is still active
         preserved_context = contextvars.copy_context()
-        items: list[Any] = []
-
-        class ContextPreservedSyncGeneratorWrapper:
-            """Sync generator wrapper that ensures each iteration runs in preserved context."""
-
-            def __init__(
-                self,
-                generator: Generator,
-                context: contextvars.Context,
-                items: list[Any],
-                span: Union[
-                    LangfuseSpan,
-                    LangfuseGeneration,
-                    LangfuseAgent,
-                    LangfuseTool,
-                    LangfuseChain,
-                    LangfuseRetriever,
-                    LangfuseEvaluator,
-                    LangfuseEmbedding,
-                    LangfuseGuardrail,
-                ],
-                transform_fn: Optional[Callable[[Iterable], str]],
-            ) -> None:
-                self.generator = generator
-                self.context = context
-                self.items = items
-                self.span = span
-                self.transform_fn = transform_fn
-
-            def __iter__(self) -> "ContextPreservedSyncGeneratorWrapper":
-                return self
-
-            def __next__(self) -> Any:
-                try:
-                    # Run the generator's __next__ in the preserved context
-                    item = self.context.run(next, self.generator)
-                    self.items.append(item)
-                    return item
-
-                except StopIteration:
-                    # Handle output and span cleanup when generator is exhausted
-                    output: Any = self.items
-
-                    if self.transform_fn is not None:
-                        output = self.transform_fn(self.items)
-                    elif all(isinstance(item, str) for item in self.items):
-                        output = "".join(self.items)
-
-                    self.span.update(output=output)
-                    self.span.end()
-                    raise  # Re-raise StopIteration
-
-        return ContextPreservedSyncGeneratorWrapper(
+
+        return _ContextPreservedSyncGeneratorWrapper(
             generator,
             preserved_context,
-            items,
             langfuse_span_or_generation,
             transform_to_string,
         )
@@ -546,75 +494,11 @@ def _wrap_async_generator_result(
         generator: AsyncGenerator,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
     ) -> Any:
-        import asyncio
-
-        # Capture the current context while the span is still active
         preserved_context = contextvars.copy_context()
-        items: list[Any] = []
-
-        class ContextPreservedAsyncGeneratorWrapper:
-            """Async generator wrapper that ensures each iteration runs in preserved context."""
-
-            def __init__(
-                self,
-                generator: AsyncGenerator,
-                context: contextvars.Context,
-                items: list[Any],
-                span: Union[
-                    LangfuseSpan,
-                    LangfuseGeneration,
-                    LangfuseAgent,
-                    LangfuseTool,
-                    LangfuseChain,
-                    LangfuseRetriever,
-                    LangfuseEvaluator,
-                    LangfuseEmbedding,
-                    LangfuseGuardrail,
-                ],
-                transform_fn: Optional[Callable[[Iterable], str]],
-            ) -> None:
-                self.generator = generator
-                self.context = context
-                self.items = items
-                self.span = span
-                self.transform_fn = transform_fn
-
-            def __aiter__(self) -> "ContextPreservedAsyncGeneratorWrapper":
-                return self
-
-            async def __anext__(self) -> Any:
-                try:
-                    # Run the generator's __anext__ in the preserved context
-                    try:
-                        # Python 3.10+ approach with context parameter
-                        item = await asyncio.create_task(
-                            self.generator.__anext__(),  # type: ignore
-                            context=self.context,
-                        )  # type: ignore
-                    except TypeError:
-                        # Python < 3.10 fallback - context parameter not supported
-                        item = await self.generator.__anext__()
-
-                    self.items.append(item)
-                    return item
-
-                except StopAsyncIteration:
-                    # Handle output and span cleanup when generator is exhausted
-                    output: Any = self.items
-
-                    if self.transform_fn is not None:
-                        output = self.transform_fn(self.items)
-                    elif all(isinstance(item, str) for item in self.items):
-                        output = "".join(self.items)
-
-                    self.span.update(output=output)
-                    self.span.end()
-                    raise  # Re-raise StopAsyncIteration
-
-        return ContextPreservedAsyncGeneratorWrapper(
+
+        return _ContextPreservedAsyncGeneratorWrapper(
             generator,
             preserved_context,
-            items,
             langfuse_span_or_generation,
             transform_to_string,
         )
@@ -623,3 +507,125 @@ async def __anext__(self) -> Any:
 _decorator = LangfuseDecorator()
 
 observe = _decorator.observe
+
+
+class _ContextPreservedSyncGeneratorWrapper:
+    """Sync generator wrapper that ensures each iteration runs in preserved context."""
+
+    def __init__(
+        self,
+        generator: Generator,
+        context: contextvars.Context,
+        span: Union[
+            LangfuseSpan,
+            LangfuseGeneration,
+            LangfuseAgent,
+            LangfuseTool,
+            LangfuseChain,
+            LangfuseRetriever,
+            LangfuseEvaluator,
+            LangfuseEmbedding,
+            LangfuseGuardrail,
+        ],
+        transform_fn: Optional[Callable[[Iterable], str]],
+    ) -> None:
+        self.generator = generator
+        self.context = context
+        self.items: List[Any] = []
+        self.span = span
+        self.transform_fn = transform_fn
+
+    def __iter__(self) -> "_ContextPreservedSyncGeneratorWrapper":
+        return self
+
+    def __next__(self) -> Any:
+        try:
+            # Run the generator's __next__ in the preserved context
+            item = self.context.run(next, self.generator)
+            self.items.append(item)
+
+            return item
+
+        except StopIteration:
+            # Handle output and span cleanup when generator is exhausted
+            output: Any = self.items
+
+            if self.transform_fn is not None:
+                output = self.transform_fn(self.items)
+
+            elif all(isinstance(item, str) for item in self.items):
+                output = "".join(self.items)
+
+            self.span.update(output=output).end()
+
+            raise  # Re-raise StopIteration
+
+        except Exception as e:
+            self.span.update(level="ERROR", status_message=str(e)).end()
+
+            raise e
+
+
+class _ContextPreservedAsyncGeneratorWrapper:
+    """Async generator wrapper that ensures each iteration runs in preserved context."""
+
+    def __init__(
+        self,
+        generator: AsyncGenerator,
+        context: contextvars.Context,
+        span: Union[
+            LangfuseSpan,
+            LangfuseGeneration,
+            LangfuseAgent,
+            LangfuseTool,
+            LangfuseChain,
+            LangfuseRetriever,
+            LangfuseEvaluator,
+            LangfuseEmbedding,
+            LangfuseGuardrail,
+        ],
+        transform_fn: Optional[Callable[[Iterable], str]],
+    ) -> None:
+        self.generator = generator
+        self.context = context
+        self.items: List[Any] = []
+        self.span = span
+        self.transform_fn = transform_fn
+
+    def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper":
+        return self
+
+    async def __anext__(self) -> Any:
+        try:
+            # Run the generator's __anext__ in the preserved context
+            try:
+                # Python 3.10+ approach with context parameter
+                item = await asyncio.create_task(
+                    self.generator.__anext__(),  # type: ignore
+                    context=self.context,
+                )  # type: ignore
+            except TypeError:
+                # Python < 3.10 fallback - context parameter not supported
+                item = await self.generator.__anext__()
+
+            self.items.append(item)
+
+            return item
+
+        except StopAsyncIteration:
+            # Handle output and span cleanup when generator is exhausted
+            output: Any = self.items
+
+            if self.transform_fn is not None:
+                output = self.transform_fn(self.items)
+
+            elif all(isinstance(item, str) for item in self.items):
+                output = "".join(self.items)
+
+            self.span.update(output=output).end()
+
+            raise  # Re-raise StopAsyncIteration
+        except Exception as e:
+            self.span.update(level="ERROR", status_message=str(e)).end()
+
+            raise e
diff --git a/tests/test_decorators.py b/tests/test_decorators.py
@@ -8,6 +8,7 @@
 import pytest
 from langchain.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
+from opentelemetry import trace
 
 from langfuse import Langfuse, get_client, observe
 from langfuse._client.environment_variables import LANGFUSE_PUBLIC_KEY
@@ -1698,8 +1699,6 @@ def test_sync_generator_context_preservation():
 
     @observe(name="sync_generator")
     def create_generator():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_info["generator_span_id"] = trace.format_span_id(
             current_span.get_span_context().span_id
@@ -1710,8 +1709,6 @@ def create_generator():
 
     @observe(name="root")
     def root_function():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_info["root_span_id"] = trace.format_span_id(
             current_span.get_span_context().span_id
@@ -1764,8 +1761,6 @@ async def test_async_generator_context_preservation():
 
     @observe(name="async_generator")
     async def create_async_generator():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_info["generator_span_id"] = trace.format_span_id(
             current_span.get_span_context().span_id
@@ -1777,8 +1772,6 @@ async def create_async_generator():
 
     @observe(name="root")
     async def root_function():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_info["root_span_id"] = trace.format_span_id(
             current_span.get_span_context().span_id
@@ -1833,8 +1826,6 @@ async def test_async_generator_context_preservation_with_trace_hierarchy():
 
     @observe(name="child_stream")
     async def child_generator():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_context = current_span.get_span_context()
         span_info["child_span_id"] = trace.format_span_id(span_context.span_id)
@@ -1846,8 +1837,6 @@ async def child_generator():
 
     @observe(name="parent_root")
     async def parent_function():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         span_context = current_span.get_span_context()
         span_info["parent_span_id"] = trace.format_span_id(span_context.span_id)
@@ -1896,8 +1885,6 @@ async def test_async_generator_exception_handling_with_context():
 
     @observe(name="failing_generator")
     async def failing_generator():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         # Verify we have valid context even when exception occurs
         assert (
@@ -1946,8 +1933,6 @@ def test_sync_generator_empty_context_preservation():
 
     @observe(name="empty_generator")
     def empty_generator():
-        from opentelemetry import trace
-
         current_span = trace.get_current_span()
         # Should have valid context even for empty generator
         assert (
@@ -1977,4 +1962,4 @@ def root_function():
     empty_obs = next(
         obs for obs in trace_data.observations if obs.name == "empty_generator"
     )
-    assert empty_obs.output == ""
+    assert empty_obs.output is None
