Skip to content

Commit 3e7d78f

Browse files
colin-sentrycolin-sentry
committed
OpenAI integration
1 parent b96f03d commit 3e7d78f

File tree

4 files changed

+382
-0
lines changed

4 files changed

+382
-0
lines changed

‎sentry_sdk/integrations/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
7878
"sentry_sdk.integrations.fastapi.FastApiIntegration",
7979
"sentry_sdk.integrations.flask.FlaskIntegration",
8080
"sentry_sdk.integrations.httpx.HttpxIntegration",
81+
"sentry_sdk.integrations.openai.OpenAIIntegration",
8182
"sentry_sdk.integrations.pyramid.PyramidIntegration",
8283
"sentry_sdk.integrations.redis.RedisIntegration",
8384
"sentry_sdk.integrations.rq.RqIntegration",

‎sentry_sdk/integrations/openai.py

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
from __future__ import absolute_import
2+
3+
from sentry_sdk._types import TYPE_CHECKING
4+
5+
if TYPE_CHECKING:
6+
from typing import Iterator, Any, TypeVar, Callable
7+
8+
F = TypeVar("F", bound=Callable[..., Any])
9+
10+
from sentry_sdk._functools import wraps
11+
from sentry_sdk.hub import Hub
12+
from sentry_sdk.integrations import DidNotEnable, Integration
13+
from sentry_sdk.utils import logger, capture_internal_exceptions
14+
15+
try:
16+
from openai.types.chat import ChatCompletionChunk
17+
from openai.resources.chat.completions import Completions
18+
from openai.resources import Embeddings
19+
except ImportError:
20+
raise DidNotEnable("OpenAI not installed")
21+
22+
try:
23+
import tiktoken
24+
25+
enc = tiktoken.get_encoding("cl100k_base")
26+
27+
def count_tokens(s):
28+
# type: (str) -> int
29+
return len(enc.encode_ordinary(s))
30+
31+
logger.debug("[OpenAI] using tiktoken to count tokens")
32+
except ImportError:
33+
logger.info(
34+
"The Sentry Python SDK requires 'tiktoken' in order to measure token usage from some OpenAI APIs"
35+
"Please install 'tiktoken' if you aren't receiving token usage in Sentry."
36+
"See https://docs.sentry.io/platforms/python/guides/openai/ for more information."
37+
)
38+
39+
def count_tokens(s):
40+
# type: (str) -> int
41+
return 0
42+
43+
44+
COMPLETION_TOKENS = "completion_tоkens"
45+
PROMPT_TOKENS = "prompt_tоkens"
46+
TOTAL_TOKENS = "total_tоkens"
47+
48+
49+
class OpenAIIntegration(Integration):
50+
identifier = "openai"
51+
52+
@staticmethod
53+
def setup_once():
54+
# TODO minimum version
55+
Completions.create = _wrap_chat_completion_create(Completions.create)
56+
Embeddings.create = _wrap_enbeddings_create(Embeddings.create)
57+
58+
59+
def _calculate_chat_completion_usage(
60+
messages, response, span, streaming_message_responses=None
61+
):
62+
completion_tokens = 0
63+
prompt_tokens = 0
64+
total_tokens = 0
65+
if hasattr(response, "usage"):
66+
if hasattr(response.usage, "completion_tokens") and isinstance(
67+
response.usage.completion_tokens, int
68+
):
69+
completion_tokens = response.usage.completion_tokens
70+
if hasattr(response.usage, "prompt_tokens") and isinstance(
71+
response.usage.prompt_tokens, int
72+
):
73+
prompt_tokens = response.usage.prompt_tokens
74+
if hasattr(response.usage, "total_tokens") and isinstance(
75+
response.usage.total_tokens, int
76+
):
77+
total_tokens = response.usage.total_tokens
78+
79+
if prompt_tokens == 0:
80+
for message in messages:
81+
if hasattr(message, "content"):
82+
prompt_tokens += count_tokens(message.content)
83+
elif "content" in message:
84+
prompt_tokens += count_tokens(message["content"])
85+
86+
if completion_tokens == 0:
87+
if streaming_message_responses is not None:
88+
for message in streaming_message_responses:
89+
completion_tokens += count_tokens(message)
90+
elif hasattr(response, "choices"):
91+
for choice in response.choices:
92+
if hasattr(choice, "message"):
93+
completion_tokens += count_tokens(choice.message)
94+
95+
if total_tokens == 0:
96+
total_tokens = prompt_tokens + completion_tokens
97+
98+
if completion_tokens != 0:
99+
span.set_data(COMPLETION_TOKENS, completion_tokens)
100+
if prompt_tokens != 0:
101+
span.set_data(PROMPT_TOKENS, prompt_tokens)
102+
if total_tokens != 0:
103+
span.set_data(TOTAL_TOKENS, total_tokens)
104+
105+
106+
def _wrap_chat_completion_create(f):
107+
# type: (F) -> F
108+
@wraps(f)
109+
def new_chat_completion(*args, **kwargs):
110+
# type: (*Any, **Any) -> Any
111+
hub = Hub.current
112+
integration = hub.get_integration(OpenAIIntegration)
113+
if integration is None:
114+
return f(*args, **kwargs)
115+
116+
if "messages" not in kwargs:
117+
# invalid call (in all versions of openai), let it return error
118+
return f(*args, **kwargs)
119+
120+
try:
121+
iter(kwargs["messages"])
122+
except TypeError:
123+
# invalid call (in all versions), messages must be iterable
124+
return f(*args, **kwargs)
125+
126+
kwargs["messages"] = list(kwargs["messages"])
127+
messages = kwargs["messages"]
128+
model = kwargs.get("model")
129+
streaming = kwargs.get("stream") # TODO handle streaming
130+
131+
span = hub.start_span(op="openai", description="Chat Completion")
132+
span.__enter__()
133+
res = f(*args, **kwargs)
134+
with capture_internal_exceptions():
135+
span.set_data("messages", messages)
136+
span.set_tag("model", model)
137+
span.set_tag("streaming", streaming)
138+
139+
if hasattr(res, "choices"):
140+
span.set_data("response", res.choices[0].message)
141+
_calculate_chat_completion_usage(messages, res, span)
142+
span.__exit__(None, None, None)
143+
elif hasattr(res, "_iterator"):
144+
data_buf: list[list[str]] = [] # one for each choice
145+
146+
old_iterator: Iterator[ChatCompletionChunk] = res._iterator
147+
148+
def new_iterator() -> Iterator[ChatCompletionChunk]:
149+
with capture_internal_exceptions():
150+
for x in old_iterator:
151+
if hasattr(x, "choices"):
152+
choice_index = 0
153+
for choice in x.choices:
154+
if hasattr(choice, "delta") and hasattr(
155+
choice.delta, "content"
156+
):
157+
content = choice.delta.content
158+
if len(data_buf) <= choice_index:
159+
data_buf.append([])
160+
data_buf[choice_index].append(content or "")
161+
choice_index += 1
162+
yield x
163+
if len(data_buf) > 0:
164+
all_responses = list(
165+
map(lambda chunk: "".join(chunk), data_buf)
166+
)
167+
span.set_data("responses", all_responses)
168+
_calculate_chat_completion_usage(
169+
messages, res, span, all_responses
170+
)
171+
span.__exit__(None, None, None)
172+
173+
res._iterator = new_iterator()
174+
else:
175+
span.set_tag("unknown_response", True)
176+
span.__exit__(None, None, None)
177+
return res
178+
179+
return new_chat_completion
180+
181+
182+
def _wrap_enbeddings_create(f):
183+
# type: (F) -> F
184+
185+
@wraps(f)
186+
def new_embeddings_create(*args, **kwargs):
187+
hub = Hub.current
188+
integration = hub.get_integration(OpenAIIntegration)
189+
if integration is None:
190+
return f(*args, **kwargs)
191+
192+
with hub.start_span(op="openai", description="Embeddings Creation") as span:
193+
if "input" in kwargs and isinstance(kwargs["input"], str):
194+
span.set_data("input", kwargs["input"])
195+
if "model" in kwargs:
196+
span.set_tag("model", kwargs["model"])
197+
if "dimensions" in kwargs:
198+
span.set_tag("dimensions", kwargs["dimensions"])
199+
response = f(*args, **kwargs)
200+
201+
prompt_tokens = 0
202+
total_tokens = 0
203+
if hasattr(response, "usage"):
204+
if hasattr(response.usage, "prompt_tokens") and isinstance(
205+
response.usage.prompt_tokens, int
206+
):
207+
prompt_tokens = response.usage.prompt_tokens
208+
if hasattr(response.usage, "total_tokens") and isinstance(
209+
response.usage.total_tokens, int
210+
):
211+
total_tokens = response.usage.total_tokens
212+
213+
if prompt_tokens == 0:
214+
prompt_tokens = count_tokens(kwargs["input"] or "")
215+
216+
if total_tokens == 0:
217+
total_tokens = prompt_tokens
218+
219+
span.set_data(PROMPT_TOKENS, prompt_tokens)
220+
span.set_data(TOTAL_TOKENS, total_tokens)
221+
222+
return response
223+
224+
return new_embeddings_create

