judgeval 0.16.9__py3-none-any.whl → 0.22.2__py3-none-any.whl

judgeval/tracer/__init__.py

@@ -17,9 +17,12 @@ from typing import (
     overload,
     Literal,
     TypedDict,
-    Iterator,
-    AsyncIterator,
+    Generator,
+    AsyncGenerator,
+    Iterable,
 )
+import contextvars
+import asyncio
 from functools import partial
 from warnings import warn
 
@@ -71,6 +74,8 @@ from judgeval.tracer.processors import (
     NoOpJudgmentSpanProcessor,
 )
 from judgeval.tracer.utils import set_span_attribute, TraceScorerConfig
+from judgeval.utils.project import _resolve_project_id
+from opentelemetry.trace import use_span
 
 C = TypeVar("C", bound=Callable)
 Cls = TypeVar("Cls", bound=Type)
@@ -101,11 +106,12 @@ class Tracer(metaclass=SingletonMeta):
         "judgment_processor",
         "tracer",
         "agent_context",
+        "customer_id",
         "_initialized",
     )
 
-    api_key: str
-    organization_id: str
+    api_key: str | None
+    organization_id: str | None
     project_name: str
     enable_monitoring: bool
     enable_evaluation: bool
@@ -114,6 +120,7 @@ class Tracer(metaclass=SingletonMeta):
     judgment_processor: JudgmentSpanProcessor
     tracer: ABCTracer
     agent_context: ContextVar[Optional[AgentContext]]
