Skip to content

Commit a9f48f3

Browse files
alizenhomalizenhom
committed
add prompts to events
1 parent 43b12ea commit a9f48f3

File tree

9 files changed

+118
-95
lines changed

9 files changed

+118
-95
lines changed

‎src/langtrace_python_sdk/instrumentation/chroma/patch.py‎

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
"""
1616

1717
from langtrace.trace_attributes import DatabaseSpanAttributes
18-
from langtrace_python_sdk.utils.llm import set_span_attributes
18+
from langtrace_python_sdk.utils import set_span_attribute
1919
from langtrace_python_sdk.utils.silently_fail import silently_fail
2020
from opentelemetry import baggage, trace
2121
from opentelemetry.trace import SpanKind
@@ -119,20 +119,20 @@ def handle_null_params(param):
119119

120120
@silently_fail
121121
def _set_chroma_add_attributes(span, kwargs):
122-
set_span_attributes(
122+
set_span_attribute(
123123
span, "db.chroma.add.ids_count", get_count_or_none(kwargs.get("ids"))
124124
)
125-
set_span_attributes(
125+
set_span_attribute(
126126
span,
127127
"db.chroma.add.embeddings_count",
128128
get_count_or_none(kwargs.get("embeddings")),
129129
)
130-
set_span_attributes(
130+
set_span_attribute(
131131
span,
132132
"db.chroma.add.metadatas_count",
133133
get_count_or_none(kwargs.get("metadatas")),
134134
)
135-
set_span_attributes(
135+
set_span_attribute(
136136
span,
137137
"db.chroma.add.documents_count",
138138
get_count_or_none(kwargs.get("documents")),
@@ -141,71 +141,71 @@ def _set_chroma_add_attributes(span, kwargs):
141141

142142
@silently_fail
143143
def _set_chroma_get_attributes(span, kwargs):
144-
set_span_attributes(
144+
set_span_attribute(
145145
span, "db.chroma.get.ids_count", get_count_or_none(kwargs.get("ids"))
146146
)
147-
set_span_attributes(
147+
set_span_attribute(
148148
span, "db.chroma.get.where", handle_null_params(kwargs.get("where"))
149149
)
150-
set_span_attributes(span, "db.chroma.get.limit", kwargs.get("limit"))
151-
set_span_attributes(span, "db.chroma.get.offset", kwargs.get("offset"))
152-
set_span_attributes(
150+
set_span_attribute(span, "db.chroma.get.limit", kwargs.get("limit"))
151+
set_span_attribute(span, "db.chroma.get.offset", kwargs.get("offset"))
152+
set_span_attribute(
153153
span,
154154
"db.chroma.get.where_document",
155155
handle_null_params(kwargs.get("where_document")),
156156
)
157-
set_span_attributes(
157+
set_span_attribute(
158158
span, "db.chroma.get.include", handle_null_params(kwargs.get("include"))
159159
)
160160

161161

162162
@silently_fail
163163
def _set_chroma_query_attributes(span, kwargs):
164-
set_span_attributes(
164+
set_span_attribute(
165165
span,
166166
"db.chroma.query.query_embeddings_count",
167167
get_count_or_none(kwargs.get("query_embeddings")),
168168
)
169-
set_span_attributes(
169+
set_span_attribute(
170170
span,
171171
"db.chroma.query.query_texts_count",
172172
get_count_or_none(kwargs.get("query_texts")),
173173
)
174-
set_span_attributes(span, "db.chroma.query.n_results", kwargs.get("n_results"))
175-
set_span_attributes(
174+
set_span_attribute(span, "db.chroma.query.n_results", kwargs.get("n_results"))
175+
set_span_attribute(
176176
span, "db.chroma.query.where", handle_null_params(kwargs.get("where"))
177177
)
178-
set_span_attributes(
178+
set_span_attribute(
179179
span,
180180
"db.chroma.query.where_document",
181181
handle_null_params(kwargs.get("where_document")),
182182
)
183-
set_span_attributes(
183+
set_span_attribute(
184184
span, "db.chroma.query.include", handle_null_params(kwargs.get("include"))
185185
)
186186

187187

188188
@silently_fail
189189
def _set_chroma_peek_attributes(span, kwargs):
190-
set_span_attributes(span, "db.chroma.peek.limit", kwargs.get("limit"))
190+
set_span_attribute(span, "db.chroma.peek.limit", kwargs.get("limit"))
191191

192192

193193
@silently_fail
194194
def _set_chroma_update_attributes(span, kwargs):
195-
set_span_attributes(
195+
set_span_attribute(
196196
span, "db.chroma.update.ids_count", get_count_or_none(kwargs.get("ids"))
197197
)
198-
set_span_attributes(
198+
set_span_attribute(
199199
span,
200200
"db.chroma.update.embeddings_count",
201201
get_count_or_none(kwargs.get("embeddings")),
202202
)
203-
set_span_attributes(
203+
set_span_attribute(
204204
span,
205205
"db.chroma.update.metadatas_count",
206206
get_count_or_none(kwargs.get("metadatas")),
207207
)
208-
set_span_attributes(
208+
set_span_attribute(
209209
span,
210210
"db.chroma.update.documents_count",
211211
get_count_or_none(kwargs.get("documents")),
@@ -214,23 +214,23 @@ def _set_chroma_update_attributes(span, kwargs):
214214

215215
@silently_fail
216216
def _set_chroma_modify_attributes(span, kwargs):
217-
set_span_attributes(span, "db.chroma.modify.name", kwargs.get("name"))
217+
set_span_attribute(span, "db.chroma.modify.name", kwargs.get("name"))
218218
# TODO: Add metadata attribute
219219

220220

221221
@silently_fail
222222
def _set_chroma_upsert_attributes(span, kwargs):
223-
set_span_attributes(
223+
set_span_attribute(
224224
span,
225225
"db.chroma.upsert.embeddings_count",
226226
get_count_or_none(kwargs.get("embeddings")),
227227
)
228-
set_span_attributes(
228+
set_span_attribute(
229229
span,
230230
"db.chroma.upsert.metadatas_count",
231231
get_count_or_none(kwargs.get("metadatas")),
232232
)
233-
set_span_attributes(
233+
set_span_attribute(
234234
span,
235235
"db.chroma.upsert.documents_count",
236236
get_count_or_none(kwargs.get("documents")),
@@ -239,13 +239,13 @@ def _set_chroma_upsert_attributes(span, kwargs):
239239

240240
@silently_fail
241241
def _set_chroma_delete_attributes(span, kwargs):
242-
set_span_attributes(
242+
set_span_attribute(
243243
span, "db.chroma.delete.ids_count", get_count_or_none(kwargs.get("ids"))
244244
)
245-
set_span_attributes(
245+
set_span_attribute(
246246
span, "db.chroma.delete.where", handle_null_params(kwargs.get("where"))
247247
)
248-
set_span_attributes(
248+
set_span_attribute(
249249
span,
250250
"db.chroma.delete.where_document",
251251
handle_null_params(kwargs.get("where_document")),

‎src/langtrace_python_sdk/instrumentation/qdrant/patch.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import json
1818
from langtrace.trace_attributes import DatabaseSpanAttributes
1919
from langtrace_python_sdk.utils.silently_fail import silently_fail
20-
from langtrace_python_sdk.utils.llm import set_span_attributes
20+
from langtrace_python_sdk.utils import set_span_attribute
2121
from opentelemetry import baggage, trace
2222
from opentelemetry.trace import SpanKind
2323
from opentelemetry.trace.status import Status, StatusCode
@@ -64,7 +64,7 @@ def traced_method(wrapped, instance, args, kwargs):
6464
) as span:
6565
collection_name = kwargs.get("collection_name") or args[0]
6666
operation = api["OPERATION"]
67-
set_span_attributes(span, "db.collection.name", collection_name)
67+
set_span_attribute(span, "db.collection.name", collection_name)
6868

6969
if operation == "add":
7070
_set_upload_attributes(span, args, kwargs, "documents")
@@ -111,7 +111,7 @@ def _set_upsert_attributes(span, args, kwargs):
111111
else:
112112
# In case of using Batch.
113113
length = len(points.ids)
114-
set_span_attributes(span, "db.upsert.points_count", length)
114+
set_span_attribute(span, "db.upsert.points_count", length)
115115

116116

117117
@silently_fail
@@ -123,16 +123,16 @@ def _set_upload_attributes(span, args, kwargs, field):
123123
# In case of using Batch.
124124
length = len(docs.ids)
125125

126-
set_span_attributes(span, f"db.upload.{field}_count", length)
126+
set_span_attribute(span, f"db.upload.{field}_count", length)
127127

128128

129129
@silently_fail
130130
def _set_search_attributes(span, args, kwargs):
131131
limit = kwargs.get("limit") or 10
132-
set_span_attributes(span, "db.query.top_k", limit)
132+
set_span_attribute(span, "db.query.top_k", limit)
133133

134134

135135
@silently_fail
136136
def _set_batch_search_attributes(span, args, kwargs, method):
137137
requests = kwargs.get("requests") or []
138-
set_span_attributes(span, f"db.{method}.requests_count", len(requests))
138+
set_span_attribute(span, f"db.{method}.requests_count", len(requests))

‎src/langtrace_python_sdk/utils/__init__.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,21 @@
11
from openai import NOT_GIVEN
22
from .sdk_version_checker import SDKVersionChecker
3+
from opentelemetry.trace import Span
4+
from langtrace.trace_attributes import SpanAttributes
35

46

5-
def set_span_attribute(span, name, value):
7+
def set_span_attribute(span: Span, name, value):
68
if value is not None:
79
if value != "" or value != NOT_GIVEN:
8-
span.set_attribute(name, value)
10+
if name == SpanAttributes.LLM_PROMPTS:
11+
span.add_event(
12+
name=SpanAttributes.LLM_CONTENT_PROMPT,
13+
attributes={
14+
SpanAttributes.LLM_PROMPTS: value,
15+
},
16+
)
17+
else:
18+
span.set_attribute(name, value)
919
return
1020

1121

‎src/langtrace_python_sdk/utils/llm.py‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""
1616

1717
from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
18+
from langtrace_python_sdk.utils import set_span_attribute
1819
from openai import NOT_GIVEN
1920
from tiktoken import get_encoding
2021

@@ -69,13 +70,6 @@ def calculate_price_from_usage(model, usage):
6970
return 0
7071

7172

72-
def set_span_attributes(span, name, value):
73-
if value is not None:
74-
if value != "":
75-
span.set_attribute(name, value)
76-
return
77-
78-
7973
def get_langtrace_attributes(version, service_provider, vendor_type="llm"):
8074
return {
8175
SpanAttributes.LANGTRACE_SDK_NAME: LANGTRACE_SDK_NAME,
@@ -151,26 +145,26 @@ def set_usage_attributes(span, usage):
151145
input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
152146
output_tokens = usage.get("output_tokens") or usage.get("completion_tokens") or 0
153147

154-
set_span_attributes(
148+
set_span_attribute(
155149
span,
156150
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
157151
input_tokens,
158152
)
159153

160-
set_span_attributes(
154+
set_span_attribute(
161155
span,
162156
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
163157
output_tokens,
164158
)
165159

166-
set_span_attributes(
160+
set_span_attribute(
167161
span,
168162
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
169163
input_tokens + output_tokens,
170164
)
171165

172166
if "search_units" in usage:
173-
set_span_attributes(
167+
set_span_attribute(
174168
span, SpanAttributes.LLM_USAGE_SEARCH_UNITS, usage["search_units"]
175169
)
176170

‎src/tests/anthropic/test_anthropic.py‎

Lines changed: 16 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@
33
import importlib
44
from langtrace_python_sdk.constants.instrumentation.anthropic import APIS
55
from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
6-
from tests.utils import assert_response_format, assert_token_count
6+
from tests.utils import (
7+
assert_langtrace_attributes,
8+
assert_prompt_in_events,
9+
assert_response_format,
10+
assert_token_count,
11+
)
712
from importlib_metadata import version as v
813

914
from langtrace.trace_attributes import SpanAttributes
@@ -27,20 +32,16 @@ def test_anthropic(anthropic_client, exporter):
2732

2833
assert completion_span.name == "anthropic.messages.create"
2934
attributes = completion_span.attributes
30-
31-
assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == LANGTRACE_SDK_NAME
32-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "Anthropic"
33-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == "llm"
34-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v("anthropic")
35-
assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v(LANGTRACE_SDK_NAME)
35+
assert_langtrace_attributes(attributes, "Anthropic")
36+
assert_prompt_in_events(completion_span.events, messages_value)
3637
assert attributes.get(SpanAttributes.LLM_URL) == "https://api.anthropic.com"
3738
assert (
3839
attributes.get(SpanAttributes.LLM_PATH) == APIS["MESSAGES_CREATE"]["ENDPOINT"]
3940
)
4041
assert attributes.get(SpanAttributes.LLM_REQUEST_MODEL) == llm_model_value
41-
assert json.loads(attributes.get(SpanAttributes.LLM_PROMPTS)) == json.dumps(
42-
messages_value
43-
)
42+
# assert json.loads(attributes.get(SpanAttributes.LLM_PROMPTS)) == json.dumps(
43+
# messages_value
44+
# )
4445
assert attributes.get(SpanAttributes.LLM_IS_STREAMING) is False
4546

4647
assert_token_count(attributes)
@@ -72,25 +73,21 @@ def test_anthropic_streaming(anthropic_client, exporter):
7273
assert streaming_span.name == "anthropic.messages.create"
7374
attributes = streaming_span.attributes
7475

75-
assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == LANGTRACE_SDK_NAME
76-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "Anthropic"
77-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == "llm"
78-
assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v("anthropic")
79-
assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v(LANGTRACE_SDK_NAME)
76+
assert_langtrace_attributes(attributes, "Anthropic")
77+
assert_prompt_in_events(streaming_span.events, messages_value)
8078

8179
assert attributes.get(SpanAttributes.LLM_URL) == "https://api.anthropic.com"
8280
assert (
8381
attributes.get(SpanAttributes.LLM_PATH) == APIS["MESSAGES_CREATE"]["ENDPOINT"]
8482
)
8583
assert attributes.get(SpanAttributes.LLM_REQUEST_MODEL) == llm_model_value
86-
assert json.loads(attributes.get(SpanAttributes.LLM_PROMPTS)) == json.dumps(
87-
messages_value
88-
)
8984
assert attributes.get(SpanAttributes.LLM_IS_STREAMING) is True
9085

9186
events = streaming_span.events
9287

93-
assert len(events) - 2 == chunk_count # -2 for start and end events
88+
assert (
89+
len(events) - 3 == chunk_count
90+
) # -2 for start and end events and prompt event
9491

9592
assert_token_count(attributes)
9693
assert_response_format(attributes)

0 commit comments

Comments
 (0)