ai-pipeline-core 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

ai_pipeline_core/__init__.py

@@ -64,7 +64,7 @@ from .prompt_manager import PromptManager
 from .settings import Settings
 from .testing import disable_run_logger, prefect_test_harness
 
-__version__ = "0.4.6"
+__version__ = "0.4.8"
 
 __all__ = [
     "AIMessageType",

ai_pipeline_core/deployment/base.py

@@ -661,7 +661,10 @@ class PipelineDeployment(Generic[TOptions, TResult]):
         except Exception as e:
             logger.warning(f"Failed to initialize observability: {e}")
             with contextlib.suppress(Exception):
-                Laminar.initialize(export_timeout_seconds=15)
+                # Use canonical initializer to ensure consistent Laminar setup
+                from ai_pipeline_core.observability import tracing
+
+                tracing._initialise_laminar()
 
         deployment = self
 
@@ -882,13 +885,34 @@ class PipelineDeployment(Generic[TOptions, TResult]):
             options: FlowOptions,
             context: DeploymentContext,
         ) -> DeploymentResult:
+            # Initialize observability for remote workers
+            try:
+                initialize_observability()
+            except Exception as e:
+                logger.warning(f"Failed to initialize observability: {e}")
+                with contextlib.suppress(Exception):
+                    # Use canonical initializer to ensure consistent Laminar setup
+                    from ai_pipeline_core.observability import tracing
+
+                    tracing._initialise_laminar()
+
+            # Set session ID from Prefect flow run for trace grouping
+            flow_run_id = str(runtime.flow_run.get_id()) if runtime.flow_run else str(uuid4())  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownArgumentType]
+            os.environ["LMNR_SESSION_ID"] = flow_run_id
+
             store = create_document_store(
                 settings,
                 summary_generator=_build_summary_generator(),
             )
             set_document_store(store)
             try:
-                return await deployment.run(project_name, documents, cast(Any, options), context)
+                # Create parent span to group all traces under a single deployment trace
+                with Laminar.start_as_current_span(
+                    name=f"{deployment.name}-{project_name}",
+                    input={"project_name": project_name, "options": options.model_dump()},
+                    session_id=flow_run_id,
+                ):
+                    return await deployment.run(project_name, documents, cast(Any, options), context)
             finally:
                 store.shutdown()
                 set_document_store(None)

ai_pipeline_core/observability/_initialization.py

@@ -8,7 +8,6 @@ import importlib
 from typing import Any, Protocol
 from uuid import UUID
 
-from lmnr import Laminar
 from opentelemetry import trace as otel_trace
 from pydantic import BaseModel, ConfigDict
 
@@ -180,10 +179,12 @@ def initialize_observability(config: ObservabilityConfig | None = None) -> None:
     if config is None:
         config = _build_config_from_settings()
 
-    # 1. Laminar
+    # 1. Laminar - use canonical initializer from tracing module
     if config.has_lmnr:
         try:
-            Laminar.initialize(project_api_key=config.lmnr_project_api_key, export_timeout_seconds=15)
+            from ai_pipeline_core.observability import tracing  # noqa: PLC0415
+
+            tracing._initialise_laminar()
             logger.info("Laminar initialized")
         except Exception as e:
             logger.warning(f"Laminar initialization failed: {e}")

ai_pipeline_core/observability/tracing.py

@@ -10,6 +10,7 @@ import contextlib
 import inspect
 import json
 import os
+import threading
 from collections.abc import Callable
 from functools import wraps
 from typing import Any, Literal, ParamSpec, TypeVar, cast, overload
@@ -220,19 +221,42 @@ class TraceInfo(BaseModel):
 # ---------------------------------------------------------------------------
 
 
+_laminar_initialized = False
+_laminar_init_lock = threading.Lock()
+
+
 def _initialise_laminar() -> None:
-    """Initialize Laminar SDK with project configuration.
+    """Initialize Laminar SDK with project configuration (lazy, once per process).
 
     Sets up the Laminar observability client with the project API key
     from settings. Disables automatic OpenAI instrumentation to avoid
     conflicts with our custom tracing.
 
