laminar integration (#2887)

skyvern/forge/sdk/trace/__init__.py

@@ -0,0 +1,92 @@
+from functools import wraps
+from typing import Any, Awaitable, Callable, ParamSpec, TypeVar
+
+from skyvern.forge.sdk.core import skyvern_context
+from skyvern.forge.sdk.settings_manager import SettingsManager
+from skyvern.forge.sdk.trace.base import BaseTrace, NoOpTrace
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class TraceManager:
+    __instance: BaseTrace = NoOpTrace()
+
+    @staticmethod
+    def traced_async(
+        *,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **trace_parameters: Any,
+    ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+        def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
+            @wraps(func)
+            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                new_metadata: dict[str, Any] = metadata or {}
+                user_id: str | None = None
+                context = skyvern_context.current()
+                if context is not None:
+                    new_metadata["request_id"] = context.request_id
+                    new_metadata["organization_id"] = context.organization_id
+                    new_metadata["task_id"] = context.task_id
+                    new_metadata["workflow_id"] = context.workflow_id
+                    new_metadata["workflow_run_id"] = context.workflow_run_id
+                    new_metadata["task_v2_id"] = context.task_v2_id
+                    new_metadata["run_id"] = context.run_id
+                    new_metadata["organization_name"] = context.organization_name
+                    user_id = context.run_id
+
+                new_tags: list[str] = tags or []
+                new_tags.append(SettingsManager.get_settings().ENV)
+
+                return await TraceManager.__instance.traced_async(
+                    name=name, metadata=new_metadata, tags=new_tags, user_id=user_id, **trace_parameters
+                )(func)(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def traced(
+        *,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **trace_parameters: Any,
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
+        def decorator(func: Callable[P, R]) -> Callable[P, R]:
+            def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                new_metadata: dict[str, Any] = metadata or {}
+                user_id: str | None = None
+                context = skyvern_context.current()
+                if context is not None:
+                    new_metadata["request_id"] = context.request_id
+                    new_metadata["organization_id"] = context.organization_id
+                    new_metadata["task_id"] = context.task_id
+                    new_metadata["workflow_id"] = context.workflow_id
+                    new_metadata["workflow_run_id"] = context.workflow_run_id
+                    new_metadata["task_v2_id"] = context.task_v2_id
+                    new_metadata["run_id"] = context.run_id
+                    new_metadata["organization_name"] = context.organization_name
+                    user_id = context.run_id
+
+                new_tags: list[str] = tags or []
+                new_tags.append(SettingsManager.get_settings().ENV)
+
+                return TraceManager.__instance.traced(
+                    name=name, metadata=new_metadata, tags=new_tags, user_id=user_id, **trace_parameters
+                )(func)(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def get_trace_provider() -> BaseTrace:
+        return TraceManager.__instance
+
+    @staticmethod
+    def set_trace_provider(trace_provider: BaseTrace) -> None:
+        TraceManager.__instance = trace_provider

skyvern/forge/sdk/trace/base.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import Any, Awaitable, Callable, ParamSpec, TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class BaseTrace(ABC):
+    @abstractmethod
+    def traced(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
+        pass
+
+    @abstractmethod
+    def traced_async(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+        pass
+
+
+class NoOpTrace(BaseTrace):
+    def traced(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
+        return lambda func: func
+
+    def traced_async(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+        return lambda func: func

skyvern/forge/sdk/trace/lmnr.py

@@ -0,0 +1,33 @@
+from typing import Any, Awaitable, Callable, ParamSpec, TypeVar
+
+import litellm
+from lmnr import Instruments, Laminar, LaminarLiteLLMCallback, observe
+
+from skyvern.forge.sdk.trace.base import BaseTrace
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class LaminarTrace(BaseTrace):
+    def __init__(self, api_key: str) -> None:
+        Laminar.initialize(project_api_key=api_key, disabled_instruments={Instruments.SKYVERN})
+        litellm.callbacks.append(LaminarLiteLLMCallback())
+
+    def traced(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
+        return observe(name=name, ignore_output=True, metadata=metadata, tags=tags, **kwargs)
+
+    def traced_async(
+        self,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+        return observe(name=name, ignore_output=True, metadata=metadata, tags=tags, **kwargs)