Preview: openai.py
Size: 53.38 KB
/opt/hc_python/lib/python3.12/site-packages/sentry_sdk/integrations/openai.py
import json
import sys
import time
from collections.abc import Iterable
from functools import wraps
from typing import TYPE_CHECKING
import sentry_sdk
from sentry_sdk import consts
from sentry_sdk.ai._openai_completions_api import (
_get_system_instructions as _get_system_instructions_completions,
)
from sentry_sdk.ai._openai_completions_api import (
_get_text_items,
_transform_system_instructions,
)
from sentry_sdk.ai._openai_completions_api import (
_is_system_instruction as _is_system_instruction_completions,
)
from sentry_sdk.ai._openai_responses_api import (
_get_system_instructions as _get_system_instructions_responses,
)
from sentry_sdk.ai._openai_responses_api import (
_is_system_instruction as _is_system_instruction_responses,
)
from sentry_sdk.ai.monitoring import record_token_usage
from sentry_sdk.ai.utils import (
get_start_span_function,
normalize_message_roles,
set_data_normalized,
truncate_and_annotate_embedding_inputs,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import SPANDATA
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import (
has_span_streaming_enabled,
should_truncate_gen_ai_input,
)
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
reraise,
safe_serialize,
)
if TYPE_CHECKING:
from typing import (
Any,
AsyncIterator,
Callable,
Iterable,
Iterator,
List,
Optional,
Union,
)
from openai import Omit
from openai.types import CompletionUsage
from openai.types.responses import (
ResponseInputParam,
ResponseStreamEvent,
SequenceNotStr,
)
from openai.types.responses.response_usage import ResponseUsage
from sentry_sdk._types import TextPart
from sentry_sdk.tracing import Span
try:
try:
from openai import NotGiven
except ImportError:
NotGiven = None
try:
from openai import Omit
except ImportError:
Omit = None
from openai import AsyncStream, Stream
from openai.resources import AsyncEmbeddings, Embeddings
from openai.resources.chat.completions import AsyncCompletions, Completions
if TYPE_CHECKING:
from openai.types.chat import (
ChatCompletionChunk,
ChatCompletionMessageParam,
)
except ImportError:
raise DidNotEnable("OpenAI not installed")
RESPONSES_API_ENABLED = True
try:
# responses API support was introduced in v1.66.0
from openai.resources.responses import AsyncResponses, Responses
from openai.types.responses.response_completed_event import ResponseCompletedEvent
except ImportError:
RESPONSES_API_ENABLED = False
class OpenAIIntegration(Integration):
identifier = "openai"
origin = f"auto.ai.{identifier}"
def __init__(
self: "OpenAIIntegration",
include_prompts: bool = True,
tiktoken_encoding_name: "Optional[str]" = None,
) -> None:
self.include_prompts = include_prompts
self.tiktoken_encoding = None
if tiktoken_encoding_name is not None:
import tiktoken # type: ignore
self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name)
@staticmethod
def setup_once() -> None:
Completions.create = _wrap_chat_completion_create(Completions.create)
AsyncCompletions.create = _wrap_async_chat_completion_create(
AsyncCompletions.create
)
Embeddings.create = _wrap_embeddings_create(Embeddings.create)
AsyncEmbeddings.create = _wrap_async_embeddings_create(AsyncEmbeddings.create)
if RESPONSES_API_ENABLED:
Responses.create = _wrap_responses_create(Responses.create)
AsyncResponses.create = _wrap_async_responses_create(AsyncResponses.create)
def count_tokens(self: "OpenAIIntegration", s: str) -> int:
if self.tiktoken_encoding is None:
return 0
try:
return len(self.tiktoken_encoding.encode_ordinary(s))
except Exception:
return 0
def _capture_exception(exc: "Any") -> None:
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "openai", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _has_attr_and_is_int(
token_usage: "Union[CompletionUsage, ResponseUsage]", attr_name: str
) -> bool:
return hasattr(token_usage, attr_name) and isinstance(
getattr(token_usage, attr_name, None), int
)
def _calculate_completions_token_usage(
messages: "Optional[Iterable[ChatCompletionMessageParam]]",
response: "Any",
span: "Union[Span, StreamedSpan]",
streaming_message_responses: "Optional[List[str]]",
streaming_message_total_token_usage: "Optional[CompletionUsage]",
count_tokens: "Callable[..., Any]",
) -> None:
"""Extract and record token usage from a Chat Completions API response."""
input_tokens: "Optional[int]" = 0
input_tokens_cached: "Optional[int]" = 0
output_tokens: "Optional[int]" = 0
output_tokens_reasoning: "Optional[int]" = 0
total_tokens: "Optional[int]" = 0
usage = None
if streaming_message_total_token_usage is not None:
usage = streaming_message_total_token_usage
elif hasattr(response, "usage"):
usage = response.usage
if usage is not None:
if _has_attr_and_is_int(usage, "prompt_tokens"):
input_tokens = usage.prompt_tokens
if _has_attr_and_is_int(usage, "completion_tokens"):
output_tokens = usage.completion_tokens
if _has_attr_and_is_int(usage, "total_tokens"):
total_tokens = usage.total_tokens
if hasattr(usage, "prompt_tokens_details"):
cached = getattr(usage.prompt_tokens_details, "cached_tokens", None)
if isinstance(cached, int):
input_tokens_cached = cached
if hasattr(usage, "completion_tokens_details"):
reasoning = getattr(
usage.completion_tokens_details, "reasoning_tokens", None
)
if isinstance(reasoning, int):
output_tokens_reasoning = reasoning
# Manually count input tokens
if input_tokens == 0:
for message in messages or []:
if isinstance(message, str):
input_tokens += count_tokens(message)
continue
elif isinstance(message, dict):
message_content = message.get("content")
if message_content is None:
continue
text_items = _get_text_items(message_content)
input_tokens += sum(count_tokens(text) for text in text_items)
continue
# Manually count output tokens
if output_tokens == 0:
if streaming_message_responses is not None:
for message in streaming_message_responses:
output_tokens += count_tokens(message)
elif hasattr(response, "choices") and response.choices is not None:
for choice in response.choices:
if hasattr(choice, "message") and hasattr(choice.message, "content"):
output_tokens += count_tokens(choice.message.content)
# Do not set token data if it is 0
input_tokens = input_tokens or None
input_tokens_cached = input_tokens_cached or None
output_tokens = output_tokens or None
output_tokens_reasoning = output_tokens_reasoning or None
total_tokens = total_tokens or None
record_token_usage(
span,
input_tokens=input_tokens,
input_tokens_cached=input_tokens_cached,
output_tokens=output_tokens,
output_tokens_reasoning=output_tokens_reasoning,
total_tokens=total_tokens,
)
def _calculate_responses_token_usage(
input: "Any",
response: "Any",
span: "Union[Span, StreamedSpan]",
streaming_message_responses: "Optional[List[str]]",
count_tokens: "Callable[..., Any]",
) -> None:
"""Extract and record token usage from a Responses API response."""
input_tokens: "Optional[int]" = 0
input_tokens_cached: "Optional[int]" = 0
output_tokens: "Optional[int]" = 0
output_tokens_reasoning: "Optional[int]" = 0
total_tokens: "Optional[int]" = 0
if hasattr(response, "usage"):
usage = response.usage
if _has_attr_and_is_int(usage, "input_tokens"):
input_tokens = usage.input_tokens
if _has_attr_and_is_int(usage, "output_tokens"):
output_tokens = usage.output_tokens
if _has_attr_and_is_int(usage, "total_tokens"):
total_tokens = usage.total_tokens
if hasattr(usage, "input_tokens_details"):
cached = getattr(usage.input_tokens_details, "cached_tokens", None)
if isinstance(cached, int):
input_tokens_cached = cached
if hasattr(usage, "output_tokens_details"):
reasoning = getattr(usage.output_tokens_details, "reasoning_tokens", None)
if isinstance(reasoning, int):
output_tokens_reasoning = reasoning
# Manually count input tokens
if input_tokens == 0:
for message in input or []:
if isinstance(message, str):
input_tokens += count_tokens(message)
continue
elif isinstance(message, dict):
message_content = message.get("content")
if message_content is None:
continue
# Deliberate use of Completions function for both Completions and Responses input format.
text_items = _get_text_items(message_content)
input_tokens += sum(count_tokens(text) for text in text_items)
continue
# Manually count output tokens
if output_tokens == 0:
if streaming_message_responses is not None:
for message in streaming_message_responses:
output_tokens += count_tokens(message)
elif hasattr(response, "output"):
for output_item in response.output:
if hasattr(output_item, "content"):
for content_item in output_item.content:
if hasattr(content_item, "text"):
output_tokens += count_tokens(content_item.text)
# Do not set token data if it is 0
input_tokens = input_tokens or None
input_tokens_cached = input_tokens_cached or None
output_tokens = output_tokens or None
output_tokens_reasoning = output_tokens_reasoning or None
total_tokens = total_tokens or None
record_token_usage(
span,
input_tokens=input_tokens,
input_tokens_cached=input_tokens_cached,
output_tokens=output_tokens,
output_tokens_reasoning=output_tokens_reasoning,
total_tokens=total_tokens,
)
def _set_responses_api_input_data(
span: "Union[Span, StreamedSpan]",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
) -> None:
explicit_instructions: "Union[Optional[str], Omit]" = kwargs.get("instructions")
messages: "Optional[Union[str, ResponseInputParam]]" = kwargs.get("input")
tools = kwargs.get("tools")
if tools is not None and _is_given(tools) and len(tools) > 0:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
)
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
model = kwargs.get("model")
if model is not None:
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
max_tokens = kwargs.get("max_output_tokens")
if max_tokens is not None and _is_given(max_tokens):
set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
temperature = kwargs.get("temperature")
if temperature is not None and _is_given(temperature):
set_on_span(SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)
top_p = kwargs.get("top_p")
if top_p is not None and _is_given(top_p):
set_on_span(SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)
conversation = kwargs.get("conversation")
if conversation is not None and _is_given(conversation):
conversation_id: "Optional[str]" = None
if isinstance(conversation, str):
conversation_id = conversation
elif isinstance(conversation, dict):
conversation_id = conversation.get("id")
if conversation_id is not None:
set_on_span(SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id)
if not should_send_default_pii() or not integration.include_prompts:
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
return
if (
messages is None
and explicit_instructions is not None
and _is_given(explicit_instructions)
):
set_on_span(
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
json.dumps(
[
{
"type": "text",
"content": explicit_instructions,
}
]
),
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
return
if messages is None:
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
return
instructions_text_parts: "list[TextPart]" = []
if explicit_instructions is not None and _is_given(explicit_instructions):
instructions_text_parts.append(
{
"type": "text",
"content": explicit_instructions,
}
)
system_instructions = _get_system_instructions_responses(messages)
# Deliberate use of function accepting completions API type because
# of shared structure FOR THIS PURPOSE ONLY.
instructions_text_parts += _transform_system_instructions(system_instructions)
if len(instructions_text_parts) > 0:
set_on_span(
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
json.dumps(instructions_text_parts),
)
if isinstance(messages, str):
normalized_messages = normalize_message_roles([messages]) # type: ignore
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
return
non_system_messages = [
message for message in messages if not _is_system_instruction_responses(message)
]
if len(non_system_messages) > 0:
normalized_messages = normalize_message_roles(non_system_messages)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
def _set_completions_api_input_data(
span: "Union[Span, StreamedSpan]",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
) -> None:
messages: "Optional[Union[str, Iterable[ChatCompletionMessageParam]]]" = kwargs.get(
"messages"
)
tools = kwargs.get("tools")
if tools is not None and _is_given(tools) and len(tools) > 0:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
)
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
model = kwargs.get("model")
if model is not None:
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
max_tokens = kwargs.get("max_tokens")
if max_tokens is not None and _is_given(max_tokens):
set_on_span(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_on_span(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_on_span(SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty)
temperature = kwargs.get("temperature")
if temperature is not None and _is_given(temperature):
set_on_span(SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)
top_p = kwargs.get("top_p")
if top_p is not None and _is_given(top_p):
set_on_span(SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)
if (
not should_send_default_pii()
or not integration.include_prompts
or messages is None
):
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
return
if isinstance(messages, str):
normalized_messages = normalize_message_roles([messages]) # type: ignore
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
return
# dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197
if not isinstance(messages, Iterable) or isinstance(messages, dict):
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
return
messages = list(messages)
kwargs["messages"] = messages
system_instructions = _get_system_instructions_completions(messages)
if len(system_instructions) > 0:
set_on_span(
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
json.dumps(_transform_system_instructions(system_instructions)),
)
non_system_messages = [
message
for message in messages
if not _is_system_instruction_completions(message)
]
if len(non_system_messages) > 0:
normalized_messages = normalize_message_roles(non_system_messages)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
def _set_embeddings_input_data(
span: "Union[Span, StreamedSpan]",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
) -> None:
messages: "Union[str, SequenceNotStr[str], Iterable[int], Iterable[Iterable[int]]]" = kwargs.get(
"input"
)
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
model = kwargs.get("model")
if model is not None:
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
if (
not should_send_default_pii()
or not integration.include_prompts
or messages is None
):
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
return
if isinstance(messages, str):
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
normalized_messages = normalize_message_roles([messages]) # type: ignore
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_embedding_inputs(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False
)
return
# dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197
if not isinstance(messages, Iterable) or isinstance(messages, dict):
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
return
messages = list(messages)
kwargs["input"] = messages
if len(messages) > 0:
normalized_messages = normalize_message_roles(messages)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_embedding_inputs(normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False
)
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
def _set_common_output_data(
span: "Union[Span, StreamedSpan]",
response: "Any",
input: "Any",
integration: "OpenAIIntegration",
finish_span: bool = True,
) -> None:
if hasattr(response, "model"):
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model)
# Chat Completions API
if hasattr(response, "choices") and response.choices is not None:
if should_send_default_pii() and integration.include_prompts:
response_text = [
choice.message.model_dump()
for choice in response.choices
if choice.message is not None
]
if len(response_text) > 0:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_text)
_calculate_completions_token_usage(
messages=input,
response=response,
span=span,
streaming_message_responses=None,
streaming_message_total_token_usage=None,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
# Responses API
elif hasattr(response, "output"):
if should_send_default_pii() and integration.include_prompts:
output_messages: "dict[str, list[Any]]" = {
"response": [],
"tool": [],
}
for output in response.output:
if output.type == "function_call":
output_messages["tool"].append(output.dict())
elif output.type == "message":
for output_message in output.content:
try:
output_messages["response"].append(output_message.text)
except AttributeError:
# Unknown output message type, just return the json
output_messages["response"].append(output_message.dict())
if len(output_messages["tool"]) > 0:
set_data_normalized(
span,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
output_messages["tool"],
unpack=False,
)
if len(output_messages["response"]) > 0:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
)
_calculate_responses_token_usage(
input=input,
response=response,
span=span,
streaming_message_responses=None,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
# Embeddings API (fallback for responses with neither choices nor output)
else:
_calculate_completions_token_usage(
messages=input,
response=response,
span=span,
streaming_message_responses=None,
streaming_message_total_token_usage=None,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
def _new_sync_chat_completion(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return f(*args, **kwargs)
if "messages" not in kwargs:
# invalid call (in all versions of openai), let it return error
return f(*args, **kwargs)
try:
iter(kwargs["messages"])
except TypeError:
# invalid call (in all versions), messages must be iterable
return f(*args, **kwargs)
model = kwargs.get("model")
# Same bool handling as in https://github.com/openai/openai-python/blob/acd0c54d8a68efeedde0e5b4e6c310eef1ce7867/src/openai/resources/completions.py#L585
is_streaming_response = kwargs.get("stream", False) or False
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name=f"chat {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_CHAT,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
SPANDATA.GEN_AI_RESPONSE_STREAMING: is_streaming_response,
},
)
else:
span = get_start_span_function()(
op=consts.OP.GEN_AI_CHAT,
name=f"chat {model}",
origin=OpenAIIntegration.origin,
)
span.__enter__()
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming_response)
_set_completions_api_input_data(span, kwargs, integration)
start_time = time.perf_counter()
try:
response = f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
span.__exit__(*exc_info)
reraise(*exc_info)
# Attribute check to fail gracefully if the attribute is not present in future `openai` versions.
if isinstance(response, Stream) and hasattr(response, "_iterator"):
messages = kwargs.get("messages")
if messages is not None and isinstance(messages, str):
messages = [messages]
response._iterator = _wrap_synchronous_completions_chunk_iterator(
span=span,
integration=integration,
start_time=start_time,
messages=messages,
response=response,
old_iterator=response._iterator,
finish_span=True,
)
else:
_set_completions_api_output_data(
span, response, kwargs, integration, finish_span=True
)
return response
async def _new_async_chat_completion(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return await f(*args, **kwargs)
if "messages" not in kwargs:
# invalid call (in all versions of openai), let it return error
return await f(*args, **kwargs)
try:
iter(kwargs["messages"])
except TypeError:
# invalid call (in all versions), messages must be iterable
return await f(*args, **kwargs)
model = kwargs.get("model")
# Same bool handling as in https://github.com/openai/openai-python/blob/acd0c54d8a68efeedde0e5b4e6c310eef1ce7867/src/openai/resources/completions.py#L585
is_streaming_response = kwargs.get("stream", False) or False
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name=f"chat {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_CHAT,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
SPANDATA.GEN_AI_RESPONSE_STREAMING: is_streaming_response,
},
)
else:
span = get_start_span_function()(
op=consts.OP.GEN_AI_CHAT,
name=f"chat {model}",
origin=OpenAIIntegration.origin,
)
span.__enter__()
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming_response)
_set_completions_api_input_data(span, kwargs, integration)
start_time = time.perf_counter()
try:
response = await f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
span.__exit__(*exc_info)
reraise(*exc_info)
# Attribute check to fail gracefully if the attribute is not present in future `openai` versions.
if isinstance(response, AsyncStream) and hasattr(response, "_iterator"):
messages = kwargs.get("messages")
if messages is not None and isinstance(messages, str):
messages = [messages]
response._iterator = _wrap_asynchronous_completions_chunk_iterator(
span=span,
integration=integration,
start_time=start_time,
messages=messages,
response=response,
old_iterator=response._iterator,
finish_span=True,
)
else:
_set_completions_api_output_data(
span, response, kwargs, integration, finish_span=True
)
return response
def _set_completions_api_output_data(
span: "Union[Span, StreamedSpan]",
response: "Any",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
finish_span: bool = True,
) -> None:
messages = kwargs.get("messages")
if messages is not None and isinstance(messages, str):
messages = [messages]
_set_common_output_data(
span,
response,
messages,
integration,
finish_span,
)
def _wrap_synchronous_completions_chunk_iterator(
span: "Union[Span, StreamedSpan]",
integration: "OpenAIIntegration",
start_time: "Optional[float]",
messages: "Optional[Iterable[ChatCompletionMessageParam]]",
response: "Stream[ChatCompletionChunk]",
old_iterator: "Iterator[ChatCompletionChunk]",
finish_span: "bool",
) -> "Iterator[ChatCompletionChunk]":
"""
Sets information received while iterating the response stream on the AI Client Span.
Compute token count based on inputs and outputs using tiktoken if token counts are not in the model response.
Responsible for closing the AI Client Span if instructed to by the `finish_span` argument.
"""
ttft = None
data_buf: "list[list[str]]" = [] # one for each choice
streaming_message_total_token_usage = None
for x in old_iterator:
if isinstance(span, StreamedSpan):
span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model)
else:
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model)
with capture_internal_exceptions():
if hasattr(x, "choices") and x.choices is not None:
choice_index = 0
for choice in x.choices:
if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
if start_time is not None and ttft is None:
ttft = time.perf_counter() - start_time
content = choice.delta.content
if len(data_buf) <= choice_index:
data_buf.append([])
data_buf[choice_index].append(content or "")
choice_index += 1
if hasattr(x, "usage"):
streaming_message_total_token_usage = x.usage
yield x
with capture_internal_exceptions():
if ttft is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft
)
all_responses = None
if len(data_buf) > 0:
all_responses = ["".join(chunk) for chunk in data_buf]
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses)
_calculate_completions_token_usage(
messages=messages,
response=response,
span=span,
streaming_message_responses=all_responses,
streaming_message_total_token_usage=streaming_message_total_token_usage,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
async def _wrap_asynchronous_completions_chunk_iterator(
span: "Union[Span, StreamedSpan]",
integration: "OpenAIIntegration",
start_time: "Optional[float]",
messages: "Optional[Iterable[ChatCompletionMessageParam]]",
response: "AsyncStream[ChatCompletionChunk]",
old_iterator: "AsyncIterator[ChatCompletionChunk]",
finish_span: "bool",
) -> "AsyncIterator[ChatCompletionChunk]":
"""
Sets information received while iterating the response stream on the AI Client Span.
Compute token count based on inputs and outputs using tiktoken if token counts are not in the model response.
Responsible for closing the AI Client Span if instructed to by the `finish_span` argument.
"""
ttft = None
data_buf: "list[list[str]]" = [] # one for each choice
streaming_message_total_token_usage = None
async for x in old_iterator:
if isinstance(span, StreamedSpan):
span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model)
else:
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model)
with capture_internal_exceptions():
if hasattr(x, "choices") and x.choices is not None:
choice_index = 0
for choice in x.choices:
if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
if start_time is not None and ttft is None:
ttft = time.perf_counter() - start_time
content = choice.delta.content
if len(data_buf) <= choice_index:
data_buf.append([])
data_buf[choice_index].append(content or "")
choice_index += 1
if hasattr(x, "usage"):
streaming_message_total_token_usage = x.usage
yield x
with capture_internal_exceptions():
if ttft is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft
)
all_responses = None
if len(data_buf) > 0:
all_responses = ["".join(chunk) for chunk in data_buf]
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses)
_calculate_completions_token_usage(
messages=messages,
response=response,
span=span,
streaming_message_responses=all_responses,
streaming_message_total_token_usage=streaming_message_total_token_usage,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
def _wrap_synchronous_responses_event_iterator(
span: "Union[Span, StreamedSpan]",
integration: "OpenAIIntegration",
start_time: "Optional[float]",
input: "Optional[Union[str, ResponseInputParam]]",
response: "Stream[ResponseStreamEvent]",
old_iterator: "Iterator[ResponseStreamEvent]",
finish_span: "bool",
) -> "Iterator[ResponseStreamEvent]":
"""
Sets information received while iterating the response stream on the AI Client Span.
Compute token count based on inputs and outputs using tiktoken if token counts are not in the model response.
Responsible for closing the AI Client Span if instructed to by the `finish_span` argument.
"""
ttft = None
data_buf: "list[list[str]]" = [] # one for each choice
count_tokens_manually = True
for x in old_iterator:
with capture_internal_exceptions():
if hasattr(x, "delta"):
if start_time is not None and ttft is None:
ttft = time.perf_counter() - start_time
if len(data_buf) == 0:
data_buf.append([])
data_buf[0].append(x.delta or "")
if isinstance(x, ResponseCompletedEvent):
if isinstance(span, StreamedSpan):
span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model)
else:
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model)
_calculate_responses_token_usage(
input=input,
response=x.response,
span=span,
streaming_message_responses=None,
count_tokens=integration.count_tokens,
)
count_tokens_manually = False
yield x
with capture_internal_exceptions():
if ttft is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft
)
if len(data_buf) > 0:
all_responses = ["".join(chunk) for chunk in data_buf]
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses)
if count_tokens_manually:
_calculate_responses_token_usage(
input=input,
response=response,
span=span,
streaming_message_responses=all_responses,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
async def _wrap_asynchronous_responses_event_iterator(
span: "Union[Span, StreamedSpan]",
integration: "OpenAIIntegration",
start_time: "Optional[float]",
input: "Optional[Union[str, ResponseInputParam]]",
response: "AsyncStream[ResponseStreamEvent]",
old_iterator: "AsyncIterator[ResponseStreamEvent]",
finish_span: "bool",
) -> "AsyncIterator[ResponseStreamEvent]":
"""
Sets information received while iterating the response stream on the AI Client Span.
Compute token count based on inputs and outputs using tiktoken if token counts are not in the model response.
Responsible for closing the AI Client Span if instructed to by the `finish_span` argument.
"""
ttft: "Optional[float]" = None
data_buf: "list[list[str]]" = [] # one for each choice
count_tokens_manually = True
async for x in old_iterator:
with capture_internal_exceptions():
if hasattr(x, "delta"):
if start_time is not None and ttft is None:
ttft = time.perf_counter() - start_time
if len(data_buf) == 0:
data_buf.append([])
data_buf[0].append(x.delta or "")
if isinstance(x, ResponseCompletedEvent):
if isinstance(span, StreamedSpan):
span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model)
else:
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model)
_calculate_responses_token_usage(
input=input,
response=x.response,
span=span,
streaming_message_responses=None,
count_tokens=integration.count_tokens,
)
count_tokens_manually = False
yield x
with capture_internal_exceptions():
if ttft is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft
)
if len(data_buf) > 0:
all_responses = ["".join(chunk) for chunk in data_buf]
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses)
if count_tokens_manually:
_calculate_responses_token_usage(
input=input,
response=response,
span=span,
streaming_message_responses=all_responses,
count_tokens=integration.count_tokens,
)
if finish_span:
span.__exit__(None, None, None)
def _set_responses_api_output_data(
span: "Union[Span, StreamedSpan]",
response: "Any",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
finish_span: bool = True,
) -> None:
input = kwargs.get("input")
if input is not None and isinstance(input, str):
input = [input]
_set_common_output_data(
span,
response,
input,
integration,
finish_span,
)
def _set_embeddings_output_data(
span: "Union[Span, StreamedSpan]",
response: "Any",
kwargs: "dict[str, Any]",
integration: "OpenAIIntegration",
finish_span: bool = True,
) -> None:
input = kwargs.get("input")
if input is not None and isinstance(input, str):
input = [input]
_set_common_output_data(
span,
response,
input,
integration,
finish_span,
)
def _wrap_chat_completion_create(f: "Callable[..., Any]") -> "Callable[..., Any]":
@wraps(f)
def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None or "messages" not in kwargs:
# no "messages" means invalid call (in all versions of openai), let it return error
return f(*args, **kwargs)
return _new_sync_chat_completion(f, *args, **kwargs)
return _sentry_patched_create_sync
def _wrap_async_chat_completion_create(f: "Callable[..., Any]") -> "Callable[..., Any]":
@wraps(f)
async def _sentry_patched_create_async(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None or "messages" not in kwargs:
# no "messages" means invalid call (in all versions of openai), let it return error
return await f(*args, **kwargs)
return await _new_async_chat_completion(f, *args, **kwargs)
return _sentry_patched_create_async
def _new_sync_embeddings_create(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return f(*args, **kwargs)
model = kwargs.get("model")
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=f"embeddings {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_EMBEDDINGS,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
},
) as span:
_set_embeddings_input_data(span, kwargs, integration)
try:
response = f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
reraise(*exc_info)
_set_embeddings_output_data(
span, response, kwargs, integration, finish_span=False
)
return response
else:
with get_start_span_function()(
op=consts.OP.GEN_AI_EMBEDDINGS,
name=f"embeddings {model}",
origin=OpenAIIntegration.origin,
) as span:
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
_set_embeddings_input_data(span, kwargs, integration)
try:
response = f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
reraise(*exc_info)
_set_embeddings_output_data(
span, response, kwargs, integration, finish_span=False
)
return response
async def _new_async_embeddings_create(
f: "Any", *args: "Any", **kwargs: "Any"
) -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return await f(*args, **kwargs)
model = kwargs.get("model")
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=f"embeddings {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_EMBEDDINGS,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
},
) as span:
_set_embeddings_input_data(span, kwargs, integration)
try:
response = await f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
reraise(*exc_info)
_set_embeddings_output_data(
span, response, kwargs, integration, finish_span=False
)
return response
else:
with get_start_span_function()(
op=consts.OP.GEN_AI_EMBEDDINGS,
name=f"embeddings {model}",
origin=OpenAIIntegration.origin,
) as span:
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
_set_embeddings_input_data(span, kwargs, integration)
try:
response = await f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
reraise(*exc_info)
_set_embeddings_output_data(
span, response, kwargs, integration, finish_span=False
)
return response
def _wrap_embeddings_create(f: "Any") -> "Any":
@wraps(f)
def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None:
return f(*args, **kwargs)
return _new_sync_embeddings_create(f, *args, **kwargs)
return _sentry_patched_create_sync
def _wrap_async_embeddings_create(f: "Any") -> "Any":
@wraps(f)
async def _sentry_patched_create_async(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None:
return await f(*args, **kwargs)
return await _new_async_embeddings_create(f, *args, **kwargs)
return _sentry_patched_create_async
def _new_sync_responses_create(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return f(*args, **kwargs)
model = kwargs.get("model")
# Same bool handling as in https://github.com/openai/openai-python/blob/acd0c54d8a68efeedde0e5b4e6c310eef1ce7867/src/openai/resources/responses/responses.py#L940
is_streaming_response = kwargs.get("stream", False) or False
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name=f"responses {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_RESPONSES,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
SPANDATA.GEN_AI_RESPONSE_STREAMING: is_streaming_response,
},
)
else:
span = get_start_span_function()(
op=consts.OP.GEN_AI_RESPONSES,
name=f"responses {model}",
origin=OpenAIIntegration.origin,
)
span.__enter__()
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming_response)
_set_responses_api_input_data(span, kwargs, integration)
start_time = time.perf_counter()
try:
response = f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
span.__exit__(*exc_info)
reraise(*exc_info)
# Attribute check to fail gracefully if the attribute is not present in future `openai` versions.
if isinstance(response, Stream) and hasattr(response, "_iterator"):
input = kwargs.get("input")
if input is not None and isinstance(input, str):
input = [input]
response._iterator = _wrap_synchronous_responses_event_iterator(
span=span,
integration=integration,
start_time=start_time,
input=input,
response=response,
old_iterator=response._iterator,
finish_span=True,
)
else:
_set_responses_api_output_data(
span, response, kwargs, integration, finish_span=True
)
return response
async def _new_async_responses_create(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(OpenAIIntegration)
if integration is None:
return await f(*args, **kwargs)
model = kwargs.get("model")
# Same bool handling as in https://github.com/openai/openai-python/blob/acd0c54d8a68efeedde0e5b4e6c310eef1ce7867/src/openai/resources/responses/responses.py#L940
is_streaming_response = kwargs.get("stream", False) or False
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name=f"responses {model}",
attributes={
"sentry.op": consts.OP.GEN_AI_RESPONSES,
"sentry.origin": OpenAIIntegration.origin,
SPANDATA.GEN_AI_SYSTEM: "openai",
SPANDATA.GEN_AI_RESPONSE_STREAMING: is_streaming_response,
},
)
else:
span = get_start_span_function()(
op=consts.OP.GEN_AI_RESPONSES,
name=f"responses {model}",
origin=OpenAIIntegration.origin,
)
span.__enter__()
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming_response)
_set_responses_api_input_data(span, kwargs, integration)
start_time = time.perf_counter()
try:
response = await f(*args, **kwargs)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)
span.__exit__(*exc_info)
reraise(*exc_info)
# Attribute check to fail gracefully if the attribute is not present in future `openai` versions.
if isinstance(response, AsyncStream) and hasattr(response, "_iterator"):
input = kwargs.get("input")
if input is not None and isinstance(input, str):
input = [input]
response._iterator = _wrap_asynchronous_responses_event_iterator(
span=span,
integration=integration,
start_time=start_time,
input=input,
response=response,
old_iterator=response._iterator,
finish_span=True,
)
else:
_set_responses_api_output_data(
span, response, kwargs, integration, finish_span=True
)
return response
def _wrap_responses_create(f: "Any") -> "Any":
@wraps(f)
def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None:
return f(*args, **kwargs)
return _new_sync_responses_create(f, *args, **kwargs)
return _sentry_patched_create_sync
def _wrap_async_responses_create(f: "Any") -> "Any":
@wraps(f)
async def _sentry_patched_responses_async(*args: "Any", **kwargs: "Any") -> "Any":
integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
if integration is None:
return await f(*args, **kwargs)
return await _new_async_responses_create(f, *args, **kwargs)
return _sentry_patched_responses_async
def _is_given(obj: "Any") -> bool:
"""
Check for givenness safely across different openai versions.
"""
if NotGiven is not None and isinstance(obj, NotGiven):
return False
if Omit is not None and isinstance(obj, Omit):
return False
return True
Directory Contents
Dirs: 10 × Files: 73