Merge branch 'development' of github.com:Scale3-Labs/langtrace-python-sdk into development

Author: karthikscale3

README.md

@@ -228,6 +228,11 @@ from langtrace_python_sdk import get_prompt_from_registry
 prompt = get_prompt_from_registry(<Registry ID>, options={"prompt_version": 1, "variables": {"foo": "bar"} })
 ```
 
+### Opt out of tracing prompt and completion data
+By default, prompt and completion data are captured. If you would like to opt out of it, set the following env var,
+
+`TRACE_PROMPT_COMPLETION_DATA=false`
+
 ## Supported integrations
 
 Langtrace automatically captures traces from the following vendors:

src/langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -25,6 +25,7 @@
     get_llm_url,
     is_streaming,
     set_event_completion,
+    set_event_completion_chunk,
     set_usage_attributes,
 )
 from opentelemetry.trace import SpanKind
@@ -119,10 +120,7 @@ def handle_streaming_response(result, span):
                 # Assuming span.add_event is part of a larger logging or event system
                 # Add event for each chunk of content
                 if content:
-                    span.add_event(
-                        Event.STREAM_OUTPUT.value,
-                        {SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: "".join(content)},
-                    )
+                    set_event_completion_chunk(span, "".join(content))
 
                 # Assuming this is part of a generator, yield chunk or aggregated content
                 yield content

src/langtrace_python_sdk/instrumentation/cohere/patch.py

@@ -22,6 +22,7 @@
     get_extra_attributes,
     get_llm_url,
     set_event_completion,
+    set_event_completion_chunk,
     set_usage_attributes,
 )
 from langtrace.trace_attributes import Event, LLMSpanAttributes
@@ -403,10 +404,7 @@ def traced_method(wrapped, instance, args, kwargs):
                         content = event.text
                     else:
                         content = ""
-                    span.add_event(
-                        Event.STREAM_OUTPUT.value,
-                        {SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: "".join(content)},
-                    )
+                    set_event_completion_chunk(span, "".join(content))
 
                     if (
                         hasattr(event, "finish_reason")

src/langtrace_python_sdk/instrumentation/groq/patch.py

@@ -30,6 +30,7 @@
     get_llm_url,
     get_langtrace_attributes,
     set_event_completion,
+    set_event_completion_chunk,
     set_usage_attributes,
 )
 from langtrace_python_sdk.constants.instrumentation.common import (
@@ -242,15 +243,14 @@ def handle_streaming_response(
                                         content = content + []
                 else:
                     content = []
-                span.add_event(
-                    Event.STREAM_OUTPUT.value,
-                    {
-                        SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
-                            "".join(content)
-                            if len(content) > 0 and content[0] is not None
-                            else ""
-                        )
-                    },
+
+                set_event_completion_chunk(
+                    span,
+                    (
+                        "".join(content)
+                        if len(content) > 0 and content[0] is not None
+                        else ""
+                    ),
                 )
                 result_content.append(content[0] if len(content) > 0 else "")
                 yield chunk

src/langtrace_python_sdk/instrumentation/ollama/patch.py

@@ -6,6 +6,7 @@
     get_llm_request_attributes,
     get_llm_url,
     set_event_completion,
+    set_event_completion_chunk,
 )
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
@@ -177,12 +178,8 @@ def _handle_streaming_response(span, response, api):
             if api == "generate":
                 accumulated_tokens["response"] += chunk["response"]
 
-            span.add_event(
-                Event.STREAM_OUTPUT.value,
-                {
-                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: chunk.get("response")
-                    or chunk.get("message").get("content"),
-                },
+            set_event_completion_chunk(
+                span, chunk.get("response") or chunk.get("message").get("content")
             )
 
         _set_response_attributes(span, chunk | accumulated_tokens)
@@ -211,12 +208,7 @@ async def _ahandle_streaming_response(span, response, api):
             if api == "generate":
                 accumulated_tokens["response"] += chunk["response"]
 
-            span.add_event(
-                Event.STREAM_OUTPUT.value,
-                {
-                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: json.dumps(chunk),
-                },
-            )
+            set_event_completion_chunk(span, chunk)
         _set_response_attributes(span, chunk | accumulated_tokens)
     finally:
         # Finalize span after processing all chunks

src/langtrace_python_sdk/utils/__init__.py

@@ -2,23 +2,32 @@
 from .sdk_version_checker import SDKVersionChecker
 from opentelemetry.trace import Span
 from langtrace.trace_attributes import SpanAttributes
+import os
 
 
 def set_span_attribute(span: Span, name, value):
     if value is not None:
         if value != "" or value != NOT_GIVEN:
             if name == SpanAttributes.LLM_PROMPTS:
-                span.add_event(
-                    name=SpanAttributes.LLM_CONTENT_PROMPT,
-                    attributes={
-                        SpanAttributes.LLM_PROMPTS: value,
-                    },
-                )
+                set_event_prompt(span, value)
             else:
                 span.set_attribute(name, value)
     return
 
 
+def set_event_prompt(span: Span, prompt):
+    enabled = os.environ.get("TRACE_PROMPT_COMPLETION_DATA", "true")
+    if enabled.lower() == "false":
+        return
+
+    span.add_event(
+        name=SpanAttributes.LLM_CONTENT_PROMPT,
+        attributes={
+            SpanAttributes.LLM_PROMPTS: prompt,
+        },
+    )
+
+
 def check_if_sdk_is_outdated():
     SDKVersionChecker().check()
     return

src/langtrace_python_sdk/utils/llm.py

@@ -30,6 +30,7 @@
 from opentelemetry import baggage
 from opentelemetry.trace import Span
 from opentelemetry.trace.status import StatusCode
+import os
 
 
 def estimate_tokens(prompt):
@@ -42,6 +43,9 @@ def estimate_tokens(prompt):
 
 
 def set_event_completion_chunk(span: Span, chunk):
+    enabled = os.environ.get("TRACE_PROMPT_COMPLETION_DATA", "true")
+    if enabled.lower() == "false":
+        return
     span.add_event(
         name=SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK,
         attributes={
@@ -203,6 +207,9 @@ def get_tool_calls(item):
 
 
 def set_event_completion(span: Span, result_content):
+    enabled = os.environ.get("TRACE_PROMPT_COMPLETION_DATA", "true")
+    if enabled.lower() == "false":
+        return
 
     span.add_event(
         name=SpanAttributes.LLM_CONTENT_COMPLETION,
@@ -352,15 +359,9 @@ def process_chunk(self, chunk):
                                 )
                                 self.completion_tokens += token_counts
                                 content.append(tool_call.function.arguments)
-            self.span.add_event(
-                Event.STREAM_OUTPUT.value,
-                {
-                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
-                        "".join(content)
-                        if len(content) > 0 and content[0] is not None
-                        else ""
-                    )
-                },
+            set_event_completion_chunk(
+                self.span,
+                "".join(content) if len(content) > 0 and content[0] is not None else "",
             )
             if content:
                 self.result_content.append(content[0])
@@ -369,16 +370,11 @@ def process_chunk(self, chunk):
             token_counts = estimate_tokens(chunk.text)
             self.completion_tokens += token_counts
             content = [chunk.text]
-            self.span.add_event(
-                Event.STREAM_OUTPUT.value,
-                {
-                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
-                        "".join(content)
-                        if len(content) > 0 and content[0] is not None
-                        else ""
-                    )
-                },
+            set_event_completion_chunk(
+                self.span,
+                "".join(content) if len(content) > 0 and content[0] is not None else "",
             )
+
             if content:
                 self.result_content.append(content[0])
 

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.2.8"
+__version__ = "2.2.9"