+    customer_id: ContextVar[Optional[str]]
     _initialized: bool
 
     def __init__(
@@ -121,8 +128,8 @@ class Tracer(metaclass=SingletonMeta):
         /,
         *,
         project_name: str,
-        api_key: Optional[str] = None,
-        organization_id: Optional[str] = None,
+        api_key: str | None = None,
+        organization_id: str | None = None,
         enable_monitoring: bool = JUDGMENT_ENABLE_MONITORING.lower() == "true",
         enable_evaluation: bool = JUDGMENT_ENABLE_EVALUATIONS.lower() == "true",
         resource_attributes: Optional[Dict[str, Any]] = None,
@@ -131,6 +138,7 @@ class Tracer(metaclass=SingletonMeta):
         if not hasattr(self, "_initialized"):
             self._initialized = False
             self.agent_context = ContextVar("current_agent_context", default=None)
+            self.customer_id = ContextVar("current_customer_id", default=None)
 
             self.project_name = project_name
             self.api_key = expect_api_key(api_key or JUDGMENT_API_KEY)
@@ -141,10 +149,14 @@ class Tracer(metaclass=SingletonMeta):
             self.enable_evaluation = enable_evaluation
             self.resource_attributes = resource_attributes
 
-            self.api_client = JudgmentSyncClient(
-                api_key=self.api_key,
-                organization_id=self.organization_id,
-            )
+            if self.api_key and self.organization_id:
+                self.api_client = JudgmentSyncClient(
+                    api_key=self.api_key, organization_id=self.organization_id
+                )
+            else:
+                judgeval_logger.error(
+                    "API Key or Organization ID is not set. You must set them in the environment variables to use the tracer."
+                )
 
             if initialize:
                 self.initialize()
@@ -155,10 +167,10 @@ class Tracer(metaclass=SingletonMeta):
 
         self.judgment_processor = NoOpJudgmentSpanProcessor()
         if self.enable_monitoring:
-            project_id = Tracer._resolve_project_id(
+            project_id = _resolve_project_id(
                 self.project_name, self.api_key, self.organization_id
             )
-            if project_id:
+            if self.api_key and self.organization_id and project_id:
                 self.judgment_processor = self.get_processor(
                     tracer=self,
                     project_name=self.project_name,
@@ -173,9 +185,10 @@ class Tracer(metaclass=SingletonMeta):
                 provider.add_span_processor(self.judgment_processor)
                 set_tracer_provider(provider)
             else:
-                judgeval_logger.error(
-                    f"Failed to resolve or autocreate project {self.project_name}, please create it first at https://app.judgmentlabs.ai/org/{self.organization_id}/projects. Skipping Judgment export."
-                )
+                if self.api_key and self.organization_id:
+                    judgeval_logger.error(
+                        f"Failed to resolve or autocreate project {self.project_name}, please create it first at https://app.judgmentlabs.ai/org/{self.organization_id}/projects. Skipping Judgment export."
+                    )
 
         self.tracer = get_tracer_provider().get_tracer(
             JUDGEVAL_TRACER_INSTRUMENTING_MODULE_NAME,
@@ -194,10 +207,19 @@ class Tracer(metaclass=SingletonMeta):
     ):
         from judgeval.tracer.exporters import JudgmentSpanExporter
 
+        api_key = api_key or JUDGMENT_API_KEY
+        organization_id = organization_id or JUDGMENT_ORG_ID
+
+        if not api_key or not organization_id:
+            judgeval_logger.error(
+                "API Key or Organization ID is not set. You must set them in the environment variables to use the tracer."
+            )
+            return None
+
         return JudgmentSpanExporter(
             endpoint=url_for("/otel/v1/traces"),
-            api_key=api_key or JUDGMENT_API_KEY,
-            organization_id=organization_id or JUDGMENT_ORG_ID,
+            api_key=api_key,
+            organization_id=organization_id,
             project_id=project_id,
         )
 
@@ -213,31 +235,24 @@ class Tracer(metaclass=SingletonMeta):
         resource_attributes: Optional[Dict[str, Any]] = None,
     ) -> JudgmentSpanProcessor:
         """Create a JudgmentSpanProcessor using the correct constructor."""
+        api_key = api_key or JUDGMENT_API_KEY
+        organization_id = organization_id or JUDGMENT_ORG_ID
+        if not api_key or not organization_id:
+            judgeval_logger.error(
+                "API Key or Organization ID is not set. You must set them in the environment variables to use the tracer."
+            )
+            return NoOpJudgmentSpanProcessor()
         return JudgmentSpanProcessor(
             tracer,
             project_name,
             project_id,
-            api_key or JUDGMENT_API_KEY,
-            organization_id or JUDGMENT_ORG_ID,
+            api_key,
+            organization_id,
             max_queue_size=max_queue_size,
             export_timeout_millis=export_timeout_millis,
             resource_attributes=resource_attributes,
         )
 
-    @dont_throw
-    @functools.lru_cache(maxsize=64)
-    @staticmethod
-    def _resolve_project_id(
-        project_name: str, api_key: str, organization_id: str
-    ) -> str:
-        """Resolve project_id from project_name using the API."""
-        client = JudgmentSyncClient(
-            api_key=api_key,
-            organization_id=organization_id,
-        )
-        response = client.projects_resolve({"project_name": project_name})
-        return response["project_id"]
-
     def get_current_span(self):
         return get_current_span()
 
@@ -247,17 +262,51 @@ class Tracer(metaclass=SingletonMeta):
     def get_current_agent_context(self):
         return self.agent_context
 
+    def get_current_customer_context(self):
+        return self.customer_id
+
     def get_span_processor(self) -> JudgmentSpanProcessor:
         """Get the internal span processor of this tracer instance."""
         return self.judgment_processor
 
+    @dont_throw
     def set_customer_id(self, customer_id: str) -> None:
+        if not customer_id:
+            judgeval_logger.warning("Customer ID is empty, skipping.")
+            return
+
         span = self.get_current_span()
+
+        if not span or not span.is_recording():
+            judgeval_logger.warning(
+                "No active span found. Customer ID will not be set."
+            )
+            return
+
+        if self.get_current_customer_context().get():
+            judgeval_logger.warning("Customer ID is already set, skipping.")
+            return
+
         if span and span.is_recording():
             set_span_attribute(span, AttributeKeys.JUDGMENT_CUSTOMER_ID, customer_id)
+            self.get_current_customer_context().set(customer_id)
+
+            self.get_span_processor().set_internal_attribute(
+                span_context=span.get_span_context(),
+                key=InternalAttributeKeys.IS_CUSTOMER_CONTEXT_OWNER,
+                value=True,
+            )
+
+    def _maybe_clear_customer_context(self, span: Span) -> None:
+        if self.get_span_processor().get_internal_attribute(
+            span_context=span.get_span_context(),
+            key=InternalAttributeKeys.IS_CUSTOMER_CONTEXT_OWNER,
+            default=False,
+        ):
+            self.get_current_customer_context().set(None)
 
     @dont_throw
-    def add_agent_attributes_to_span(self, span):
+    def _add_agent_attributes_to_span(self, span):
         """Add agent ID, class name, and instance name to span if they exist in context"""
         current_agent_context = self.agent_context.get()
         if not current_agent_context:
@@ -289,7 +338,7 @@ class Tracer(metaclass=SingletonMeta):
         current_agent_context["is_agent_entry_point"] = False
 
     @dont_throw
-    def record_instance_state(self, record_point: Literal["before", "after"], span):
+    def _record_instance_state(self, record_point: Literal["before", "after"], span):
         current_agent_context = self.agent_context.get()
 
         if current_agent_context and current_agent_context.get("track_state"):
@@ -318,6 +367,17 @@ class Tracer(metaclass=SingletonMeta):
                 safe_serialize(attributes),
             )
 
+    @dont_throw
+    def _add_customer_id_to_span(self, span):
+        customer_id = self.get_current_customer_context().get()
+        if customer_id:
+            set_span_attribute(span, AttributeKeys.JUDGMENT_CUSTOMER_ID, customer_id)
+
+    @dont_throw
+    def _inject_judgment_context(self, span):
+        self._add_agent_attributes_to_span(span)
+        self._add_customer_id_to_span(span)
+
     def _set_pending_trace_eval(
         self,
         span: Span,
@@ -381,91 +441,37 @@ class Tracer(metaclass=SingletonMeta):
 
     def _create_traced_sync_generator(
         self,
-        generator: Iterator[Any],
+        generator: Generator,
         main_span: Span,
-        base_name: str,
-        attributes: Optional[Dict[str, Any]],
+        disable_generator_yield_span: bool = False,
     ):
         """Create a traced synchronous generator that wraps each yield in a span."""
-        try:
-            while True:
-                yield_span_name = f"{base_name}_yield"
-                yield_attributes = {
-                    AttributeKeys.JUDGMENT_SPAN_KIND: "generator_yield",
-                    **(attributes or {}),
-                }
-
-                with sync_span_context(
-                    self, yield_span_name, yield_attributes, disable_partial_emit=True
-                ) as yield_span:
-                    self.add_agent_attributes_to_span(yield_span)
-
-                    try:
-                        value = next(generator)
-                    except StopIteration:
-                        # Mark span as cancelled so it won't be exported
-                        self.judgment_processor.set_internal_attribute(
-                            span_context=yield_span.get_span_context(),
-                            key=InternalAttributeKeys.CANCELLED,
-                            value=True,
-                        )
-                        break
-
-                    set_span_attribute(
-                        yield_span,
-                        AttributeKeys.JUDGMENT_OUTPUT,
-                        safe_serialize(value),
-                    )
-
-                yield value
-        except Exception as e:
-            main_span.record_exception(e)
-            main_span.set_status(Status(StatusCode.ERROR, str(e)))
-            raise
+        preserved_context = contextvars.copy_context()
+        return _ContextPreservedSyncGeneratorWrapper(
+            self,
+            generator,
+            preserved_context,
+            main_span,
+            None,
+            disable_generator_yield_span,
+        )
 
-    async def _create_traced_async_generator(
+    def _create_traced_async_generator(
         self,
-        async_generator: AsyncIterator[Any],
+        async_generator: AsyncGenerator,
         main_span: Span,
-        base_name: str,
-        attributes: Optional[Dict[str, Any]],
+        disable_generator_yield_span: bool = False,
     ):
         """Create a traced asynchronous generator that wraps each yield in a span."""
-        try:
-            while True:
-                yield_span_name = f"{base_name}_yield"
-                yield_attributes = {
-                    AttributeKeys.JUDGMENT_SPAN_KIND: "async_generator_yield",
-                    **(attributes or {}),
-                }
-
-                async with async_span_context(
-                    self, yield_span_name, yield_attributes, disable_partial_emit=True
-                ) as yield_span:
-                    self.add_agent_attributes_to_span(yield_span)
-
-                    try:
-                        value = await async_generator.__anext__()
-                    except StopAsyncIteration:
-                        # Mark span as cancelled so it won't be exported
-                        self.judgment_processor.set_internal_attribute(
-                            span_context=yield_span.get_span_context(),
-                            key=InternalAttributeKeys.CANCELLED,
-                            value=True,
-                        )
-                        break
-
-                    set_span_attribute(
-                        yield_span,
-                        AttributeKeys.JUDGMENT_OUTPUT,
-                        safe_serialize(value),
-                    )
-
-                yield value
-        except Exception as e:
-            main_span.record_exception(e)
-            main_span.set_status(Status(StatusCode.ERROR, str(e)))
-            raise
+        preserved_context = contextvars.copy_context()
+        return _ContextPreservedAsyncGeneratorWrapper(
+            self,
+            async_generator,
+            preserved_context,
+            main_span,
+            None,
+            disable_generator_yield_span,
+        )
 
     def _wrap_sync(
         self,
@@ -473,19 +479,16 @@ class Tracer(metaclass=SingletonMeta):
         name: Optional[str],
         attributes: Optional[Dict[str, Any]],
         scorer_config: TraceScorerConfig | None = None,
+        disable_generator_yield_span: bool = False,
     ):
-        # Check if this is a generator function - if so, wrap it specially
-        if inspect.isgeneratorfunction(f):
-            return self._wrap_sync_generator_function(
-                f, name, attributes, scorer_config
-            )
-
         @functools.wraps(f)
         def wrapper(*args, **kwargs):
             n = name or f.__qualname__
             with sync_span_context(self, n, attributes) as span:
-                self.add_agent_attributes_to_span(span)
-                self.record_instance_state("before", span)
+                is_return_type_generator = False
+
+                self._inject_judgment_context(span)
+                self._record_instance_state("before", span)
                 try:
                     set_span_attribute(
                         span,
@@ -499,73 +502,40 @@ class Tracer(metaclass=SingletonMeta):
                         self._set_pending_trace_eval(span, scorer_config, args, kwargs)
 
                     result = f(*args, **kwargs)
-                except Exception as user_exc:
-                    span.record_exception(user_exc)
-                    span.set_status(Status(StatusCode.ERROR, str(user_exc)))
-                    raise
-
-                if inspect.isgenerator(result):
-                    set_span_attribute(
-                        span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
-                    )
-                    self.record_instance_state("after", span)
-                    return self._create_traced_sync_generator(
-                        result, span, n, attributes
-                    )
-                else:
-                    set_span_attribute(
-                        span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
-                    )
-                    self.record_instance_state("after", span)
-                    return result
-
-        return wrapper
-
-    def _wrap_sync_generator_function(
-        self,
-        f: Callable,
-        name: Optional[str],
-        attributes: Optional[Dict[str, Any]],
-        scorer_config: TraceScorerConfig | None = None,
-    ):
-        """Wrap a generator function to trace nested function calls within each yield."""
-
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            n = name or f.__qualname__
 
-            with sync_span_context(self, n, attributes) as main_span:
-                self.add_agent_attributes_to_span(main_span)
-                self.record_instance_state("before", main_span)
-
-                try:
-                    set_span_attribute(
-                        main_span,
-                        AttributeKeys.JUDGMENT_INPUT,
-                        safe_serialize(format_inputs(f, args, kwargs)),
-                    )
-
-                    self.judgment_processor.emit_partial()
-
-                    if scorer_config:
-                        self._set_pending_trace_eval(
-                            main_span, scorer_config, args, kwargs
+                    if inspect.isgenerator(result):
+                        is_return_type_generator = True
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
                         )
-
-                    generator = f(*args, **kwargs)
-                    set_span_attribute(
-                        main_span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
-                    )
-                    self.record_instance_state("after", main_span)
-
-                    return self._create_traced_sync_generator(
-                        generator, main_span, n, attributes
-                    )
-
+                        self._record_instance_state("after", span)
+                        return self._create_traced_sync_generator(
+                            result, span, disable_generator_yield_span
+                        )
+                    elif inspect.isasyncgen(result):
+                        is_return_type_generator = True
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
+                        )
+                        self._record_instance_state("after", span)
+                        return self._create_traced_async_generator(
+                            result, span, disable_generator_yield_span
+                        )
+                    else:
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
+                        )
+                        self._record_instance_state("after", span)
+                        self._maybe_clear_customer_context(span)
+                        return result
                 except Exception as user_exc:
-                    main_span.record_exception(user_exc)
-                    main_span.set_status(Status(StatusCode.ERROR, str(user_exc)))
+                    span.record_exception(user_exc)
+                    span.set_status(Status(StatusCode.ERROR, str(user_exc)))
+                    self._maybe_clear_customer_context(span)
                     raise
+                finally:
+                    if not is_return_type_generator:
+                        span.end()
 
         return wrapper
 
@@ -575,19 +545,15 @@ class Tracer(metaclass=SingletonMeta):
         name: Optional[str],
         attributes: Optional[Dict[str, Any]],
         scorer_config: TraceScorerConfig | None = None,
+        disable_generator_yield_span: bool = False,
     ):
-        # Check if this is an async generator function - if so, wrap it specially
-        if inspect.isasyncgenfunction(f):
-            return self._wrap_async_generator_function(
-                f, name, attributes, scorer_config
-            )
-
         @functools.wraps(f)
         async def wrapper(*args, **kwargs):
             n = name or f.__qualname__
             async with async_span_context(self, n, attributes) as span:
-                self.add_agent_attributes_to_span(span)
-                self.record_instance_state("before", span)
+                is_return_type_generator = False
+                self._inject_judgment_context(span)
+                self._record_instance_state("before", span)
                 try:
                     set_span_attribute(
                         span,
@@ -601,73 +567,39 @@ class Tracer(metaclass=SingletonMeta):
                         self._set_pending_trace_eval(span, scorer_config, args, kwargs)
 
                     result = await f(*args, **kwargs)
+                    if inspect.isasyncgen(result):
+                        is_return_type_generator = True
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
+                        )
+                        self._record_instance_state("after", span)
+                        return self._create_traced_async_generator(
+                            result, span, disable_generator_yield_span
+                        )
+                    elif inspect.isgenerator(result):
+                        is_return_type_generator = True
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
+                        )
+                        self._record_instance_state("after", span)
+                        return self._create_traced_sync_generator(
+                            result, span, disable_generator_yield_span
+                        )
+                    else:
+                        set_span_attribute(
+                            span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
+                        )
+                        self._record_instance_state("after", span)
+                        self._maybe_clear_customer_context(span)
+                        return result
                 except Exception as user_exc:
                     span.record_exception(user_exc)
                     span.set_status(Status(StatusCode.ERROR, str(user_exc)))