‎tests/integrations/openai/test_openai.py

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
from openai import OpenAI, Stream
2+
from openai.types import CompletionUsage, CreateEmbeddingResponse, Embedding
3+
from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionChunk
4+
from openai.types.chat.chat_completion_chunk import ChoiceDelta, Choice
5+
from openai.types.create_embedding_response import Usage as EmbeddingTokenUsage
6+
7+
from sentry_sdk import start_transaction
8+
from sentry_sdk.integrations.openai import OpenAIIntegration
9+
10+
try:
11+
from unittest import mock # python 3.3 and above
12+
except ImportError:
13+
import mock # python < 3.3
14+
15+
COMPLETION_TOKENS = "completion_tоkens"
16+
PROMPT_TOKENS = "prompt_tоkens"
17+
TOTAL_TOKENS = "total_tоkens"
18+
19+
20+
def test_nonstreaming_chat_completion(sentry_init, capture_events):
21+
sentry_init(integrations=[OpenAIIntegration()], traces_sample_rate=1.0)
22+
events = capture_events()
23+
24+
client = OpenAI(api_key="z")
25+
returned_chat = ChatCompletion(
26+
id="chat-id",
27+
choices=[
28+
Choice(
29+
index=0,
30+
finish_reason="stop",
31+
message=ChatCompletionMessage(role="assistant", content="response"),
32+
)
33+
],
34+
created=10000000,
35+
model="model-id",
36+
object="chat.completion",
37+
usage=CompletionUsage(
38+
completion_tokens=10,
39+
prompt_tokens=20,
40+
total_tokens=30,
41+
),
42+
)
43+
44+
client.chat.completions._post = mock.Mock(return_value=returned_chat)
45+
with start_transaction(name="openai tx"):
46+
response = (
47+
client.chat.completions.create(
48+
model="some-model", messages=[{"role": "system", "content": "hello"}]
49+
)
50+
.choices[0]
51+
.message.content
52+
)
53+
54+
assert response == "response"
55+
tx = events[0]
56+
assert tx["type"] == "transaction"
57+
span = tx["spans"][0]
58+
assert span["op"] == "openai"
59+
60+
assert span["data"][COMPLETION_TOKENS] == 10
61+
assert span["data"][PROMPT_TOKENS] == 20
62+
assert span["data"][TOTAL_TOKENS] == 30
63+
64+
65+
# noinspection PyTypeChecker
66+
def test_streaming_chat_completion(sentry_init, capture_events):
67+
sentry_init(integrations=[OpenAIIntegration()], traces_sample_rate=1.0)
68+
events = capture_events()
69+
70+
client = OpenAI(api_key="z")
71+
returned_stream = Stream(cast_to=None, response=None, client=None)
72+
returned_stream._iterator = [
73+
ChatCompletionChunk(
74+
id="1",
75+
choices=[Choice(index=0, delta=ChoiceDelta(content="hel"))],
76+
created=100000,
77+
model="model-id",
78+
object="chat.completion.chunk",
79+
),
80+
ChatCompletionChunk(
81+
id="1",
82+
choices=[Choice(index=1, delta=ChoiceDelta(content="lo "))],
83+
created=100000,
84+
model="model-id",
85+
object="chat.completion.chunk",
86+
),
87+
ChatCompletionChunk(
88+
id="1",
89+
choices=[
90+
Choice(
91+
index=2, delta=ChoiceDelta(content="world"), finish_reason="stop"
92+
)
93+
],
94+
created=100000,
95+
model="model-id",
96+
object="chat.completion.chunk",
97+
),
98+
]
99+
100+
client.chat.completions._post = mock.Mock(return_value=returned_stream)
101+
with start_transaction(name="openai tx"):
102+
response_stream = client.chat.completions.create(
103+
model="some-model", messages=[{"role": "system", "content": "hello"}]
104+
)
105+
response_string = "".join(
106+
map(lambda x: x.choices[0].delta.content, response_stream)
107+
)
108+
assert response_string == "hello world"
109+
tx = events[0]
110+
assert tx["type"] == "transaction"
111+
span = tx["spans"][0]
112+
assert span["op"] == "openai"
113+
assert span["data"][COMPLETION_TOKENS] == 2
114+
assert span["data"][PROMPT_TOKENS] == 1
115+
assert span["data"][TOTAL_TOKENS] == 3
116+
117+
118+
def test_embeddings_create(sentry_init, capture_events):
119+
sentry_init(integrations=[OpenAIIntegration()], traces_sample_rate=1.0)
120+
events = capture_events()
121+
122+
client = OpenAI(api_key="z")
123+
124+
returned_embedding = CreateEmbeddingResponse(
125+
data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])],
126+
model="some-model",
127+
object="list",
128+
usage=EmbeddingTokenUsage(
129+
prompt_tokens=20,
130+
total_tokens=30,
131+
),
132+
)
133+
134+
client.embeddings._post = mock.Mock(return_value=returned_embedding)
135+
with start_transaction(name="openai tx"):
136+
response = client.embeddings.create(
137+
input="test", model="text-embedding-3-large"
138+
)
139+
140+
assert len(response.data[0].embedding) == 3
141+
142+
tx = events[0]
143+
assert tx["type"] == "transaction"
144+
span = tx["spans"][0]
145+
assert span["op"] == "openai"
146+
147+
assert span["data"][PROMPT_TOKENS] == 20
148+
assert span["data"][TOTAL_TOKENS] == 30

0 commit comments

Comments
 (0)