-    Called once per process. Multiple calls are safe (Laminar handles idempotency).
+    IMPORTANT: This is called lazily at first trace execution (not at decoration time)
+    to allow LMNR_SPAN_CONTEXT environment variable to be set before initialization.
+    Laminar reads LMNR_SPAN_CONTEXT during initialize() to establish parent context
+    for cross-process tracing.
+
+    Uses double-checked locking pattern for thread safety. The flag is set AFTER
+    successful initialization to prevent permanently disabled tracing on init failure.
     """
-    if settings.lmnr_project_api_key:
-        Laminar.initialize(
-            project_api_key=settings.lmnr_project_api_key, disabled_instruments=[Instruments.OPENAI] if Instruments.OPENAI else [], export_timeout_seconds=15
-        )
+    global _laminar_initialized  # noqa: PLW0603
+
+    # Fast path: already initialized (no lock needed)
+    if _laminar_initialized:
+        return
+
+    with _laminar_init_lock:
+        # Double-check inside lock
+        if _laminar_initialized:
+            return
+
+        if settings.lmnr_project_api_key:
+            disabled = [Instruments.OPENAI] if Instruments.OPENAI else []
+            Laminar.initialize(project_api_key=settings.lmnr_project_api_key, disabled_instruments=disabled, export_timeout_seconds=15)
+
+        # Set flag AFTER successful initialization
+        _laminar_initialized = True
 
 
 # Overload for calls like @trace(name="...", level="debug")
@@ -400,7 +424,9 @@ def trace(  # noqa: UP047
             return f
 
         # --- Pre-computation (done once when the function is decorated) ---
-        _initialise_laminar()
+        # NOTE: _initialise_laminar() is NOT called here (at decoration/import time)
+        # to allow LMNR_SPAN_CONTEXT to be set before Laminar.initialize() runs.
+        # It's called lazily in the wrapper functions at first execution.
         sig = inspect.signature(f)
         is_coroutine = inspect.iscoroutinefunction(f)
         observe_name = name or f.__name__
@@ -550,6 +576,9 @@ def trace(  # noqa: UP047
             Returns:
                 The result of the wrapped function.
             """
+            # Lazy initialization: called at first execution, not at decoration time.
+            # This allows LMNR_SPAN_CONTEXT to be set before Laminar.initialize().
+            _initialise_laminar()
             observe_params = _prepare_and_get_observe_params(kwargs)
             observed_func = bound_observe(**observe_params)(f)
             return observed_func(*args, **kwargs)
