lmnr 0.6.2__py3-none-any.whl → 0.6.4__py3-none-any.whl

lmnr/sdk/client/synchronous/resources/agent.py

@@ -1,6 +1,6 @@
 """Agent resource for interacting with Laminar agents."""
 
-from typing import Generator, Literal, Optional, Union, overload
+from typing import Generator, Literal, overload
 import uuid
 
 from lmnr.sdk.client.synchronous.resources.base import BaseResource
@@ -23,44 +23,44 @@ class Agent(BaseResource):
         self,
         prompt: str,
         stream: Literal[True],
-        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
-        model_provider: Optional[ModelProvider] = None,
-        model: Optional[str] = None,
+        parent_span_context: LaminarSpanContext | str | None = None,
+        model_provider: ModelProvider | None = None,
+        model: str | None = None,
         enable_thinking: bool = True,
-        agent_state: Optional[str] = None,
-        storage_state: Optional[str] = None,
+        agent_state: str | None = None,
+        storage_state: str | None = None,
         return_screenshots: bool = False,
         return_agent_state: bool = False,
         return_storage_state: bool = False,
         disable_give_control: bool = False,
-        timeout: Optional[int] = None,
-        cdp_url: Optional[str] = None,
-        max_steps: Optional[int] = None,
-        thinking_token_budget: Optional[int] = None,
-        start_url: Optional[str] = None,
-        user_agent: Optional[str] = None,
+        timeout: int | None = None,
+        cdp_url: str | None = None,
+        max_steps: int | None = None,
+        thinking_token_budget: int | None = None,
+        start_url: str | None = None,
+        user_agent: str | None = None,
     ) -> Generator[RunAgentResponseChunk, None, None]:
         """Run Laminar index agent in streaming mode.
 
         Args:
             prompt (str): prompt for the agent
             stream (Literal[True]): whether to stream the agent's response
-            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
-            model_provider (Optional[ModelProvider], optional): LLM model provider
-            model (Optional[str], optional): LLM model name
+            parent_span_context (LaminarSpanContext | str | None, optional): span context if the agent is part of a trace
+            model_provider (ModelProvider | None, optional): LLM model provider
+            model (str | None, optional): LLM model name
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
-            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
-            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
+            agent_state (str | None, optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (str | None, optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
             return_agent_state (bool, optional): whether to return the agent's state. Default to False.
             return_storage_state (bool, optional): whether to return the storage state. Default to False.
             disable_give_control (bool, optional): whether to NOT give the agent additional direction to give control to the user for tasks such as login. Default to False.
-            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
-            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
-            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
-            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
-            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
-            user_agent (Optional[str], optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
+            timeout (int | None, optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (str | None, optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (int | None, optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (int | None, optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (str | None, optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+            user_agent (str | None, optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
         Returns:
             Generator[RunAgentResponseChunk, None, None]: a generator of response chunks
         """
