Skip to content

Commit a130443

Browse files
alizenhomalizenhom
authored
Merge pull request #218 from Scale3-Labs/development
Fix StreamWrapper
2 parents 0fbec4c + 5686642 commit a130443

File tree

5 files changed

+118
-200803
lines changed

5 files changed

+118
-200803
lines changed

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

Lines changed: 55 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -276,54 +276,83 @@ def traced_method(wrapped, instance, args, kwargs):
276276

277277

278278
class StreamWrapper:
279-
def __init__(self, stream, span, prompt_tokens, function_call=False, tool_calls=False):
279+
def __init__(
280+
self, stream, span, prompt_tokens, function_call=False, tool_calls=False
281+
):
280282
self.stream = stream
281283
self.span = span
282284
self.prompt_tokens = prompt_tokens
283285
self.function_call = function_call
284286
self.tool_calls = tool_calls
285287
self.result_content = []
286288
self.completion_tokens = 0
289+
self._span_started = False
290+
self._start_span()
291+
292+
def _start_span(self):
293+
if not self._span_started:
294+
self.span.add_event(Event.STREAM_START.value)
295+
self._span_started = True
296+
297+
def _end_span(self):
298+
if self._span_started:
299+
self.span.add_event(Event.STREAM_END.value)
300+
self.span.set_attribute(
301+
"llm.token.counts",
302+
json.dumps(
303+
{
304+
"input_tokens": self.prompt_tokens,
305+
"output_tokens": self.completion_tokens,
306+
"total_tokens": self.prompt_tokens + self.completion_tokens,
307+
}
308+
),
309+
)
310+
self.span.set_attribute(
311+
"llm.responses",
312+
json.dumps(
313+
[
314+
{
315+
"role": "assistant",
316+
"content": "".join(self.result_content),
317+
}
318+
]
319+
),
320+
)
321+
self.span.set_status(StatusCode.OK)
322+
self.span.end()
323+
self._span_started = False
287324

288325
def __enter__(self):
289-
self.span.add_event(Event.STREAM_START.value)
326+
self._start_span()
290327
return self
291328

292329
def __exit__(self, exc_type, exc_val, exc_tb):
293-
self.span.add_event(Event.STREAM_END.value)
294-
self.span.set_attribute(
295-
"llm.token.counts",
296-
json.dumps(
297-
{
298-
"input_tokens": self.prompt_tokens,
299-
"output_tokens": self.completion_tokens,
300-
"total_tokens": self.prompt_tokens + self.completion_tokens,
301-
}
302-
),
303-
)
304-
self.span.set_attribute(
305-
"llm.responses",
306-
json.dumps(
307-
[
308-
{
309-
"role": "assistant",
310-
"content": "".join(self.result_content),
311-
}
312-
]
313-
),
314-
)
315-
self.span.set_status(StatusCode.OK)
316-
self.span.end()
330+
self._end_span()
317331

318332
def __iter__(self):
333+
self._start_span()
334+
return self
335+
336+
def __aiter__(self):
337+
self._start_span()
319338
return self
320339

340+
async def __anext__(self):
341+
try:
342+
chunk = await self.stream.__anext__()
343+
self.process_chunk(chunk)
344+
return chunk
345+
except StopIteration:
346+
self._end_span()
347+
raise
348+
321349
def __next__(self):
322350
try:
323351
chunk = next(self.stream)
324352
self.process_chunk(chunk)
325353
return chunk
326354
except StopIteration:
355+
self._end_span()
327356
raise
328357

329358
def process_chunk(self, chunk):

‎src/langtrace_python_sdk/version.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
__version__ = "2.1.22"
1+
__version__ = "2.1.23"

‎src/tests/openai/cassettes/test_async_chat_completion_streaming.yaml‎

