feat(sdk): client, annotations (#2452)

Author: doronkopit5

packages/sample-app/sample_app/methods_decorated_app.py

@@ -68,7 +68,17 @@ def joke_workflow():
     eng_joke = create_joke()
     pirate_joke = translate_joke_to_pirate(eng_joke)
     signature = generate_signature(pirate_joke)
-    print(pirate_joke + "\n\n" + signature)
+
+    traceloop_client = Traceloop.get()
+    traceloop_client.user_feedback.create(
+        "sample-annotation-task",
+        "user_12345",
+        {"sentiment": "positive", "score": 0.95, "tones": ["happy", "surprised"]},
+    )
+
+    result = pirate_joke + "\n\n" + signature
+    print(result)
+    return result
 
 
 joke_workflow()

packages/traceloop-sdk/tests/test_client.py

@@ -0,0 +1,48 @@
+import pytest
+from traceloop.sdk.client import Client
+from traceloop.sdk.client.http import HTTPClient
+from traceloop.sdk.annotation.user_feedback import UserFeedback
+
+
+def test_client_initialization():
+    """Test basic client initialization"""
+    client = Client(api_key="test-key", app_name="test-app")
+
+    assert client.app_name == "test-app"
+    assert client.api_key == "test-key"
+    assert client.api_endpoint == "https://api.traceloop.com"
+    assert isinstance(client._http, HTTPClient)
+
+
+def test_client_custom_endpoint():
+    """Test client initialization with custom endpoint"""
+    client = Client(api_key="test-key", app_name="test-app", api_endpoint="https://custom.endpoint.com")
+
+    assert client.api_endpoint == "https://custom.endpoint.com"
+    assert client._http.base_url == "https://custom.endpoint.com"
+
+
+def test_client_default_app_name():
+    """Test client initialization with default app_name"""
+    client = Client(api_key="test-key")
+
+    # Default app_name should be sys.argv[0]
+    import sys
+
+    assert client.app_name == sys.argv[0]
+
+
+@pytest.mark.parametrize("api_key", [None, "", " "])
+def test_client_requires_api_key(api_key):
+    """Test that client requires a valid API key"""
+    with pytest.raises(ValueError, match="API key is required"):
+        Client(api_key=api_key)
+
+
+def test_user_feedback_initialization():
+    """Test user_feedback is properly initialized"""
+    client = Client(api_key="test-key", app_name="test-app")
+
+    assert isinstance(client.user_feedback, UserFeedback)
+    assert client.user_feedback._http == client._http
+    assert client.user_feedback._app_name == client.app_name

packages/traceloop-sdk/tests/test_user_feedback.py

@@ -0,0 +1,89 @@
+"""
+Tests for the UserFeedback class.
+
+These tests verify:
+1. Proper initialization of UserFeedback instances
+2. Basic feedback submission with minimal parameters
+3. Handling of complex tag structures
+4. Proper API endpoint construction and payload formatting
+"""
+
+import pytest
+from unittest.mock import Mock
+from traceloop.sdk.annotation.user_feedback import UserFeedback
+from traceloop.sdk.client.http import HTTPClient
+
+
+@pytest.fixture
+def mock_http():
+    """Create a mock HTTP client"""
+    http = Mock(spec=HTTPClient)
+    http.post.return_value = {"status": "success"}
+    return http
+
+
+@pytest.fixture
+def user_feedback(mock_http):
+    """Create a UserFeedback instance with mock HTTP client"""
+    return UserFeedback(mock_http, "test-app")
+
+
+def test_user_feedback_initialization(mock_http):
+    """Test UserFeedback is properly initialized"""
+    feedback = UserFeedback(mock_http, "test-app")
+    assert feedback._http == mock_http
+    assert feedback._app_name == "test-app"
+
+
+def test_create_basic_feedback(user_feedback, mock_http):
+    """Test creating basic user feedback"""
+    user_feedback.create(
+        annotation_task="task_123", entity_id="instance_456", tags={"sentiment": "positive"}
+    )
+
+    mock_http.post.assert_called_once_with(
+        "annotation-tasks/task_123/annotations",
+        {
+            "entity_instance_id": "instance_456",
+            "tags": {"sentiment": "positive"},
+            "source": "sdk",
+            "flow": "user_feedback",
+            "actor": {
+                "type": "service",
+                "id": "test-app",
+            },
+        },
+    )
+
+
+def test_create_feedback_complex_tags(user_feedback, mock_http):
+    """Test creating user feedback with complex tags"""
+    tags = {"sentiment": "positive", "relevance": 0.95, "tones": ["happy", "nice"]}
+
+    user_feedback.create(annotation_task="task_123", entity_id="instance_456", tags=tags)
+
+    mock_http.post.assert_called_once_with(
+        "annotation-tasks/task_123/annotations",
+        {
+            "entity_instance_id": "instance_456",
+            "tags": tags,
+            "source": "sdk",
+            "flow": "user_feedback",
+            "actor": {
+                "type": "service",
+                "id": "test-app",
+            },
+        },
+    )
+
+
+def test_create_feedback_parameter_validation(user_feedback):
+    """Test parameter validation for feedback creation"""
+    with pytest.raises(ValueError, match="annotation_task is required"):
+        user_feedback.create(annotation_task="", entity_id="instance_456", tags={"sentiment": "positive"})
+
+    with pytest.raises(ValueError, match="entity_id is required"):
+        user_feedback.create(annotation_task="task_123", entity_id="", tags={"sentiment": "positive"})
+
+    with pytest.raises(ValueError, match="tags cannot be empty"):
+        user_feedback.create(annotation_task="task_123", entity_id="instance_456", tags={})

packages/traceloop-sdk/traceloop/sdk/__init__.py

@@ -30,6 +30,7 @@
     set_external_prompt_tracing_context,
 )
 from typing import Dict
