Enrich Laminar data (#3650)

skyvern/forge/sdk/trace/__init__.py

@@ -90,3 +90,13 @@ class TraceManager:
     @staticmethod
     def set_trace_provider(trace_provider: BaseTrace) -> None:
         TraceManager.__instance = trace_provider
+
+    @staticmethod
+    def add_task_completion_tag(status: str) -> None:
+        """Add a completion tag to the current trace based on task/workflow status."""
+        TraceManager.__instance.add_task_completion_tag(status)
+
+    @staticmethod
+    def add_experiment_metadata(experiment_data: dict[str, Any]) -> None:
+        """Add experiment metadata to the current trace."""
+        TraceManager.__instance.add_experiment_metadata(experiment_data)

skyvern/forge/sdk/trace/base.py

@@ -26,6 +26,12 @@ class BaseTrace(ABC):
     ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
         pass
 
+    def add_task_completion_tag(self, status: str) -> None:
+        """Add a completion tag to the current trace based on task/workflow status."""
+
+    def add_experiment_metadata(self, experiment_data: dict[str, Any]) -> None:
+        """Add experiment metadata to the current trace."""
+
 
 class NoOpTrace(BaseTrace):
     def traced(

skyvern/forge/sdk/trace/experiment_utils.py

@@ -0,0 +1,97 @@
+"""Utilities for collecting and formatting experiment data for Laminar tracing."""
+
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from skyvern.forge.sdk.core import skyvern_context
+
+if TYPE_CHECKING:
+    from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider
+
+LOG = structlog.get_logger()
+
+
+async def collect_experiment_metadata_safely(
+    experimentation_provider: "BaseExperimentationProvider",
+) -> dict[str, Any]:
+    """
+    Safely collect experiment-related metadata from the current context.
+
+    This is a safe wrapper around collect_experiment_metadata() that ensures
+    any exceptions are caught and handled gracefully.
+
+    Args:
+        experimentation_provider: The experimentation provider to use for fetching experiment data.
+
+    Returns:
+        Dictionary containing experiment data, or empty dict if collection fails.
+    """
+    try:
+        return await collect_experiment_metadata(experimentation_provider)
+    except Exception:
+        LOG.warning("Failed to collect experiment metadata", exc_info=True)
+        return {}
+
+
+async def collect_experiment_metadata(
+    experimentation_provider: "BaseExperimentationProvider",
+) -> dict[str, Any]:
+    """
+    Collect experiment-related metadata from the current context.
+
+    Args:
+        experimentation_provider: The experimentation provider to use for fetching experiment data.
+
+    Returns:
+        Dictionary containing experiment data that can be added to Laminar traces.
+    """
+    # Get the current context
+    context = skyvern_context.current()
+    if not context or not context.run_id:
+        return {}
+
+    # Use run_id as the distinct_id for experiments
+    distinct_id = context.run_id
+    organization_id = context.organization_id
+
+    if not distinct_id or not organization_id:
+        return {}
+
+    experiment_metadata: dict[str, Any] = {}
+
+    try:
+        # Only collect critical experiment flags that are relevant for tracing
+        experiment_flags = [
+            "LLM_NAME",
+            "LLM_SECONDARY_NAME",
+            # Add more experiment flags as needed
+            "PROMPT_CACHING_ENABLED",
+            "THINKING_BUDGET_OPTIMIZATION",
+        ]
+
+        for flag in experiment_flags:
+            try:
+                # Get the experiment value (already cached by experimentation provider)
+                value = experimentation_provider.get_value_cached(
+                    flag, distinct_id, properties={"organization_id": organization_id}
+                )
+
+                # Get the payload if available (already cached by experimentation provider)
+                payload = experimentation_provider.get_payload_cached(
+                    flag, distinct_id, properties={"organization_id": organization_id}
+                )
+
+                # Only include if we have actual experiment data
+                if value is not None or payload is not None:
+                    experiment_metadata[f"experiment_{flag}"] = {"value": value, "payload": payload}
+
+            except Exception:
+                # Silently skip failed experiments
+                continue
+
+    except Exception:
+        # Silently fail if experimentation provider is not available
+        pass
+
+    return experiment_metadata

skyvern/forge/sdk/trace/lmnr.py

@@ -31,3 +31,35 @@ class LaminarTrace(BaseTrace):
         **kwargs: Any,
     ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
         return observe(name=name, ignore_output=True, metadata=metadata, tags=tags, **kwargs)
+
+    def add_task_completion_tag(self, status: str) -> None:
+        """Add a completion tag to the current trace based on task/workflow status."""
+        try:
+            # Get the current trace ID
+            trace_id = Laminar.get_trace_id()
+            if trace_id is None:
+                return
+
+            # Map status to appropriate tag
+            status_tag_map = {
+                "completed": "COMPLETION",
+                "failed": "FAILURE",
+                "timed_out": "TIMEOUT",
+                "canceled": "CANCELED",
+                "terminated": "TERMINATED",
+            }
+
+            tag = status_tag_map.get(status, "FAILURE")
+            Laminar.set_span_tags([tag])
+        except Exception:
+            # Silently fail if tracing is not available or there's an error
+            pass
+
+    def add_experiment_metadata(self, experiment_data: dict[str, Any]) -> None:
+        """Add experiment metadata to the current trace."""
+        try:
+            # Add experiment metadata to the current trace
+            Laminar.set_trace_metadata(experiment_data)
+        except Exception:
+            # Silently fail if tracing is not available or there's an error
+            pass