Lines changed: 29 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ interactions:
1616
host:
1717
- api.openai.com
1818
user-agent:
19-
- AsyncOpenAI/Python 1.23.2
19+
- AsyncOpenAI/Python 1.34.0
2020
x-stainless-arch:
2121
- arm64
2222
x-stainless-async:
@@ -26,75 +26,75 @@ interactions:
2626
x-stainless-os:
2727
- MacOS
2828
x-stainless-package-version:
29-
- 1.23.2
29+
- 1.34.0
3030
x-stainless-runtime:
3131
- CPython
3232
x-stainless-runtime-version:
33-
- 3.11.5
33+
- 3.12.3
3434
method: POST
3535
uri: https://api.openai.com/v1/chat/completions
3636
response:
3737
body:
38-
string: 'data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,"finish_reason":null}]}
38+
string: 'data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,"finish_reason":null}]}
3939
4040
41-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"This"},"logprobs":null,"finish_reason":null}]}
41+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"This"},"logprobs":null,"finish_reason":null}]}
4242
4343
44-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
44+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
4545
is"},"logprobs":null,"finish_reason":null}]}
4646
4747
48-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
48+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
4949
a"},"logprobs":null,"finish_reason":null}]}
5050
5151
52-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
52+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
5353
test"},"logprobs":null,"finish_reason":null}]}
5454
5555
56-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
56+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
5757
5858
59-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
59+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
6060
This"},"logprobs":null,"finish_reason":null}]}
6161
6262
63-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
63+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
6464
is"},"logprobs":null,"finish_reason":null}]}
6565
6666
67-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
67+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
6868
a"},"logprobs":null,"finish_reason":null}]}
6969
7070
71-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
71+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
7272
test"},"logprobs":null,"finish_reason":null}]}
7373
7474
75-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
75+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
7676
7777
78-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
78+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
7979
This"},"logprobs":null,"finish_reason":null}]}
8080
8181
82-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
82+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
8383
is"},"logprobs":null,"finish_reason":null}]}
8484
8585
86-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
86+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
8787
a"},"logprobs":null,"finish_reason":null}]}
8888
8989
90-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
90+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"
9191
test"},"logprobs":null,"finish_reason":null}]}
9292
9393
94-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
94+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}]}
9595
9696
97-
data: {"id":"chatcmpl-9G7iUffCm37bRsnrZcMjcD55dwZcN","object":"chat.completion.chunk","created":1713629938,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}
97+
data: {"id":"chatcmpl-9deeLdQVwFn9pGfxPrHrnGhzMINed","object":"chat.completion.chunk","created":1719238077,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}
9898
9999
100100
data: [DONE]
@@ -105,39 +105,33 @@ interactions:
105105
CF-Cache-Status:
106106
- DYNAMIC
107107
CF-RAY:
108-
- 8776740aaa3e73a3-MRS
109-
Cache-Control:
110-
- no-cache, must-revalidate
108+
- 898d497cff180da3-MRS
111109
Connection:
112110
- keep-alive
113111
Content-Type:
114-
- text/event-stream
112+
- text/event-stream; charset=utf-8
115113
Date:
116-
- Sat, 20 Apr 2024 16:18:59 GMT
114+
- Mon, 24 Jun 2024 14:07:57 GMT
117115
Server:
118116
- cloudflare
119117
Set-Cookie:
120-
- __cf_bm=VDKSK7GoYoxisSawjMl9W0b7YZarMfekW_Y69gq5ons-1713629939-1.0.1.1-fH6Sc.9fQ0Kb4MvzvnBRlAmk_cXfYNeNbDd_K6pZeMNxnmMy3qiDlS.olHx3Y7rfDhYg7a3FffrCHr.Xu8j_Uw;
121-
path=/; expires=Sat, 20-Apr-24 16:48:59 GMT; domain=.api.openai.com; HttpOnly;
118+
- __cf_bm=9.DjXDRH6yqveelPQR44ZtyVUYJzgnA0G.nzqFBokzU-1719238077-1.0.1.1-aMCvEWGLFuE6ItsOru0NDaEtwnP4gVQChxLODAqfb5jmClZJa8lKqg.tNgoeMSGnxhnih1YJRhw2gq02fThn1g;
119+
path=/; expires=Mon, 24-Jun-24 14:37:57 GMT; domain=.api.openai.com; HttpOnly;
122120
Secure; SameSite=None
123-
- _cfuvid=GmZ5rJWGxt2nlOYm6_pDkhIo_8V.YkD6O9_B1qloO7g-1713629939085-0.0.1.1-604800000;
121+
- _cfuvid=O9yzYcvRGvChbCerFoeKPvIqX_jwhxHPRVMMxoF4wLY-1719238077679-0.0.1.1-604800000;
124122
path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
125123
Transfer-Encoding:
126124
- chunked
127-
access-control-allow-origin:
128-
- '*'
129125
alt-svc:
130126
- h3=":443"; ma=86400
131-
openai-model:
132-
- gpt-4-0613
133127
openai-organization:
134128
- scale3-1
135129
openai-processing-ms:
136-
- '247'
130+
- '193'
137131
openai-version:
138132
- '2020-10-01'
139133
strict-transport-security:
140-
- max-age=15724800; includeSubDomains
134+
- max-age=31536000; includeSubDomains
141135
x-ratelimit-limit-requests:
142136
- '10000'
143137
x-ratelimit-limit-tokens:
@@ -151,7 +145,7 @@ interactions:
151145
x-ratelimit-reset-tokens:
152146
- 5ms
153147
x-request-id:
154-
- req_4c6f987df5c44fb9c842d54126d2608d
148+
- ae303b10004d584d9da9cb432917ef64
155149
status:
156150
code: 200
157151
message: OK

0 commit comments

Comments
 (0)