Skip to content

Commit a9d5e15

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
Merge branch 'webb/langchain/tool-pipeline-name' into webb/langchain/agent-run-name
2 parents 45b3dae + d6a28b7 commit a9d5e15

4 files changed

Lines changed: 160 additions & 34 deletions

File tree

‎tests/conftest.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1266,7 +1266,11 @@ def streaming_chat_completions_model_response():
12661266
@pytest.fixture
12671267
def nonstreaming_chat_completions_model_response():
12681268
def inner(
1269-
response_id: str, response_model: str, message_content: str, created: int
1269+
response_id: str,
1270+
response_model: str,
1271+
message_content: str,
1272+
created: int,
1273+
usage: openai.types.CompletionUsage,
12701274
):
12711275
return openai.types.chat.ChatCompletion(
12721276
id=response_id,
@@ -1282,11 +1286,7 @@ def inner(
12821286
created=created,
12831287
model=response_model,
12841288
object="chat.completion",
1285-
usage=openai.types.CompletionUsage(
1286-
prompt_tokens=10,
1287-
completion_tokens=20,
1288-
total_tokens=30,
1289-
),
1289+
usage=usage,
12901290
)
12911291

12921292
return inner

‎tests/integrations/langchain/test_langchain.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,8 +189,13 @@ def test_langchain_chat_with_run_name(
189189
nonstreaming_chat_completions_model_response(
190190
response_id="chat-id",
191191
response_model="response-model-id",
192-
response_content="the model response",
192+
message_content="the model response",
193193
created=10000000,
194+
usage=CompletionUsage(
195+
prompt_tokens=20,
196+
completion_tokens=10,
197+
total_tokens=30,
198+
),
194199
),
195200
serialize_pydantic=True,
196201
request_headers=request_headers,

‎tests/integrations/litellm/test_litellm.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ async def __call__(self, *args, **kwargs):
3333
from sentry_sdk.utils import package_version
3434

3535
from openai import OpenAI, AsyncOpenAI
36+
from openai.types import CompletionUsage
3637

3738
from concurrent.futures import ThreadPoolExecutor
3839

@@ -165,6 +166,11 @@ def test_nonstreaming_chat_completion(
165166
response_model="gpt-3.5-turbo",
166167
message_content="Test response",
167168
created=1234567890,
169+
usage=CompletionUsage(
170+
prompt_tokens=10,
171+
completion_tokens=20,
172+
total_tokens=30,
173+
),
168174
),
169175
serialize_pydantic=True,
170176
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -252,6 +258,11 @@ async def test_async_nonstreaming_chat_completion(
252258
response_model="gpt-3.5-turbo",
253259
message_content="Test response",
254260
created=1234567890,
261+
usage=CompletionUsage(
262+
prompt_tokens=10,
263+
completion_tokens=20,
264+
total_tokens=30,
265+
),
255266
),
256267
serialize_pydantic=True,
257268
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -919,6 +930,11 @@ def test_span_origin(
919930
response_model="gpt-3.5-turbo",
920931
message_content="Test response",
921932
created=1234567890,
933+
usage=CompletionUsage(
934+
prompt_tokens=10,
935+
completion_tokens=20,
936+
total_tokens=30,
937+
),
922938
),
923939
serialize_pydantic=True,
924940
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -969,6 +985,11 @@ def test_multiple_providers(
969985
response_model="gpt-3.5-turbo",
970986
message_content="Test response",
971987
created=1234567890,
988+
usage=CompletionUsage(
989+
prompt_tokens=10,
990+
completion_tokens=20,
991+
total_tokens=30,
992+
),
972993
),
973994
serialize_pydantic=True,
974995
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1068,6 +1089,11 @@ async def test_async_multiple_providers(
10681089
response_model="gpt-3.5-turbo",
10691090
message_content="Test response",
10701091
created=1234567890,
1092+
usage=CompletionUsage(
1093+
prompt_tokens=10,
1094+
completion_tokens=20,
1095+
total_tokens=30,
1096+
),
10711097
),
10721098
serialize_pydantic=True,
10731099
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1168,6 +1194,11 @@ def test_additional_parameters(
11681194
response_model="gpt-3.5-turbo",
11691195
message_content="Test response",
11701196
created=1234567890,
1197+
usage=CompletionUsage(
1198+
prompt_tokens=10,
1199+
completion_tokens=20,
1200+
total_tokens=30,
1201+
),
11711202
),
11721203
serialize_pydantic=True,
11731204
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1231,6 +1262,11 @@ async def test_async_additional_parameters(
12311262
response_model="gpt-3.5-turbo",
12321263
message_content="Test response",
12331264
created=1234567890,
1265+
usage=CompletionUsage(
1266+
prompt_tokens=10,
1267+
completion_tokens=20,
1268+
total_tokens=30,
1269+
),
12341270
),
12351271
serialize_pydantic=True,
12361272
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1294,6 +1330,11 @@ def test_no_integration(
12941330
response_model="gpt-3.5-turbo",
12951331
message_content="Test response",
12961332
created=1234567890,
1333+
usage=CompletionUsage(
1334+
prompt_tokens=10,
1335+
completion_tokens=20,
1336+
total_tokens=30,
1337+
),
12971338
),
12981339
serialize_pydantic=True,
12991340
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1346,6 +1387,11 @@ async def test_async_no_integration(
13461387
response_model="gpt-3.5-turbo",
13471388
message_content="Test response",
13481389
created=1234567890,
1390+
usage=CompletionUsage(
1391+
prompt_tokens=10,
1392+
completion_tokens=20,
1393+
total_tokens=30,
1394+
),
13491395
),
13501396
serialize_pydantic=True,
13511397
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1528,6 +1574,11 @@ def test_binary_content_encoding_image_url(
15281574
response_model="gpt-3.5-turbo",
15291575
message_content="Test response",
15301576
created=1234567890,
1577+
usage=CompletionUsage(
1578+
prompt_tokens=10,
1579+
completion_tokens=20,
1580+
total_tokens=30,
1581+
),
15311582
),
15321583
serialize_pydantic=True,
15331584
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1611,6 +1662,11 @@ async def test_async_binary_content_encoding_image_url(
16111662
response_model="gpt-3.5-turbo",
16121663
message_content="Test response",
16131664
created=1234567890,
1665+
usage=CompletionUsage(
1666+
prompt_tokens=10,
1667+
completion_tokens=20,
1668+
total_tokens=30,
1669+
),
16141670
),
16151671
serialize_pydantic=True,
16161672
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1696,6 +1752,11 @@ def test_binary_content_encoding_mixed_content(
16961752
response_model="gpt-3.5-turbo",
16971753
message_content="Test response",
16981754
created=1234567890,
1755+
usage=CompletionUsage(
1756+
prompt_tokens=10,
1757+
completion_tokens=20,
1758+
total_tokens=30,
1759+
),
16991760
),
17001761
serialize_pydantic=True,
17011762
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1768,6 +1829,11 @@ async def test_async_binary_content_encoding_mixed_content(
17681829
response_model="gpt-3.5-turbo",
17691830
message_content="Test response",
17701831
created=1234567890,
1832+
usage=CompletionUsage(
1833+
prompt_tokens=10,
1834+
completion_tokens=20,
1835+
total_tokens=30,
1836+
),
17711837
),
17721838
serialize_pydantic=True,
17731839
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1839,6 +1905,11 @@ def test_binary_content_encoding_uri_type(
18391905
response_model="gpt-3.5-turbo",
18401906
message_content="Test response",
18411907
created=1234567890,
1908+
usage=CompletionUsage(
1909+
prompt_tokens=10,
1910+
completion_tokens=20,
1911+
total_tokens=30,
1912+
),
18421913
),
18431914
serialize_pydantic=True,
18441915
request_headers={"X-Stainless-Raw-Response": "true"},
@@ -1916,6 +1987,11 @@ async def test_async_binary_content_encoding_uri_type(
19161987
response_model="gpt-3.5-turbo",
19171988
message_content="Test response",
19181989
created=1234567890,
1990+
usage=CompletionUsage(
1991+
prompt_tokens=10,
1992+
completion_tokens=20,
1993+
total_tokens=30,
1994+
),
19191995
),
19201996
serialize_pydantic=True,
19211997
request_headers={"X-Stainless-Raw-Response": "true"},

‎tests/integrations/openai/test_openai.py‎

Lines changed: 72 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,15 @@ def test_nonstreaming_chat_completion_no_prompts(
128128
client = OpenAI(api_key="z")
129129
client.chat.completions._post = mock.Mock(
130130
return_value=nonstreaming_chat_completions_model_response(
131-
response_id="chatcmpl-test",
131+
response_id="chat-id",
132132
response_model="gpt-3.5-turbo",
133-
message_content="Test response",
134-
created=1234567890,
133+
message_content="the model response",
134+
created=10000000,
135+
usage=CompletionUsage(
136+
prompt_tokens=20,
137+
completion_tokens=10,
138+
total_tokens=30,
139+
),
135140
)
136141
)
137142

@@ -237,10 +242,15 @@ def test_nonstreaming_chat_completion(
237242
client = OpenAI(api_key="z")
238243
client.chat.completions._post = mock.Mock(
239244
return_value=nonstreaming_chat_completions_model_response(
240-
response_id="chatcmpl-test",
245+
response_id="chat-id",
241246
response_model="gpt-3.5-turbo",
242-
message_content="Test response",
243-
created=1234567890,
247+
message_content="the model response",
248+
created=10000000,
249+
usage=CompletionUsage(
250+
prompt_tokens=20,
251+
completion_tokens=10,
252+
total_tokens=30,
253+
),
244254
)
245255
)
246256

@@ -328,10 +338,15 @@ async def test_nonstreaming_chat_completion_async_no_prompts(
328338
client = AsyncOpenAI(api_key="z")
329339
client.chat.completions._post = mock.AsyncMock(
330340
return_value=nonstreaming_chat_completions_model_response(
331-
response_id="chatcmpl-test",
341+
response_id="chat-id",
332342
response_model="gpt-3.5-turbo",
333-
message_content="Test response",
334-
created=1234567890,
343+
message_content="the model response",
344+
created=10000000,
345+
usage=CompletionUsage(
346+
prompt_tokens=20,
347+
completion_tokens=10,
348+
total_tokens=30,
349+
),
335350
)
336351
)
337352

@@ -435,10 +450,15 @@ async def test_nonstreaming_chat_completion_async(
435450
client = AsyncOpenAI(api_key="z")
436451
client.chat.completions._post = AsyncMock(
437452
return_value=nonstreaming_chat_completions_model_response(
438-
response_id="chatcmpl-test",
453+
response_id="chat-id",
439454
response_model="gpt-3.5-turbo",
440-
message_content="Test response",
441-
created=1234567890,
455+
message_content="the model response",
456+
created=10000000,
457+
usage=CompletionUsage(
458+
prompt_tokens=20,
459+
completion_tokens=10,
460+
total_tokens=30,
461+
),
442462
)
443463
)
444464

@@ -1888,10 +1908,15 @@ def test_span_origin_nonstreaming_chat(
18881908
client = OpenAI(api_key="z")
18891909
client.chat.completions._post = mock.Mock(
18901910
return_value=nonstreaming_chat_completions_model_response(
1891-
response_id="chatcmpl-test",
1911+
response_id="chat-id",
18921912
response_model="gpt-3.5-turbo",
1893-
message_content="Test response",
1894-
created=1234567890,
1913+
message_content="the model response",
1914+
created=10000000,
1915+
usage=CompletionUsage(
1916+
prompt_tokens=20,
1917+
completion_tokens=10,
1918+
total_tokens=30,
1919+
),
18951920
)
18961921
)
18971922

@@ -1919,10 +1944,15 @@ async def test_span_origin_nonstreaming_chat_async(
19191944
client = AsyncOpenAI(api_key="z")
19201945
client.chat.completions._post = AsyncMock(
19211946
return_value=nonstreaming_chat_completions_model_response(
1922-
response_id="chatcmpl-test",
1947+
response_id="chat-id",
19231948
response_model="gpt-3.5-turbo",
1924-
message_content="Test response",
1925-
created=1234567890,
1949+
message_content="the model response",
1950+
created=10000000,
1951+
usage=CompletionUsage(
1952+
prompt_tokens=20,
1953+
completion_tokens=10,
1954+
total_tokens=30,
1955+
),
19261956
)
19271957
)
19281958

@@ -3686,10 +3716,15 @@ def test_empty_tools_in_chat_completion(
36863716
client = OpenAI(api_key="z")
36873717
client.chat.completions._post = mock.Mock(
36883718
return_value=nonstreaming_chat_completions_model_response(
3689-
response_id="chatcmpl-test",
3719+
response_id="chat-id",
36903720
response_model="gpt-3.5-turbo",
3691-
message_content="Test response",
3692-
created=1234567890,
3721+
message_content="the model response",
3722+
created=10000000,
3723+
usage=CompletionUsage(
3724+
prompt_tokens=20,
3725+
completion_tokens=10,
3726+
total_tokens=30,
3727+
),
36933728
)
36943729
)
36953730

@@ -3740,10 +3775,15 @@ def test_openai_message_role_mapping(
37403775
client = OpenAI(api_key="z")
37413776
client.chat.completions._post = mock.Mock(
37423777
return_value=nonstreaming_chat_completions_model_response(
3743-
response_id="chatcmpl-test",
3778+
response_id="chat-id",
37443779
response_model="gpt-3.5-turbo",
3745-
message_content="Test response",
3746-
created=1234567890,
3780+
message_content="the model response",
3781+
created=10000000,
3782+
usage=CompletionUsage(
3783+
prompt_tokens=20,
3784+
completion_tokens=10,
3785+
total_tokens=30,
3786+
),
37473787
)
37483788
)
37493789

@@ -3780,10 +3820,15 @@ def test_openai_message_truncation(
37803820
client = OpenAI(api_key="z")
37813821
client.chat.completions._post = mock.Mock(
37823822
return_value=nonstreaming_chat_completions_model_response(
3783-
response_id="chatcmpl-test",
3823+
response_id="chat-id",
37843824
response_model="gpt-3.5-turbo",
3785-
message_content="Test response",
3786-
created=1234567890,
3825+
message_content="the model response",
3826+
created=10000000,
3827+
usage=CompletionUsage(
3828+
prompt_tokens=20,
3829+
completion_tokens=10,
3830+
total_tokens=30,
3831+
),
37873832
)
37883833
)
37893834

0 commit comments

Comments
 (0)