Skip to content
20 changes: 8 additions & 12 deletions sentry_sdk/integrations/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,36 +356,32 @@ def _set_completions_api_input_data(
)

model = kwargs.get("model")
if model is not None and _is_given(model):
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model)
if model is not None:
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)

stream = kwargs.get("stream")
if stream is not None and _is_given(stream):
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_STREAMING, stream)
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, stream)

max_tokens = kwargs.get("max_tokens")
if max_tokens is not None and _is_given(max_tokens):
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
span.set_data(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)

presence_penalty = kwargs.get("presence_penalty")
if presence_penalty is not None and _is_given(presence_penalty):
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, presence_penalty
)
span.set_data(SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, presence_penalty)

frequency_penalty = kwargs.get("frequency_penalty")
if frequency_penalty is not None and _is_given(frequency_penalty):
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty
)
span.set_data(SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty)

temperature = kwargs.get("temperature")
if temperature is not None and _is_given(temperature):
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)
span.set_data(SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)

top_p = kwargs.get("top_p")
if top_p is not None and _is_given(top_p):
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)
span.set_data(SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)

if (
not should_send_default_pii()
Expand Down
100 changes: 100 additions & 0 deletions tests/integrations/openai/test_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,11 @@ def test_nonstreaming_chat_completion_no_prompts(
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "hello"},
],
max_tokens=100,
presence_penalty=0.1,
frequency_penalty=0.2,
temperature=0.7,
top_p=0.9,
)
.choices[0]
.message.content
Expand All @@ -163,6 +168,13 @@ def test_nonstreaming_chat_completion_no_prompts(
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.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"]
Expand Down Expand Up @@ -231,6 +243,11 @@ def test_nonstreaming_chat_completion(sentry_init, capture_events, messages, req
client.chat.completions.create(
model="some-model",
messages=messages,
max_tokens=100,
presence_penalty=0.1,
frequency_penalty=0.2,
temperature=0.7,
top_p=0.9,
)
.choices[0]
.message.content
Expand All @@ -243,6 +260,13 @@ def test_nonstreaming_chat_completion(sentry_init, capture_events, messages, req
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

param_id = request.node.callspec.id
if "blocks" in param_id:
assert json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == [
Expand Down Expand Up @@ -300,6 +324,11 @@ async def test_nonstreaming_chat_completion_async_no_prompts(
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "hello"},
],
max_tokens=100,
presence_penalty=0.1,
frequency_penalty=0.2,
temperature=0.7,
top_p=0.9,
)
response = response.choices[0].message.content

Expand All @@ -310,6 +339,13 @@ async def test_nonstreaming_chat_completion_async_no_prompts(
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.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"]
Expand Down Expand Up @@ -380,6 +416,11 @@ async def test_nonstreaming_chat_completion_async(
response = await client.chat.completions.create(
model="some-model",
messages=messages,
max_tokens=100,
presence_penalty=0.1,
frequency_penalty=0.2,
temperature=0.7,
top_p=0.9,
)
response = response.choices[0].message.content

Expand All @@ -390,6 +431,13 @@ async def test_nonstreaming_chat_completion_async(
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

param_id = request.node.callspec.id
if "blocks" in param_id:
assert json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == [
Expand Down Expand Up @@ -498,6 +546,11 @@ def test_streaming_chat_completion_no_prompts(
{"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 = "".join(
map(lambda x: x.choices[0].delta.content, response_stream)
Expand All @@ -508,6 +561,14 @@ def test_streaming_chat_completion_no_prompts(
span = tx["spans"][0]
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id"

Expand Down Expand Up @@ -626,6 +687,11 @@ def test_streaming_chat_completion(sentry_init, capture_events, messages, reques
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)
Expand All @@ -636,6 +702,14 @@ def test_streaming_chat_completion(sentry_init, capture_events, messages, reques
span = tx["spans"][0]
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

param_id = request.node.callspec.id
if "blocks" in param_id:
Expand Down Expand Up @@ -753,6 +827,11 @@ async def test_streaming_chat_completion_async_no_prompts(
{"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 = ""
Expand All @@ -765,6 +844,14 @@ async def test_streaming_chat_completion_async_no_prompts(
span = tx["spans"][0]
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id"

Expand Down Expand Up @@ -891,6 +978,11 @@ async def test_streaming_chat_completion_async(
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 = ""
Expand All @@ -903,6 +995,14 @@ async def test_streaming_chat_completion_async(
span = tx["spans"][0]
assert span["op"] == "gen_ai.chat"
assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai"
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True

assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9

assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id"

Expand Down
Loading