fix: input/output for workflows & tasks

traceloop · Feb 13, 2024 · a1b2d2d · a1b2d2d
1 parent 301dc79
commit a1b2d2d
Showing 1 changed file with 60 additions and 4 deletions.
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")