@@ -70,43 +70,43 @@ class Agent(BaseResource):
     def run(
         self,
         prompt: str,
-        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
-        model_provider: Optional[ModelProvider] = None,
-        model: Optional[str] = None,
+        parent_span_context: LaminarSpanContext | str | None = None,
+        model_provider: ModelProvider | None = None,
+        model: str | None = None,
         enable_thinking: bool = True,
-        agent_state: Optional[str] = None,
-        storage_state: Optional[str] = None,
+        agent_state: str | None = None,
+        storage_state: str | None = None,
         return_screenshots: bool = False,
         return_agent_state: bool = False,
         disable_give_control: bool = False,
         return_storage_state: bool = False,
-        timeout: Optional[int] = None,
-        cdp_url: Optional[str] = None,
-        max_steps: Optional[int] = None,
-        thinking_token_budget: Optional[int] = None,
-        start_url: Optional[str] = None,
-        user_agent: Optional[str] = None,
+        timeout: int | None = None,
+        cdp_url: str | None = None,
+        max_steps: int | None = None,
+        thinking_token_budget: int | None = None,
+        start_url: str | None = None,
+        user_agent: str | None = None,
     ) -> AgentOutput:
         """Run Laminar index agent.
 
         Args:
             prompt (str): prompt for the agent
-            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
-            model_provider (Optional[ModelProvider], optional): LLM model provider
-            model (Optional[str], optional): LLM model name
+            parent_span_context (LaminarSpanContext | str | None, optional): span context if the agent is part of a trace
+            model_provider (ModelProvider | None, optional): LLM model provider
+            model (str | None, optional): LLM model name
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
-            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
-            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
+            agent_state (str | None, optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (str | None, optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
             return_agent_state (bool, optional): whether to return the agent's state. Default to False.
             return_storage_state (bool, optional): whether to return the storage state. Default to False.
             disable_give_control (bool, optional): whether to NOT give the agent additional direction to give control to the user for tasks such as login. Default to False.
-            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
-            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
-            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
-            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
-            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
-            user_agent (Optional[str], optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
+            timeout (int | None, optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (str | None, optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (int | None, optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (int | None, optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (str | None, optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+            user_agent (str | None, optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
 
         Returns:
             AgentOutput: agent output
@@ -117,45 +117,45 @@ class Agent(BaseResource):
     def run(
         self,
         prompt: str,
-        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
-        model_provider: Optional[ModelProvider] = None,
-        model: Optional[str] = None,
+        parent_span_context: LaminarSpanContext | str | None = None,
+        model_provider: ModelProvider | None = None,
+        model: str | None = None,
         stream: Literal[False] = False,
         enable_thinking: bool = True,
-        agent_state: Optional[str] = None,
-        storage_state: Optional[str] = None,
+        agent_state: str | None = None,
+        storage_state: str | None = None,
         return_screenshots: bool = False,
         return_agent_state: bool = False,
         return_storage_state: bool = False,
         disable_give_control: bool = False,
-        timeout: Optional[int] = None,
-        cdp_url: Optional[str] = None,
-        max_steps: Optional[int] = None,
-        thinking_token_budget: Optional[int] = None,
-        start_url: Optional[str] = None,
-        user_agent: Optional[str] = None,
+        timeout: int | None = None,
+        cdp_url: str | None = None,
+        max_steps: int | None = None,
+        thinking_token_budget: int | None = None,
+        start_url: str | None = None,
+        user_agent: str | None = None,
     ) -> AgentOutput:
         """Run Laminar index agent.
 
         Args:
             prompt (str): prompt for the agent