@@ -561,6 +590,9 @@ def trace(  # noqa: UP047
             Returns:
                 The result of the wrapped function.
             """
+            # Lazy initialization: called at first execution, not at decoration time.
+            # This allows LMNR_SPAN_CONTEXT to be set before Laminar.initialize().
+            _initialise_laminar()
             observe_params = _prepare_and_get_observe_params(kwargs)
             observed_func = bound_observe(**observe_params)(f)
             return await observed_func(*args, **kwargs)  # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType]

{ai_pipeline_core-0.4.7.dist-info → ai_pipeline_core-0.4.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.4.7
+Version: 0.4.8
 Summary: Core utilities for AI-powered processing pipelines using prefect
 Project-URL: Homepage, https://github.com/bbarwik/ai-pipeline-core
 Project-URL: Repository, https://github.com/bbarwik/ai-pipeline-core

{ai_pipeline_core-0.4.7.dist-info → ai_pipeline_core-0.4.8.dist-info}/RECORD

@@ -1,11 +1,11 @@
-ai_pipeline_core/__init__.py,sha256=QAQSyrKafsNov4dy9vNqTVarh2nDdrthMfYh7X-3Mcg,3270
+ai_pipeline_core/__init__.py,sha256=aJwyMqt4ESan14iAS9guaHbDRk1F97PbOeHBvxShhD4,3270
 ai_pipeline_core/exceptions.py,sha256=csAl7vq6xjSFBF8-UM9WZODCbhsOdOG5zH6IbA8iteM,1280
 ai_pipeline_core/prompt_manager.py,sha256=3wFkL5rrjtUT1cLInkgyhS8hKnO4MeD1cdXAEuLhgoE,9459
 ai_pipeline_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_pipeline_core/settings.py,sha256=BUz8JEFfJQrdE4rNOhQWwxnTrfekLjWkoy-3wDZQ7PY,5142
 ai_pipeline_core/testing.py,sha256=jIRrLxNvTwdamucfJoHET2qMeRhhMZV9uEJXO5vAfis,279
 ai_pipeline_core/deployment/__init__.py,sha256=wTkVK6gcEQvqBajFMTAuodRONpN25yHbR1jtcumf0WQ,900
-ai_pipeline_core/deployment/base.py,sha256=bGSnDdrw6cLM_TItAiwptnwApbw5wkoIGY9pnwDvOTQ,37485
+ai_pipeline_core/deployment/base.py,sha256=JhaEYok3USOOGm9VrqayZTQD-i_VyOch7Gdv-jQfqsQ,38785
 ai_pipeline_core/deployment/contract.py,sha256=a1qbHhneTGB27oSOUy79CUIhOIzOoq37M63XoIMzA4Y,1952
 ai_pipeline_core/deployment/deploy.py,sha256=y5FxKMm7nGwkjzA74pTffO7A82MaDuajx6LHGTem8bI,24662
 ai_pipeline_core/deployment/helpers.py,sha256=yVtGFUs4AFXkpLkiQ_ale0nXXt5btfWSb5PAbikQHNs,3312
@@ -48,10 +48,10 @@ ai_pipeline_core/logging/logging_config.py,sha256=JnTarGSSkpi7eqR7N13TLKeuwNCvZg
 ai_pipeline_core/logging/logging_mixin.py,sha256=Jn3x0xvSwSjbAMfWELMOEfffWBB1u4IeIr7M2-55CJs,7191
 ai_pipeline_core/observability/__init__.py,sha256=km2nIiY3aYH13s2m4nR91erQG4qKnGuvQkrKDdVW3bw,720
 ai_pipeline_core/observability/_document_tracking.py,sha256=tXv6rbGIuxOYdq22aVbyn9Ve5EhYHPnrYCE-kj2NGXI,5428
-ai_pipeline_core/observability/_initialization.py,sha256=GfwRHpg90Og3PzmG1ZUilJVXoFx9BIWpbMgXxJ5Alqk,6747
+ai_pipeline_core/observability/_initialization.py,sha256=kg5erwWLyKi-kWl21A8jXZTod05LtP-vP0VOd2zhAzo,6790
 ai_pipeline_core/observability/_logging_bridge.py,sha256=T3PpkgoI0YKN2vvBJEHzR5rFMFNHq9REHJs7PQX2VQk,2053
 ai_pipeline_core/observability/_summary.py,sha256=GAZXzXVkwUcubSiGb5DgkHfO1gGwx6pYoDz6RUJmL5k,3390
-ai_pipeline_core/observability/tracing.py,sha256=KhIXSl5fe39UE1Eokz9-1fe5biX6anKbwZDmXY_Z2LU,27050
+ai_pipeline_core/observability/tracing.py,sha256=RZ702N392_nP4HZvxVWv6x3c6hDBZzkdNOtqPTqmFMs,28521
 ai_pipeline_core/observability/_debug/__init__.py,sha256=V8pbgdQOx-7oFKQ_sNzAZ1-oq5c73P4kVjEClZDXe8k,942
 ai_pipeline_core/observability/_debug/_auto_summary.py,sha256=LMvETvx_RPKF8srewCKwjigTiWs3KfDmQAYYSuVybIM,2687
 ai_pipeline_core/observability/_debug/_config.py,sha256=CWfnK-F3knUuOQ34y_CjmU3l67J85NIZ3siftYhevc0,3367
@@ -70,7 +70,7 @@ ai_pipeline_core/observability/_tracking/_writer.py,sha256=xZjwYyIxDzzzPxqkKjYAY
 ai_pipeline_core/pipeline/__init__.py,sha256=uMv1jwSyq8Ym8Hbn5097twBJLdwN1iMeqnVM4EWyrhA,282
 ai_pipeline_core/pipeline/decorators.py,sha256=CDJAeOjGLt5Ewc0Jc9zEuwLZwKyutOv89LSRS9dcXmI,37456
 ai_pipeline_core/pipeline/options.py,sha256=KF4FcT085-IwX8r649v0a9ua5xnApM0qG2wJHWbq39A,438
-ai_pipeline_core-0.4.7.dist-info/METADATA,sha256=yFjXJ9fHXFtmrF2jIFx62k5spfR1PEipR_Uekbn3bmo,29947
-ai_pipeline_core-0.4.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ai_pipeline_core-0.4.7.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.4.7.dist-info/RECORD,,
+ai_pipeline_core-0.4.8.dist-info/METADATA,sha256=Ftytzz5IBhleZK7ce8HbE4XMc8pcjdVdOe2oN-fpluA,29947
+ai_pipeline_core-0.4.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ai_pipeline_core-0.4.8.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.4.8.dist-info/RECORD,,