+                    self._maybe_clear_customer_context(span)
                     raise
-
-                if inspect.isasyncgen(result):
-                    set_span_attribute(
-                        span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
-                    )
-                    self.record_instance_state("after", span)
-                    return self._create_traced_async_generator(
-                        result, span, n, attributes
-                    )
-                else:
-                    set_span_attribute(
-                        span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
-                    )
-                    self.record_instance_state("after", span)
-                    return result
-
-        return wrapper
-
-    def _wrap_async_generator_function(
-        self,
-        f: Callable,
-        name: Optional[str],
-        attributes: Optional[Dict[str, Any]],
-        scorer_config: TraceScorerConfig | None = None,
-    ):
-        """Wrap an async generator function to trace nested function calls within each yield."""
-
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            n = name or f.__qualname__
-
-            with sync_span_context(self, n, attributes) as main_span:
-                self.add_agent_attributes_to_span(main_span)
-                self.record_instance_state("before", main_span)
-
-                try:
-                    set_span_attribute(
-                        main_span,
-                        AttributeKeys.JUDGMENT_INPUT,
-                        safe_serialize(format_inputs(f, args, kwargs)),
-                    )
-
-                    self.judgment_processor.emit_partial()
-
-                    if scorer_config:
-                        self._set_pending_trace_eval(
-                            main_span, scorer_config, args, kwargs
-                        )
-
-                    async_generator = f(*args, **kwargs)
-                    set_span_attribute(
-                        main_span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
-                    )
-                    self.record_instance_state("after", main_span)
-
-                    return self._create_traced_async_generator(
-                        async_generator, main_span, n, attributes
-                    )
-
-                except Exception as user_exc:
-                    main_span.record_exception(user_exc)
-                    main_span.set_status(Status(StatusCode.ERROR, str(user_exc)))
-                    raise
+                finally:
+                    if not is_return_type_generator:
+                        span.end()
 
         return wrapper
 