-            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
-            model_provider (Optional[ModelProvider], optional): LLM model provider
-            model (Optional[str], optional): LLM model name
+            parent_span_context (LaminarSpanContext | str | None, optional): span context if the agent is part of a trace
+            model_provider (ModelProvider | None, optional): LLM model provider
+            model (str | None, optional): LLM model name
             stream (Literal[False], optional): whether to stream the agent's response
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
-            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
-            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
+            agent_state (str | None, optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (str | None, optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
             return_agent_state (bool, optional): whether to return the agent's state. Default to False.
             return_storage_state (bool, optional): whether to return the storage state. Default to False.
             disable_give_control (bool, optional): whether to NOT give the agent additional direction to give control to the user for tasks such as login. Default to False.
-            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
-            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
-            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
-            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
-            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
-            user_agent (Optional[str], optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
+            timeout (int | None, optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (str | None, optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (int | None, optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (int | None, optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (str | None, optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+            user_agent (str | None, optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
 
         Returns:
             AgentOutput: agent output
@@ -165,48 +165,48 @@ class Agent(BaseResource):
     def run(
         self,
         prompt: str,
-        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
-        model_provider: Optional[ModelProvider] = None,
-        model: Optional[str] = None,
+        parent_span_context: LaminarSpanContext | str | None = None,
+        model_provider: ModelProvider | None = None,
+        model: str | None = None,
         stream: bool = False,
         enable_thinking: bool = True,
-        agent_state: Optional[str] = None,
-        storage_state: Optional[str] = None,
+        agent_state: str | None = None,
+        storage_state: str | None = None,
         return_screenshots: bool = False,
         return_agent_state: bool = False,
         return_storage_state: bool = False,
         disable_give_control: bool = False,
-        timeout: Optional[int] = None,
-        cdp_url: Optional[str] = None,
-        max_steps: Optional[int] = None,
-        thinking_token_budget: Optional[int] = None,
-        start_url: Optional[str] = None,
-        user_agent: Optional[str] = None,
-    ) -> Union[AgentOutput, Generator[RunAgentResponseChunk, None, None]]:
+        timeout: int | None = None,
+        cdp_url: str | None = None,
+        max_steps: int | None = None,
+        thinking_token_budget: int | None = None,
+        start_url: str | None = None,
+        user_agent: str | None = None,
+    ) -> AgentOutput | Generator[RunAgentResponseChunk, None, None]:
         """Run Laminar index agent.
 
         Args:
             prompt (str): prompt for the agent
-            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
-            model_provider (Optional[ModelProvider], optional): LLM model provider
-            model (Optional[str], optional): LLM model name
+            parent_span_context (LaminarSpanContext | str | None, optional): span context if the agent is part of a trace
+            model_provider (ModelProvider | None, optional): LLM model provider
+            model (str | None, optional): LLM model name
             stream (bool, optional): whether to stream the agent's response
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
-            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
-            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
+            agent_state (str | None, optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (str | None, optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
             return_agent_state (bool, optional): whether to return the agent's state. Default to False.
             return_storage_state (bool, optional): whether to return the storage state. Default to False.
             disable_give_control (bool, optional): whether to NOT give the agent additional direction to give control to the user for tasks such as login. Default to False.
-            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
-            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
-            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
-            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
-            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
-            user_agent (Optional[str], optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
+            timeout (int | None, optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (str | None, optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (int | None, optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (int | None, optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (str | None, optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+            user_agent (str | None, optional): the user to be sent to the browser. If not specified, Laminar uses the default user agent. Default to None.
 
         Returns:
-            Union[AgentOutput, Generator[RunAgentResponseChunk, None, None]]: agent output or a generator of response chunks
+            AgentOutput | Generator[RunAgentResponseChunk, None, None]: agent output or a generator of response chunks
         """
         if parent_span_context is None:
             span = trace.get_current_span()

lmnr/sdk/client/synchronous/resources/evals.py

@@ -2,12 +2,12 @@
 
 import uuid
 import urllib.parse
-from typing import Optional
 
 from lmnr.sdk.client.synchronous.resources.base import BaseResource
 from lmnr.sdk.types import (
     InitEvaluationResponse,
     EvaluationResultDatapoint,
+    PartialEvaluationDatapoint,
     GetDatapointsResponse,
 )
 
@@ -16,13 +16,13 @@ class Evals(BaseResource):
     """Resource for interacting with Laminar evaluations API."""
 
     def init(
-        self, name: Optional[str] = None, group_name: Optional[str] = None
+        self, name: str | None = None, group_name: str | None = None
     ) -> InitEvaluationResponse:
         """Initialize a new evaluation.
 
         Args:
-            name (Optional[str], optional): Name of the evaluation. Defaults to None.
-            group_name (Optional[str], optional): Group name for the evaluation. Defaults to None.
+            name (str | None, optional): Name of the evaluation. Defaults to None.
+            group_name (str | None, optional): Group name for the evaluation. Defaults to None.
 
         Returns:
             InitEvaluationResponse: The response from the initialization request.
@@ -41,15 +41,15 @@ class Evals(BaseResource):
     def save_datapoints(
         self,
         eval_id: uuid.UUID,
-        datapoints: list[EvaluationResultDatapoint],
-        group_name: Optional[str] = None,
+        datapoints: list[EvaluationResultDatapoint | PartialEvaluationDatapoint],
+        group_name: str | None = None,
     ):
         """Save evaluation datapoints.
 
         Args:
             eval_id (uuid.UUID): The evaluation ID.
-            datapoints (list[EvaluationResultDatapoint]): The datapoints to save.
-            group_name (Optional[str], optional): Group name for the datapoints. Defaults to None.
+            datapoints (list[EvaluationResultDatapoint | PartialEvaluationDatapoint]): The datapoints to save.
+            group_name (str | None, optional): Group name for the datapoints. Defaults to None.
 
         Raises:
             ValueError: If there's an error saving the datapoints.

lmnr/sdk/client/synchronous/resources/tags.py

@@ -0,0 +1,89 @@
+"""Resource for tagging traces."""
+
+import json
+import uuid
+
+from lmnr.sdk.client.synchronous.resources.base import BaseResource
+from lmnr.sdk.log import get_default_logger
+
+logger = get_default_logger(__name__)
+
+
+class Tags(BaseResource):
+    """Resource for tagging traces."""
+
+    def tag(
+        self,
+        trace_id: str | int | uuid.UUID,
+        tags: list[str] | str,
+    ):
+        """Tag a trace with a list of tags. Note that the trace must be ended
+        before tagging it. You may want to call `Laminar.flush()` after the
+        trace that you want to tag.
+
+        Args:
+            trace_id (str | int | uuid.UUID): The trace id to tag.
+            tags (list[str] | str): The tag or list of tags to add to the trace.
+
+        Raises:
+            ValueError: If the trace id is not a valid UUID.
+
+        Returns:
+            list[dict]: The response from the server.
+
+        Example:
+        ```python
+        from lmnr import Laminar, LaminarClient, observe
+
+        Laminar.initialize()
+        client = LaminarClient()
+        trace_id = None
+
+        @observe()
+        def foo():
+            trace_id = Laminar.get_trace_id()
+            pass
+
+        # make sure `foo` is called outside a trace context
+        foo()
+
+        # or make sure the trace is ended by this point
+        Laminar.flush()
+
+        client.tags.tag(trace_id, "my_tag")
+        ```
+        """
+        trace_tags = tags if isinstance(tags, list) else [tags]
+        if isinstance(trace_id, uuid.UUID):
+            trace_id = str(trace_id)
+        elif isinstance(trace_id, int):
+            trace_id = str(uuid.UUID(int=trace_id))
+        elif isinstance(trace_id, str):
+            uuid.UUID(trace_id)
+        else:
+            raise ValueError(f"Invalid trace id: {trace_id}")
+
+        url = self._base_url + "/v1/tag"
+        payload = {
+            "traceId": trace_id,
+            "names": trace_tags,
+        }
+        response = self._client.post(
+            url,
+            content=json.dumps(payload),
+            headers={
+                **self._headers(),
+            },
+        )
+
+        if response.status_code == 404:
+            logger.warning(
+                f"Trace {trace_id} not found. The trace may have not been ended yet."
+            )
+            return []
+
+        if response.status_code != 200:
+            raise ValueError(
+                f"Failed to tag trace: [{response.status_code}] {response.text}"
+            )
+        return response.json()

lmnr/sdk/client/synchronous/sync_client.py

@@ -4,13 +4,14 @@ Laminar HTTP client. Used to send data to/from the Laminar API.
 
 import httpx
 import re
-from typing import Optional, TypeVar
+from typing import TypeVar
 from types import TracebackType
 
 from lmnr.sdk.client.synchronous.resources import (
     Agent,
     BrowserEvents,
     Evals,
+    Tags,
 )
 from lmnr.sdk.utils import from_env
 
@@ -23,22 +24,26 @@ class LaminarClient:
     __client: httpx.Client = None
 
     # Resource properties
-    __agent: Optional[Agent] = None
-    __evals: Optional[Evals] = None
+    __agent: Agent | None = None
+    __evals: Evals | None = None
+    __tags: Tags | None = None
 
     def __init__(
         self,
-        base_url: Optional[str] = None,
-        project_api_key: Optional[str] = None,
-        port: Optional[int] = None,
+        base_url: str | None = None,
+        project_api_key: str | None = None,
+        port: int | None = None,
         timeout: int = 3600,
     ):
         """Initializer for the Laminar HTTP client.
 
         Args:
-            base_url (str): base URL of the Laminar API.
-            project_api_key (str): Laminar project API key
-            port (Optional[int], optional): port of the Laminar API HTTP server.\
+            base_url (str | None): base URL of the Laminar API. If not\
+                provided, the LMNR_BASE_URL environment variable is used or we\
+                default to "https://api.lmnr.ai".
+            project_api_key (str | None): Laminar project API key. If not\
+                provided, the LMNR_PROJECT_API_KEY environment variable is used.
+            port (int | None, optional): port of the Laminar API HTTP server.\
                 Overrides any port in the base URL.
                 Defaults to None. If none is provided, the default port (443) will
                 be used.
@@ -72,6 +77,7 @@ class LaminarClient:
         self.__browser_events = BrowserEvents(
             self.__client, self.__base_url, self.__project_api_key
         )
+        self.__tags = Tags(self.__client, self.__base_url, self.__project_api_key)
 
     @property
     def agent(self) -> Agent:
@@ -100,6 +106,15 @@ class LaminarClient:
         """
         return self.__browser_events
 
+    @property
+    def tags(self) -> Tags:
+        """Get the Tags resource.
+
+        Returns:
+            Tags: The Tags resource instance.
+        """
+        return self.__tags
+
     def shutdown(self):
         """Shutdown the client by closing underlying connections."""
         self.__client.close()
@@ -119,7 +134,7 @@ class LaminarClient:
         """
         # If an error is thrown while constructing a client, self._client
         # may not be present
-        if hasattr(self, "_client"):
+        if hasattr(self, "__client"):
             self.__client.close()
 
     def __enter__(self: _T) -> _T:
@@ -127,9 +142,9 @@ class LaminarClient:
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         self.close()
 

lmnr/sdk/decorators.py

@@ -1,16 +1,14 @@
 from lmnr.opentelemetry_lib.decorators import (
     entity_method,
     aentity_method,
+    json_dumps,
 )
 from opentelemetry.trace import INVALID_SPAN, get_current_span
 
-from typing import Callable, Literal, Optional, TypeVar, Union, cast
+from typing import Any, Callable, Literal, TypeVar, cast
 from typing_extensions import ParamSpec
 
 from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID
-from lmnr.opentelemetry_lib.tracing.context_properties import (
-    update_association_properties,
-)
 
 from .utils import is_async
 
@@ -21,37 +19,42 @@ R = TypeVar("R")
 
 def observe(
     *,
-    name: Optional[str] = None,
-    session_id: Optional[str] = None,
+    name: str | None = None,
+    session_id: str | None = None,
+    user_id: str | None = None,
     ignore_input: bool = False,
     ignore_output: bool = False,
-    span_type: Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]] = "DEFAULT",
-    ignore_inputs: Optional[list[str]] = None,
+    span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT",
+    ignore_inputs: list[str] | None = None,
+    metadata: dict[str, Any] | None = None,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """The main decorator entrypoint for Laminar. This is used to wrap
     functions and methods to create spans.
 
     Args:
-        name (Optional[str], optional): Name of the span. Function
-                        name is used if not specified.
-                        Defaults to None.
-        session_id (Optional[str], optional): Session ID to associate with the\
-                        span and the following context. Defaults to None.
+        name (str | None, optional): Name of the span. Function name is used if\
+            not specified. Defaults to None.
+        session_id (str | None, optional): Session ID to associate with the\
+            span and the following context. Defaults to None.
+        user_id (str | None, optional): User ID to associate with the span and\
+            the following context. This is different from ID of a Laminar user.
+            Defaults to None.
         ignore_input (bool, optional): Whether to ignore ALL input of the\
-                        wrapped function. Defaults to False.
+            wrapped function. Defaults to False.
         ignore_output (bool, optional): Whether to ignore ALL output of the\
-                        wrapped function. Defaults to False.
-        span_type (Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]], optional): Type of the span.
-                        Defaults to "DEFAULT".
-        ignore_inputs (Optional[list[str]], optional): List of input keys to ignore.
-                        For example, if the wrapped function takes three arguments,\
-                        def foo(a, b, `sensitive_data`), and you want to ignore the\
-                        `sensitive_data` argument, you can pass ["sensitive_data"] to\
-                        this argument.
-                        Defaults to None.
+            wrapped function. Defaults to False.
+        span_type (Literal["DEFAULT", "LLM", "TOOL"], optional): Type of the span.
+            Defaults to "DEFAULT".
+        ignore_inputs (list[str] | None, optional): List of input keys to\
+            ignore. For example, if the wrapped function takes three arguments\
+            def foo(a, b, `sensitive_data`), and you want to ignore the\
+            `sensitive_data` argument, you can pass ["sensitive_data"] to\
+            this argument. Defaults to None.
+        metadata (dict[str, Any] | None, optional): Metadata to associate with\
+            the trace. Must be JSON serializable. Defaults to None.
     Raises:
-        Exception: re-raises the exception if the wrapped function raises
-                   an exception
+        Exception: re-raises the exception if the wrapped function raises an\
+            exception
 
     Returns:
         R: Returns the result of the wrapped function
@@ -65,14 +68,25 @@ def observe(
         association_properties = {}
         if session_id is not None:
             association_properties["session_id"] = session_id
-        update_association_properties(association_properties)
-        return (
+        if user_id is not None:
+            association_properties["user_id"] = user_id
+        if metadata is not None:
+            association_properties.update(
+                {
+                    f"metadata.{k}": (
+                        v if isinstance(v, (str, int, float, bool)) else json_dumps(v)
+                    )
+                    for k, v in metadata.items()
+                }
+            )
+        result = (
             aentity_method(
                 name=name,
                 ignore_input=ignore_input,
                 ignore_output=ignore_output,
                 span_type=span_type,
                 ignore_inputs=ignore_inputs,
+                association_properties=association_properties,
             )(func)
             if is_async(func)
             else entity_method(
@@ -81,7 +95,9 @@ def observe(
                 ignore_output=ignore_output,
                 span_type=span_type,
                 ignore_inputs=ignore_inputs,
+                association_properties=association_properties,
             )(func)
         )
+        return result
 
     return cast(Callable, decorator)