+from .client import Client
 
 
 class Traceloop:
@@ -39,31 +40,33 @@ class Traceloop:
     AUTO_CREATED_URL = str(Path.home() / ".cache" / "traceloop" / "auto_created_url")
 
     __tracer_wrapper: TracerWrapper
-    __fetcher: Fetcher = None
+    __fetcher: Optional[Fetcher] = None
+    __app_name: Optional[str] = None
+    __client: Optional[Client] = None
 
     @staticmethod
     def init(
-        app_name: Optional[str] = sys.argv[0],
+        app_name: str = sys.argv[0],
         api_endpoint: str = "https://api.traceloop.com",
-        api_key: str = None,
+        api_key: Optional[str] = None,
         enabled: bool = True,
         headers: Dict[str, str] = {},
         disable_batch=False,
         telemetry_enabled: bool = True,
-        exporter: SpanExporter = None,
+        exporter: Optional[SpanExporter] = None,
         metrics_exporter: MetricExporter = None,
         metrics_headers: Dict[str, str] = None,
         logging_exporter: LogExporter = None,
         logging_headers: Dict[str, str] = None,
-        processor: SpanProcessor = None,
+        processor: Optional[SpanProcessor] = None,
         propagator: TextMapPropagator = None,
         traceloop_sync_enabled: bool = False,
         should_enrich_metrics: bool = True,
         resource_attributes: dict = {},
         instruments: Optional[Set[Instruments]] = None,
         block_instruments: Optional[Set[Instruments]] = None,
         image_uploader: Optional[ImageUploader] = None,
-    ) -> None:
+    ) -> Optional[Client]:
         if not enabled:
             TracerWrapper.set_disabled(True)
             print(
@@ -82,13 +85,14 @@ def init(
 
         api_endpoint = os.getenv("TRACELOOP_BASE_URL") or api_endpoint
         api_key = os.getenv("TRACELOOP_API_KEY") or api_key
+        Traceloop.__app_name = app_name
 
         if (
             traceloop_sync_enabled
             and api_endpoint.find("traceloop.com") != -1
             and api_key
-            and not exporter
-            and not processor
+            and (exporter is None)
+            and (processor is None)
         ):
             Traceloop.__fetcher = Fetcher(base_url=api_endpoint, api_key=api_key)
             Traceloop.__fetcher.run()
@@ -186,32 +190,32 @@ def init(
             )
             Traceloop.__logger_wrapper = LoggerWrapper(exporter=logging_exporter)
 
+        if not api_key:
+            return
+        Traceloop.__client = Client(api_key=api_key, app_name=app_name, api_endpoint=api_endpoint)
+        return Traceloop.__client
+
     def set_association_properties(properties: dict) -> None:
         set_association_properties(properties)
 
     def set_prompt(template: str, variables: dict, version: int):
         set_external_prompt_tracing_context(template, variables, version)
 
-    def report_score(
-        association_property_name: str,
-        association_property_id: str,
-        score: float,
-    ):
-        if not Traceloop.__fetcher:
-            print(
-                Fore.RED
-                + "Error: Cannot report score. Missing Traceloop API key,"
-                + " go to https://app.traceloop.com/settings/api-keys to create one"
+    @staticmethod
+    def get():
+        """
+        Returns the shared SDK client instance, using the current global configuration.
+
+        To use the SDK as a singleton, first make sure you have called :func:`Traceloop.init()`
+        at startup time. Then ``get()`` will return the same shared :class:`Traceloop.client.Client`
+        instance each time. The client will be initialized if it has not been already.
+
+        If you need to create multiple client instances with different configurations, instead of this
+        singleton approach you can call the :class:`Traceloop.client.Client` constructor directly instead.
+        """
+        if not Traceloop.__client:
+            raise Exception(
+                "Client not initialized, you should call Traceloop.init() first. "
+                "If you are still getting this error - you are missing the api key"
             )
-            print("Set the TRACELOOP_API_KEY environment variable to the key")
-            print(Fore.RESET)
-            return
-
-        Traceloop.__fetcher.post(
-            "score",
-            {
-                "entity_name": f"traceloop.association.properties.{association_property_name}",
-                "entity_id": association_property_id,
-                "score": score,
-            },
-        )
+        return Traceloop.__client

packages/traceloop-sdk/traceloop/sdk/annotation/__init__.py

@@ -0,0 +1,4 @@
+from .base_annotation import BaseAnnotation
+from .user_feedback import UserFeedback
+
+__all__ = ["BaseAnnotation", "UserFeedback"]

packages/traceloop-sdk/traceloop/sdk/annotation/base_annotation.py

@@ -0,0 +1,71 @@
+from typing import Dict, Any
+
+from ..client.http import HTTPClient
+
+
+class BaseAnnotation:
+    """
+    Annotation class for creating annotations in Traceloop.
+
+    This class provides functionality to create annotations for specific tasks.
+    """
+
+    _http: HTTPClient
+    _app_name: str
+
+    def __init__(self, http: HTTPClient, app_name: str, flow: str):
+        self._http = http
+        self._app_name = app_name
+        self._flow = flow
+
+    def create(
+        self,
+        annotation_task: str,
+        entity_id: str,
+        tags: Dict[str, Any],
+    ) -> None:
+        """Create an user feedback annotation for a specific task.
+
+        Args:
+            annotation_task (str): The ID/slug of the annotation task to report to.
+                Can be found at app.traceloop.com/annotation_tasks/:annotation_task_id
+            entity_id (str): The ID of the specific entity instance being annotated, should be reported
+                in the association properties
+            tags (Dict[str, Any]): Dictionary containing the tags to be reported.
+                Should match the tags defined in the annotation task
+
+        Example:
+            ```python
+            client = Client(api_key="your-key")
+            client.annotation.create(
+                annotation_task="task_123",
+                entity_id="instance_456",
+                tags={
+                    "sentiment": "positive",
+                    "relevance": 0.95,
+                    "tones": ["happy", "nice"]
+                },
+            )
+            ```
+        """
+
+        if not annotation_task:
+            raise ValueError("annotation_task is required")
+        if not entity_id:
+            raise ValueError("entity_id is required")
+        if not tags:
+            raise ValueError("tags cannot be empty")
+
+        self._http.post(
+            f"annotation-tasks/{annotation_task}/annotations",
+            {
+                "entity_instance_id": entity_id,
+                "tags": tags,
+                "source": "sdk",
+                "flow": self._flow,
+                "actor": {
+                    "type": "service",
+                    "id": self._app_name,
+                },
+            },
+        )

packages/traceloop-sdk/traceloop/sdk/annotation/user_feedback.py

@@ -0,0 +1,43 @@
+from typing import Any, Dict
+
+from traceloop.sdk.client.http import HTTPClient
+from .base_annotation import BaseAnnotation
+
+
+class UserFeedback(BaseAnnotation):
+    def __init__(self, http: HTTPClient, app_name: str):
+        super().__init__(http, app_name, "user_feedback")
+
+
+def create(
+    self,
+    annotation_task: str,
+    entity_instance_id: str,
+    tags: Dict[str, Any],
+) -> None:
+    """Create an annotation for a specific task.
+
+    Args:
+        annotation_task (str): The ID/slug of the annotation task to report to.
+            Can be found at app.traceloop.com/annotation_tasks/:annotation_task_id
+        entity_instance_id (str): The ID of the specific entity instance being annotated, should be reported
+            in the association properties
+        tags (Dict[str, Any]): Dictionary containing the tags to be reported.
+            Should match the tags defined in the annotation task
+
+    Example:
+        ```python
+        client = Client(api_key="your-key")
+        client.annotation.create(
+            annotation_task="task_123",
+            entity_instance_id="instance_456",
+            tags={
+                "sentiment": "positive",
+                "relevance": 0.95,
+                "tones": ["happy", "nice"]
+            },
+        )
+        ```
+    """
+
+    return BaseAnnotation.create(self, annotation_task, entity_instance_id, tags)