diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index d06bfecbc0..0f342019fe 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -16,6 +16,7 @@ from openai import AsyncOpenAI, OpenAI, AsyncStream, Stream, OpenAIError from openai.types import CompletionUsage, CreateEmbeddingResponse, Embedding +from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionChunk from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import ChoiceDelta, Choice as DeltaChoice @@ -489,7 +490,12 @@ def tiktoken_encoding_if_installed(): ], ) def test_streaming_chat_completion_no_prompts( - sentry_init, capture_events, send_default_pii, include_prompts + sentry_init, + capture_events, + send_default_pii, + include_prompts, + get_model_response, + server_side_event_chunks, ): sentry_init( integrations=[ @@ -504,62 +510,97 @@ def test_streaming_chat_completion_no_prompts( events = capture_events() client = OpenAI(api_key="z") - returned_stream = Stream(cast_to=None, response=None, client=client) - returned_stream._iterator = [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, delta=ChoiceDelta(content="world"), finish_reason="stop" - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] - - client.chat.completions._post = mock.Mock(return_value=returned_stream) - with start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( - model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, + returned_stream = get_model_response( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetails( + audio_tokens=10, + cached_tokens=20, + ), + completion_tokens_details=CompletionTokensDetails( + reasoning_tokens=5, + audio_tokens=3, + accepted_prediction_tokens=7, + rejected_prediction_tokens=2, + ), + ), + ), ], - stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + include_event_type=False, ) - response_string = "".join( - map(lambda x: x.choices[0].delta.content, response_stream) - ) - assert response_string == "hello world" + ) + + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ): + with start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) + for _ in response_stream: + pass + tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] @@ -576,6 +617,10 @@ def test_streaming_chat_completion_no_prompts( assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" + assert span["data"]["gen_ai.usage.output_tokens"] == 2 + assert span["data"]["gen_ai.usage.input_tokens"] == 7 + assert span["data"]["gen_ai.usage.total_tokens"] == 9 + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["data"] assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] @@ -634,7 +679,14 @@ def test_streaming_chat_completion_no_prompts( ), ], ) -def test_streaming_chat_completion(sentry_init, capture_events, messages, request): +def test_streaming_chat_completion( + sentry_init, + capture_events, + messages, + request, + get_model_response, + server_side_event_chunks, +): sentry_init( integrations=[ OpenAIIntegration( @@ -648,59 +700,84 @@ def test_streaming_chat_completion(sentry_init, capture_events, messages, reques events = capture_events() client = OpenAI(api_key="z") - returned_stream = Stream(cast_to=None, response=None, client=client) - returned_stream._iterator = [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, delta=ChoiceDelta(content="world"), finish_reason="stop" - ) + returned_stream = get_model_response( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + ), + ), ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] - - client.chat.completions._post = mock.Mock(return_value=returned_stream) - with start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( - model="some-model", - messages=messages, - stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, - ) - response_string = "".join( - map(lambda x: x.choices[0].delta.content, response_stream) + include_event_type=False, ) - assert response_string == "hello world" + ) + + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ): + with start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", + messages=messages, + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) + for _ in response_stream: + pass + tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] @@ -766,7 +843,13 @@ def test_streaming_chat_completion(sentry_init, capture_events, messages, reques ], ) async def test_streaming_chat_completion_async_no_prompts( - sentry_init, capture_events, send_default_pii, include_prompts, async_iterator + sentry_init, + capture_events, + send_default_pii, + include_prompts, + get_model_response, + async_iterator, + server_side_event_chunks, ): sentry_init( integrations=[ @@ -781,66 +864,78 @@ async def test_streaming_chat_completion_async_no_prompts( events = capture_events() client = AsyncOpenAI(api_key="z") - returned_stream = AsyncStream(cast_to=None, response=None, client=client) - returned_stream._iterator = async_iterator( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] + include_event_type=False, + ) + ) ) - client.chat.completions._post = AsyncMock(return_value=returned_stream) - with start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], - stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, - ) + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ): + with start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content assert response_string == "hello world" tx = events[0] @@ -920,7 +1015,13 @@ async def test_streaming_chat_completion_async_no_prompts( ], ) async def test_streaming_chat_completion_async( - sentry_init, capture_events, messages, request, async_iterator + sentry_init, + capture_events, + messages, + request, + get_model_response, + async_iterator, + server_side_event_chunks, ): sentry_init( integrations=[ @@ -935,63 +1036,76 @@ async def test_streaming_chat_completion_async( events = capture_events() client = AsyncOpenAI(api_key="z") - returned_stream = AsyncStream(cast_to=None, response=None, client=client) - returned_stream._iterator = async_iterator( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) + + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] + include_event_type=False, + ) + ) ) - client.chat.completions._post = AsyncMock(return_value=returned_stream) - with start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=messages, - stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, - ) + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ): + with start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=messages, + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content assert response_string == "hello world" tx = events[0]