66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
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, Instruments.PATCHRIGHT})
|
|
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)
|
|
|
|
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": "COMPLETED",
|
|
"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
|