fix: input/output for workflows & tasks

traceloop · Feb 13, 2024 · be771a3 · be771a3
1 parent 301dc79
commit be771a3
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 28 deletions.
diff --git a/packages/traceloop-sdk/tests/cassettes/test_workflows/test_simple_workflow.yaml b/packages/traceloop-sdk/tests/cassettes/test_workflows/test_simple_workflow.yaml
@@ -1,6 +1,6 @@
 interactions:
 - request:
-    body: '{"messages": [{"role": "user", "content": "Tell me a joke about opentelemetry"}],
+    body: '{"messages": [{"role": "user", "content": "Tell me a joke about OpenTelemetry"}],
       "model": "gpt-3.5-turbo"}'
     headers:
       accept:
@@ -36,18 +36,18 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAA1SRT08CMRDF7/spxp5ZsgsBlIvRmwf/HEgwEUNKd9ittp2mnVUJ4bubLgvopYf3
-        +l5/M91nAEJXYg5CNZKV9Sa/Dov4opZfZmSL7d1DOa6nr1UsVPP4tGjFICVo84GKT6mhIusNsiZ3
-        tFVAyZhay1kxm02mN6OyMyxVaFKs9pyPh5Oc27ChvJiW4z7ZkFYYxRzeMgCAfXcmRlfhj5hDMTgp
-        FmOUNYr5+RKACGSSImSMOrJ0LAYXU5FjdB32stlBpSt49ugYDVrksINGBwQJiixWWrrblVu5e1Sy
-        jQjc4A6+pWOsYIPMGICDVBiBtmBkWzeM4Ur0rx3OmIZqH2iTRnKtMWd9q52OzTqgjOQSUmTyx/gh
-        A3jv1tH+m1D4QNbzmukTXSosJ8c6cfmAP+asN5lYmos+HmU9n4i7yGjXW+1qDD7objeJMjtkvwAA
-        AP//AwB7h+a6GgIAAA==
+        H4sIAAAAAAAAA1SRT08CMRDF7/spxp6B7AorKxcTD16I4WLUKIaU7rBb6XaadmBFwnc3Xf7FSw/v
+        zZv83nSfAAhdigkIVUtWjTP9wrvq6Xfa5uOpbovtNv14neXrIntevrtW9GKClt+o+JwaKGqcQdZk
+        j7byKBnj1mycjotROh7mndFQiSbGKsf94SDv88YvqZ/eZcNTsiatMIgJfCYAAPvujYy2xB8xgbR3
+        VhoMQVYoJpchAOHJREXIEHRgaVn0rqYiy2g77Ld6B60MwDXCzKF9QYMNst+B8xS7gbZbzVgCE0hj
+        ukEnPWsMDzC3c/uISm4CgmZYW2oD1NTGYfZSIUioiEpg3eCNOBEcLuiGKudpGWvajTEXfaWtDvXC
+        owxkI2Zgcsf4IQH46k60+ddaOE+N4wXTGm1cmOXHdeL6KVfzdnQymViaqz68T058IuwCY7NYaVuh
+        d15394qUySH5AwAA//8DAJrMQ4wuAgAA
     headers:
       CF-Cache-Status:
       - DYNAMIC
       CF-RAY:
-      - 85465bd8cea909c9-HFA
+      - 854e5a16dda409cd-HFA
       Cache-Control:
       - no-cache, must-revalidate
       Connection:
@@ -57,14 +57,14 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Mon, 12 Feb 2024 16:55:22 GMT
+      - Tue, 13 Feb 2024 16:12:16 GMT
       Server:
       - cloudflare
       Set-Cookie:
-      - __cf_bm=TkrINgoCmmTaqg05VQ1WHEolPihzU6u57hugghHXWwc-1707756922-1-ATkknLOEsgPnsTzmYjxPFIpBt23UFZQ6s/363rmVG1No8qV6Hf6NG/eS+wuP0S8QLYUa6XR+wZ2Ea+e/wnE+sFg=;
-        path=/; expires=Mon, 12-Feb-24 17:25:22 GMT; domain=.api.openai.com; HttpOnly;
+      - __cf_bm=0bvI3Z_vwqHjtAzex9Nyslo5MvRZHnGmqYqW9mdgUCc-1707840736-1.0-AVAkZCXx5XhwzBIfTNklRI8zl7mXHePuOiErG2Sh6v9/h1pKHcmLlAQvF8FvKGQTRVSiLZyT5Qlj4bOHm6XRsnk=;
+        path=/; expires=Tue, 13-Feb-24 16:42:16 GMT; domain=.api.openai.com; HttpOnly;
         Secure; SameSite=None
