34 lines
1.1 KiB
Python
34 lines
1.1 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)
|