feat: first pass at carrying contextvars though async flows (#878)

* first pass at carrying contextvars though async flows

* update docstring

* utilize with context instead of iter token

* add todo regarding log context and contextfilter

Author: AngeloDanducci

mellea/backends/huggingface.py

@@ -55,6 +55,7 @@
     instrument_generate_from_raw,
     start_generate_span,
 )
+from ..telemetry.context import generate_request_id, with_context
 from .adapters import (
     AdapterMixin,
     AdapterType,
@@ -389,71 +390,80 @@ async def _generate_from_context(
         span = start_generate_span(
             backend=self, action=action, ctx=ctx, format=format, tool_calls=tool_calls
         )
-        await self.do_generate_walk(action)
 
-        # Upsert model options.
-        model_opts = self._simplify_and_merge(model_options)
+        with with_context(
+            request_id=generate_request_id(),
+            model_id=str(getattr(self, "model_id", "unknown")),
+        ):
+            await self.do_generate_walk(action)
 
-        # Requirements can be automatically rerouted to a requirement adapter.
-        if isinstance(action, Requirement):
-            # See docs/dev/requirement_aLoRA_rerouting.md
-            reroute_to_alora = self.default_to_constraint_checking_alora
-            adapter_name = "requirement_check"
+            # Upsert model options.
+            model_opts = self._simplify_and_merge(model_options)
 
-            if isinstance(action, ALoraRequirement):
-                reroute_to_alora = True
-                adapter_name = action.intrinsic_name
-                alora_action = action
-            else:
-                assert action.description is not None, (
-                    "must have a description when generating from a requirement"
-                )
-                alora_action = ALoraRequirement(action.description, adapter_name)
+            # Requirements can be automatically rerouted to a requirement adapter.
+            if isinstance(action, Requirement):
+                # See docs/dev/requirement_aLoRA_rerouting.md
+                reroute_to_alora = self.default_to_constraint_checking_alora
+                adapter_name = "requirement_check"
 
-            # Check if a requirement_check (or AloraRequirement specified) adapter
-            # exists.
-            alora_req_adapter = get_adapter_for_intrinsic(
-                adapter_name, [AdapterType.ALORA], self._added_adapters
-            )
-            if alora_req_adapter is None:
-                # Log a warning if using an AloraRequirement but no adapter fit.
-                if reroute_to_alora and isinstance(action, ALoraRequirement):
-                    MelleaLogger.get_logger().warning(
-                        f"attempted to use an AloraRequirement but backend {self} doesn't have the specified adapter added {adapter_name}; defaulting to regular generation"
+                if isinstance(action, ALoraRequirement):
+                    reroute_to_alora = True
+                    adapter_name = action.intrinsic_name
+                    alora_action = action
+                else:
+                    assert action.description is not None, (
+                        "must have a description when generating from a requirement"
                     )
-                reroute_to_alora = False
+                    alora_action = ALoraRequirement(action.description, adapter_name)
+
+                # Check if a requirement_check (or AloraRequirement specified) adapter
+                # exists.
+                alora_req_adapter = get_adapter_for_intrinsic(
+                    adapter_name, [AdapterType.ALORA], self._added_adapters
+                )
+                if alora_req_adapter is None:
+                    # Log a warning if using an AloraRequirement but no adapter fit.
+                    if reroute_to_alora and isinstance(action, ALoraRequirement):
+                        MelleaLogger.get_logger().warning(
+                            f"attempted to use an AloraRequirement but backend {self} doesn't have the specified adapter added {adapter_name}; defaulting to regular generation"
+                        )
+                    reroute_to_alora = False
 
-            if issubclass(type(action), LLMaJRequirement):
-                reroute_to_alora = False
+                if issubclass(type(action), LLMaJRequirement):
+                    reroute_to_alora = False
 
-            if reroute_to_alora:
-                # Keep the alora requirement handling separate for now.
+                if reroute_to_alora:
+                    # Keep the alora requirement handling separate for now.
+                    mot = await self._generate_from_intrinsic(
+                        alora_action, ctx, model_options=model_opts
+                    )
+                    # Store span for telemetry
+                    if span is not None:
+                        mot._meta["_telemetry_span"] = span
+                    return mot, ctx.add(alora_action).add(mot)
+
+            elif isinstance(action, Intrinsic):
                 mot = await self._generate_from_intrinsic(
-                    alora_action, ctx, model_options=model_opts
+                    action, ctx, model_options=model_opts
                 )
                 # Store span for telemetry
                 if span is not None:
                     mot._meta["_telemetry_span"] = span
-                return mot, ctx.add(alora_action).add(mot)
+                return mot, ctx.add(action).add(mot)
 
-        elif isinstance(action, Intrinsic):
-            mot = await self._generate_from_intrinsic(
-                action, ctx, model_options=model_opts
+            mot = await self._generate_from_context_standard(
+                action,
+                ctx,
+                _format=format,
+                model_options=model_opts,
+                tool_calls=tool_calls,
             )
-            # Store span for telemetry
+
+            # Store span in metadata for post_processing to record telemetry
             if span is not None:
                 mot._meta["_telemetry_span"] = span
-            return mot, ctx.add(action).add(mot)
 
-        mot = await self._generate_from_context_standard(
-            action, ctx, _format=format, model_options=model_opts, tool_calls=tool_calls
-        )
-
-        # Store span in metadata for post_processing to record telemetry
-        if span is not None:
-            mot._meta["_telemetry_span"] = span
-
-        return mot, ctx.add(action).add(mot)
+            return mot, ctx.add(action).add(mot)
 
     def _generate_with_adapter_lock(
         self, adapter_name: str, generate_func: Callable, *args, **kwargs

mellea/backends/litellm.py

@@ -45,6 +45,7 @@
     instrument_generate_from_raw,
     start_generate_span,
 )
+from ..telemetry.context import generate_request_id, with_context
 from .backend import FormatterBackend
 from .model_options import ModelOption
 from .tools import (
@@ -166,13 +167,16 @@ async def _generate_from_context(
         span = start_generate_span(
             backend=self, action=action, ctx=ctx, format=format, tool_calls=tool_calls
         )
-        mot = await self._generate_from_chat_context_standard(
-            action,
-            ctx,
-            _format=format,
-            model_options=model_options,
-            tool_calls=tool_calls,
-        )
+
+        _model_id_str = str(getattr(self, "model_id", "unknown"))
+        with with_context(request_id=generate_request_id(), model_id=_model_id_str):
+            mot = await self._generate_from_chat_context_standard(
+                action,
+                ctx,
+                _format=format,
+                model_options=model_options,
+                tool_calls=tool_calls,
+            )
 
         # Store span for telemetry recording in post_processing
         if span is not None:

mellea/backends/ollama.py

@@ -30,7 +30,9 @@
 from ..telemetry.backend_instrumentation import (
     instrument_generate_from_context,
     instrument_generate_from_raw,
+    start_generate_span,
 )
+from ..telemetry.context import generate_request_id, with_context
 from .backend import FormatterBackend
 from .model_options import ModelOption
 from .tools import add_tools_from_context_actions, add_tools_from_model_options
@@ -286,21 +288,22 @@ async def _generate_from_context(
             tuple[ModelOutputThunk[C], Context]: A thunk holding the (lazy) model output
                 and an updated context that includes ``action`` and the new output.
         """
-        from ..telemetry.backend_instrumentation import start_generate_span
-
         # Start span without auto-closing (will be closed in post_processing)
         span = start_generate_span(self, action, ctx, format, tool_calls)
 
         assert ctx.is_chat_context, (
             "The ollama backend only supports chat-like contexts."
         )
-        mot = await self.generate_from_chat_context(
-            action,
-            ctx,
-            _format=format,
-            model_options=model_options,
-            tool_calls=tool_calls,
-        )
+
+        _model_id_str = str(getattr(self, "model_id", "unknown"))
+        with with_context(request_id=generate_request_id(), model_id=_model_id_str):
+            mot = await self.generate_from_chat_context(
+                action,
+                ctx,
+                _format=format,
+                model_options=model_options,
+                tool_calls=tool_calls,
+            )
 
         # Store span for telemetry recording and closing in post_processing
         if span is not None:

mellea/backends/openai.py

@@ -47,7 +47,9 @@
 from ..telemetry.backend_instrumentation import (
     instrument_generate_from_context,
     instrument_generate_from_raw,
+    start_generate_span,
 )
+from ..telemetry.context import generate_request_id, with_context
 from .backend import FormatterBackend
 from .model_options import ModelOption
 from .tools import (
@@ -357,8 +359,6 @@ async def _generate_from_context(
             tuple[ModelOutputThunk[C], Context]: A thunk holding the (lazy) model output
                 and an updated context that includes ``action`` and the new output.
         """
-        from ..telemetry.backend_instrumentation import start_generate_span
-
         assert ctx.is_chat_context, NotImplementedError(
             "The Openai backend only supports chat-like contexts."
         )
@@ -372,13 +372,15 @@ async def _generate_from_context(
             backend=self, action=action, ctx=ctx, format=format, tool_calls=tool_calls
         )
 
-        result = await self.generate_from_chat_context(
-            action,
-            ctx,
-            _format=format,
-            model_options=model_options,
-            tool_calls=tool_calls,
-        )
+        _model_id_str = str(getattr(self, "model_id", "unknown"))
+        with with_context(request_id=generate_request_id(), model_id=_model_id_str):
+            result = await self.generate_from_chat_context(
+                action,
+                ctx,
+                _format=format,
+                model_options=model_options,
+                tool_calls=tool_calls,
+            )
         # Store span in ModelOutputThunk for later use in post_processing
         mot, new_ctx = result
         if span is not None:

mellea/backends/watsonx.py

@@ -48,6 +48,7 @@
     instrument_generate_from_raw,
     start_generate_span,
 )
+from ..telemetry.context import generate_request_id, with_context
 from .backend import FormatterBackend
 from .model_options import ModelOption
 from .tools import (
@@ -304,13 +305,16 @@ async def _generate_from_context(
         span = start_generate_span(
             backend=self, action=action, ctx=ctx, format=format, tool_calls=tool_calls
         )
-        mot = await self.generate_from_chat_context(
-            action,
-            ctx,
-            _format=format,
-            model_options=model_options,
-            tool_calls=tool_calls,
-        )
+
+        _model_id_str = str(getattr(self, "model_id", "unknown"))
+        with with_context(request_id=generate_request_id(), model_id=_model_id_str):
+            mot = await self.generate_from_chat_context(
+                action,
+                ctx,
+                _format=format,
+                model_options=model_options,
+                tool_calls=tool_calls,
+            )
 
         # Store span in metadata for post_processing to record telemetry
         if span is not None:

mellea/core/utils.py

@@ -36,6 +36,8 @@
 except ImportError:
     _OTEL_AVAILABLE = False
 
+from ..telemetry.context import _CONTEXT_VARS as _telemetry_vars, MelleaContextFilter
+
 # ---------------------------------------------------------------------------
 # Per-task/coroutine context fields (safe for asyncio — each Task gets its own copy)
 # ---------------------------------------------------------------------------
@@ -376,13 +378,22 @@ def _build_log_dict(self, record: logging.LogRecord) -> dict[str, Any]:
         # Static extra fields (constructor-level)
         log_record.update(self._extra)
 
-        # Dynamic context fields — prefer record attributes (set by
-        # ContextFilter) but fall back to ContextVar storage so the
-        # formatter works standalone without a filter attached.
+        # Dynamic context fields — prefer record attributes (set by ContextFilter /
+        # MelleaContextFilter) but fall back to ContextVar storage so the formatter
+        # works standalone without filters attached.
         context_fields: dict[str, Any] = _log_context.get()
         for key, value in context_fields.items():
             log_record[key] = getattr(record, key, value)
 
+        # Telemetry context fields (session_id, request_id, model_id, sampling_iteration).
+        # MelleaContextFilter stamps these onto the record before formatters run; read
+        # them back off the record here so they appear in JSON output.  Fall back to the
+        # ContextVar directly so the formatter still works without the filter attached.
+        for key, var in _telemetry_vars.items():
+            value = getattr(record, key, var.get())
+            if value is not None:
+                log_record.setdefault(key, value)
+
         return log_record
 
     def format(self, record: logging.LogRecord) -> str:
@@ -521,6 +532,9 @@ def get_logger() -> logging.Logger:
                     logger.addFilter(ContextFilter())
                     logger.addFilter(OtelTraceFilter())
 
+                    # Inject telemetry context fields (session_id, request_id, etc.)
+                    logger.addFilter(MelleaContextFilter())
+
                     # Only set default level if user hasn't already configured it
                     if logger.level == logging.NOTSET:
                         logger.setLevel(MelleaLogger._resolve_log_level())