-      - _cfuvid=UFML5ZGMwerkaoj2b3b9N4PRATR9slp7OOJ3eSiyuW4-1707756922521-0-604800000;
+      - _cfuvid=JqbEPJNLYeI1J8b90KVpYInAy4kTq.43kv7nw0j_LNQ-1707840736482-0.0-604800000;
         path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
       Transfer-Encoding:
       - chunked
@@ -77,7 +77,7 @@ interactions:
       openai-organization:
       - traceloop
       openai-processing-ms:
-      - '332'
+      - '460'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -95,7 +95,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 9ms
       x-request-id:
-      - req_d671bdc2c19fd05b41f607444a181b1f
+      - req_e802a10fcb6c17271f65febb97334b8f
     status:
       code: 200
       message: OK

diff --git a/packages/traceloop-sdk/tests/test_workflows.py b/packages/traceloop-sdk/tests/test_workflows.py
@@ -1,4 +1,5 @@
 import pytest
+import json
 from openai import OpenAI
 from traceloop.sdk.decorators import workflow, task
 
@@ -10,31 +11,37 @@ def openai_client():
 
 @pytest.mark.vcr
 def test_simple_workflow(exporter, openai_client):
-    @task(name="joke_creation")
-    def create_joke():
+    @task(name="something_creator")
+    def create_something(what: str, subject: str):
         completion = openai_client.chat.completions.create(
             model="gpt-3.5-turbo",
-            messages=[
-                {"role": "user", "content": "Tell me a joke about opentelemetry"}
-            ],
+            messages=[{"role": "user", "content": f"Tell me a {what} about {subject}"}],
         )
         return completion.choices[0].message.content
 
     @workflow(name="pirate_joke_generator")
     def joke_workflow():
-        create_joke()
+        return create_something("joke", subject="OpenTelemetry")
 
-    joke_workflow()
+    joke = joke_workflow()
 
     spans = exporter.get_finished_spans()
     assert [span.name for span in spans] == [
         "openai.chat",
-        "joke_creation.task",
+        "something_creator.task",
         "pirate_joke_generator.workflow",
     ]
     open_ai_span = spans[0]
     assert (
         open_ai_span.attributes["llm.prompts.0.content"]
-        == "Tell me a joke about opentelemetry"
+        == "Tell me a joke about OpenTelemetry"
     )
     assert open_ai_span.attributes.get("llm.completions.0.content")
+
+    task_span = spans[1]
+    assert json.loads(task_span.attributes["traceloop.entity.input"]) == {
+        "args": ["joke"],
+        "kwargs": {"subject": "OpenTelemetry"},
+    }
+
+    assert task_span.attributes.get("traceloop.entity.output") == joke
diff --git a/packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py b/packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py
@@ -1,6 +1,9 @@
+import json
 from functools import wraps
+import os
 from typing import Optional
 
+from opentelemetry import context as context_api
 from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
 
 from traceloop.sdk.tracing import get_tracer, set_workflow_name
@@ -42,7 +45,19 @@ def wrap(*args, **kwargs):
                         SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value
                     )
                     span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
-                    return fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(
+                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
+                            json.dumps({"args": args, "kwargs": kwargs}),
+                        )
+
+                    res = fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, res)
+
+                    return res
 
         return wrap
 
@@ -99,7 +114,18 @@ def wrap(*args, **kwargs):
                     )
                     span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
 
-                    return fn(*args, **kwargs)
+                    if _should_send_prompts():
+                        span.set_attribute(
+                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
+                            json.dumps({"args": args, "kwargs": kwargs}),
+                        )
+
+                    res = fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, res)
+
+                    return res
 
         return wrap
 
@@ -169,7 +195,19 @@ async def wrap(*args, **kwargs):
                         SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value
                     )
                     span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
-                    return await fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(
+                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
+                            json.dumps({"args": args, "kwargs": kwargs}),
+                        )
+
+                    res = await fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, res)
+
+                    return res
 
         return wrap
 
@@ -230,7 +268,19 @@ async def wrap(*args, **kwargs):
                         span.set_attribute(
                             SpanAttributes.TRACELOOP_CORRELATION_ID, correlation_id
                         )
-                    return await fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(
+                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
+                            json.dumps({"args": args, "kwargs": kwargs}),
+                        )
+
+                    res = await fn(*args, **kwargs)
+
+                    if _should_send_prompts():
+                        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, res)
+
+                    return res
 
         return wrap
 
@@ -263,3 +313,9 @@ def atool(name: Optional[str] = None, method_name: Optional[str] = None):
     return atask(
         name=name, method_name=method_name, tlp_span_kind=TraceloopSpanKindValues.TOOL
     )
+
+
+def _should_send_prompts():
+    return (
+        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")