v0.6.4 allow tagging traces (#114)

* v0.6.4 allow tagging traces

* Update src/lmnr/sdk/client/asynchronous/resources/tags.py

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* ellipsis comment

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: dinmukhamedm

pyproject.toml

@@ -6,7 +6,7 @@
 
 [project]
 name = "lmnr"
-version = "0.6.3"
+version = "0.6.4"
 description = "Python SDK for Laminar"
 authors = [
   { name = "lmnr.ai", email = "[email protected]" }

src/lmnr/sdk/client/asynchronous/async_client.py

@@ -11,6 +11,7 @@
     AsyncAgent,
     AsyncBrowserEvents,
     AsyncEvals,
+    AsyncTags,
 )
 from lmnr.sdk.utils import from_env
 
@@ -76,6 +77,7 @@ def __init__(
         self.__browser_events = AsyncBrowserEvents(
             self.__client, self.__base_url, self.__project_api_key
         )
+        self.__tags = AsyncTags(self.__client, self.__base_url, self.__project_api_key)
 
     @property
     def agent(self) -> AsyncAgent:
@@ -104,6 +106,15 @@ def _browser_events(self) -> AsyncBrowserEvents:
         """
         return self.__browser_events
 
+    @property
+    def tags(self) -> AsyncTags:
+        """Get the Tags resource.
+
+        Returns:
+            AsyncTags: The Tags resource instance.
+        """
+        return self.__tags
+
     def is_closed(self) -> bool:
         return self.__client.is_closed
 

src/lmnr/sdk/client/asynchronous/resources/__init__.py

@@ -1,9 +1,11 @@
 from lmnr.sdk.client.asynchronous.resources.agent import AsyncAgent
 from lmnr.sdk.client.asynchronous.resources.browser_events import AsyncBrowserEvents
 from lmnr.sdk.client.asynchronous.resources.evals import AsyncEvals
+from lmnr.sdk.client.asynchronous.resources.tags import AsyncTags
 
 __all__ = [
     "AsyncAgent",
     "AsyncEvals",
     "AsyncBrowserEvents",
+    "AsyncTags",
 ]

src/lmnr/sdk/client/asynchronous/resources/tags.py

@@ -0,0 +1,89 @@
+"""Resource for tagging traces."""
+
+import json
+import uuid
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+from lmnr.sdk.log import get_default_logger
+
+logger = get_default_logger(__name__)
+
+
+class AsyncTags(BaseAsyncResource):
+    """Resource for tagging traces."""
+
+    async 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, AsyncLaminarClient, observe
+
+        Laminar.initialize()
+        client = AsyncLaminarClient()
+        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()
+
+        await 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)  # Will raise ValueError if invalid
+        else:
+            raise ValueError(f"Invalid trace id: {trace_id}")
+
+        url = self._base_url + "/v1/tag"
+        payload = {
+            "traceId": trace_id,
+            "names": trace_tags,
+        }
+        response = await 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()

src/lmnr/sdk/client/synchronous/resources/__init__.py

@@ -1,5 +1,6 @@
 from lmnr.sdk.client.synchronous.resources.agent import Agent
 from lmnr.sdk.client.synchronous.resources.browser_events import BrowserEvents
 from lmnr.sdk.client.synchronous.resources.evals import Evals
+from lmnr.sdk.client.synchronous.resources.tags import Tags
 
-__all__ = ["Agent", "Evals", "BrowserEvents"]
+__all__ = ["Agent", "Evals", "BrowserEvents", "Tags"]

src/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()

src/lmnr/sdk/client/synchronous/sync_client.py

@@ -11,6 +11,7 @@
     Agent,
     BrowserEvents,
     Evals,
+    Tags,
 )
 from lmnr.sdk.utils import from_env
 
@@ -25,6 +26,7 @@ class LaminarClient:
     # Resource properties
     __agent: Agent | None = None
     __evals: Evals | None = None
+    __tags: Tags | None = None
 
     def __init__(
         self,
@@ -75,6 +77,7 @@ def __init__(
         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:
@@ -103,6 +106,15 @@ def _browser_events(self) -> BrowserEvents:
         """
         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()

src/lmnr/sdk/laminar.py

@@ -12,6 +12,7 @@
 from lmnr.opentelemetry_lib.decorators import json_dumps
 from opentelemetry import context as context_api, trace
 from opentelemetry.context import attach, detach
+from opentelemetry.trace import INVALID_TRACE_ID
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.util.types import AttributeValue
 
@@ -703,6 +704,20 @@ def get_base_http_url(cls):
     def get_project_api_key(cls):
         return cls.__project_api_key
 
+    @classmethod
+    def get_trace_id(cls) -> uuid.UUID | None:
+        """Get the trace id for the current active span represented as a UUID.
+        Returns None if there is no active span.
+
+        Returns:
+            uuid.UUID | None: The trace id for the current span, or None if\
+            there is no active span.
+        """
+        trace_id = trace.get_current_span().get_span_context().trace_id
+        if trace_id == INVALID_TRACE_ID:
+            return None
+        return uuid.UUID(int=trace_id)
+
     @classmethod
     def _headers(cls):
         assert cls.__project_api_key is not None, "Project API key is not set"

src/lmnr/version.py

@@ -3,7 +3,7 @@
 from packaging import version
 
 
-__version__ = "0.6.3"
+__version__ = "0.6.4"
 PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"