Skip to content

Commit

Permalink
feat(sdk): client, annotations (#2452)
Browse files Browse the repository at this point in the history
  • Loading branch information
doronkopit5 authored Jan 13, 2025
1 parent 36dd9b3 commit 8b0ad84
Show file tree
Hide file tree
Showing 11 changed files with 399 additions and 36 deletions.
12 changes: 11 additions & 1 deletion packages/sample-app/sample_app/methods_decorated_app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
48 changes: 48 additions & 0 deletions packages/traceloop-sdk/tests/test_client.py
Original file line number Diff line number Diff line change
@@ -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
89 changes: 89 additions & 0 deletions packages/traceloop-sdk/tests/test_user_feedback.py
Original file line number Diff line number Diff line change
@@ -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={})
64 changes: 34 additions & 30 deletions packages/traceloop-sdk/traceloop/sdk/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
set_external_prompt_tracing_context,
)
from typing import Dict
from .client import Client


class Traceloop:
Expand All @@ -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(
Expand All @@ -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()
Expand Down Expand Up @@ -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
4 changes: 4 additions & 0 deletions packages/traceloop-sdk/traceloop/sdk/annotation/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
from .base_annotation import BaseAnnotation
from .user_feedback import UserFeedback

__all__ = ["BaseAnnotation", "UserFeedback"]
71 changes: 71 additions & 0 deletions packages/traceloop-sdk/traceloop/sdk/annotation/base_annotation.py
Original file line number Diff line number Diff line change
@@ -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,
},
},
)
43 changes: 43 additions & 0 deletions packages/traceloop-sdk/traceloop/sdk/annotation/user_feedback.py
Original file line number Diff line number Diff line change
@@ -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)
Loading

0 comments on commit 8b0ad84

Please sign in to comment.