@@ -704,6 +636,7 @@ class Tracer(metaclass=SingletonMeta):
         span_name: str | None = None,
         attributes: Optional[Dict[str, Any]] = None,
         scorer_config: TraceScorerConfig | None = None,
+        disable_generator_yield_span: bool = False,
     ) -> Callable | None:
         if func is None:
             return partial(
@@ -712,6 +645,7 @@ class Tracer(metaclass=SingletonMeta):
                 span_name=span_name,
                 attributes=attributes,
                 scorer_config=scorer_config,
+                disable_generator_yield_span=disable_generator_yield_span,
             )
 
         if not self.enable_monitoring:
@@ -724,10 +658,14 @@ class Tracer(metaclass=SingletonMeta):
             **(attributes or {}),
         }
 
-        if inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func):
-            return self._wrap_async(func, name, func_attributes, scorer_config)
+        if inspect.iscoroutinefunction(func):
+            return self._wrap_async(
+                func, name, func_attributes, scorer_config, disable_generator_yield_span
+            )
         else:
-            return self._wrap_sync(func, name, func_attributes, scorer_config)
+            return self._wrap_sync(
+                func, name, func_attributes, scorer_config, disable_generator_yield_span
+            )
 
     @overload
     def agent(
@@ -1003,6 +941,182 @@ def format_inputs(
         return {}
 
 
+class _ContextPreservedSyncGeneratorWrapper:
+    """Sync generator wrapper that ensures each iteration runs in preserved context."""
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        generator: Generator,
+        context: contextvars.Context,
+        span: Span,
+        transform_fn: Optional[Callable[[Iterable], str]],
+        disable_generator_yield_span: bool = False,
+    ) -> None:
+        self.tracer = tracer
+        self.generator = generator
+        self.context = context
+        self.span = span
+        self.transform_fn = transform_fn
+        self._finished = False
+        self.disable_generator_yield_span = disable_generator_yield_span
+
+    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)
+
+            if not self.disable_generator_yield_span:
+                with use_span(self.span):
+                    span_name = (
+                        str(self.span.name)
+                        if hasattr(self.span, "name")
+                        else "generator_item"
+                    )  # type: ignore[attr-defined]
+                    with self.tracer.get_tracer().start_as_current_span(
+                        span_name,
+                        attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "generator_item"},
+                        end_on_exit=True,
+                    ) as child_span:
+                        set_span_attribute(
+                            child_span,
+                            AttributeKeys.JUDGMENT_OUTPUT,
+                            safe_serialize(item),
+                        )
+
+            return item
+
+        except StopIteration:
+            # Handle output and span cleanup when generator is exhausted
+            if not self._finished:
+                set_span_attribute(
+                    self.span, AttributeKeys.JUDGMENT_SPAN_KIND, "generator"
+                )
+                self.span.end()
+                self._finished = True
+
+            raise  # Re-raise StopIteration
+
+        except Exception as e:
+            if not self._finished:
+                self.span.record_exception(e)
+                self.span.set_status(
+                    Status(StatusCode.ERROR, str(e) or type(e).__name__)
+                )
+                self.tracer._maybe_clear_customer_context(self.span)
+                self.span.end()
+                self._finished = True
+
+            raise
+
+    def close(self) -> None:
+        """Close the generator (minimal implementation)."""
+        try:
+            self.generator.close()
+        finally:
+            if not self._finished:
+                set_span_attribute(
+                    self.span, AttributeKeys.JUDGMENT_SPAN_KIND, "generator"
+                )
+                self.tracer._maybe_clear_customer_context(self.span)
+                self.span.end()
+                self._finished = True
+
+
+class _ContextPreservedAsyncGeneratorWrapper:
+    """Async generator wrapper that ensures each iteration runs in preserved context."""
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        generator: AsyncGenerator,
+        context: contextvars.Context,
+        span: Span,
+        transform_fn: Optional[Callable[[Iterable], str]],
+        disable_generator_yield_span: bool = False,
+    ) -> None:
+        self.tracer = tracer
+        self.generator = generator
+        self.context = context
+        self.span = span
+        self.transform_fn = transform_fn
+        self._finished = False
+        self.disable_generator_yield_span = disable_generator_yield_span
+
+    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__()
+
+            if not self.disable_generator_yield_span:
+                with use_span(self.span):
+                    span_name = (
+                        str(self.span.name)
+                        if hasattr(self.span, "name")
+                        else "generator_item"
+                    )  # type: ignore[attr-defined]
+                    with self.tracer.get_tracer().start_as_current_span(
+                        span_name,
+                        attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "generator_item"},
+                        end_on_exit=True,
+                    ) as child_span:
+                        set_span_attribute(
+                            child_span,
+                            AttributeKeys.JUDGMENT_OUTPUT,
+                            safe_serialize(item),
+                        )
+
+            return item
+
+        except StopAsyncIteration:
+            # Handle output and span cleanup when generator is exhausted
+            if not self._finished:
+                set_span_attribute(
+                    self.span, AttributeKeys.JUDGMENT_SPAN_KIND, "generator"
+                )
+                self.span.end()
+                self._finished = True
+            raise  # Re-raise StopAsyncIteration
+        except Exception as e:
+            if not self._finished:
+                self.span.record_exception(e)
+                self.span.set_status(
+                    Status(StatusCode.ERROR, str(e) or type(e).__name__)
+                )
+                self.tracer._maybe_clear_customer_context(self.span)
+                self.span.end()
+                self._finished = True
+
+            raise
+
+    async def aclose(self) -> None:
+        """Close the async generator (minimal implementation)."""
+        try:
+            await self.generator.aclose()
+        finally:
+            if not self._finished:
+                set_span_attribute(
+                    self.span, AttributeKeys.JUDGMENT_SPAN_KIND, "generator"
+                )
+                self.tracer._maybe_clear_customer_context(self.span)
+                self.span.end()
+                self._finished = True
+
+
 __all__ = [
     "Tracer",
     "wrap",