Uh oh!
There was an error while loading. Please reload this page.
- Notifications
You must be signed in to change notification settings - Fork 657
OpenAI integration#2791
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Uh oh!
There was an error while loading. Please reload this page.
Merged
OpenAI integration #2791
Changes from all commits
Commits
Show all changes
22 commits
Select commit
Hold shift + click to select a range
6d79ebf
OpenAI integration
colin-sentry edf396c
Fix linting errors
colin-sentry 05ecefa
Fix CI
colin-sentry a3ccbcd
Fix lint
colin-sentry 6003902
Fix more CI issues
colin-sentry 09013f2
Run tests on version pinned OpenAI too
colin-sentry 2f9667a
Fix pydantic issue in test
colin-sentry ec84c04
Import type in TYPE_CHECKING gate
colin-sentry e45420e
PR feedback fixes
colin-sentry 0e319f1
Fix tiktoken test variant
colin-sentry 3acc9ab
PII gate the request and response
colin-sentry 65c311d
Rename set_data tags
colin-sentry 4b33a4e
Move doc location
colin-sentry 56d2679
Add "exclude prompts" flag as optional
colin-sentry 72d4b5a
Change prompts to be excluded by default
colin-sentry b206e4c
Set flag in tests
colin-sentry 4bc4310
Fix tiktoken tox.ini extra dash
colin-sentry c6f5cd2
Change strip PII semantics
colin-sentry d1eae09
More test coverage for PII
colin-sentry c96e0e5
notiktoken
colin-sentry 0278e2c
Merge branch 'master' into openai
antonpirker 452b03b
Merge branch 'master' into openai
colin-sentry File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Jump to
Jump to file
Failed to load files.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -70,6 +70,7 @@ | ||
| "beam", | ||
| "celery", | ||
| "huey", | ||
| "openai", | ||
| "rq", | ||
| ], | ||
| "Databases": [ | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,279 @@ | ||
| from sentry_sdk import consts | ||
| from sentry_sdk._types import TYPE_CHECKING | ||
| if TYPE_CHECKING: | ||
| from typing import Any, Iterable, List, Optional, Callable, Iterator | ||
| from sentry_sdk.tracing import Span | ||
| import sentry_sdk | ||
| from sentry_sdk._functools import wraps | ||
| from sentry_sdk.hub import Hub, _should_send_default_pii | ||
| from sentry_sdk.integrations import DidNotEnable, Integration | ||
| from sentry_sdk.utils import logger, capture_internal_exceptions, event_from_exception | ||
| try: | ||
| from openai.resources.chat.completions import Completions | ||
| from openai.resources import Embeddings | ||
| if TYPE_CHECKING: | ||
| from openai.types.chat import ChatCompletionMessageParam, ChatCompletionChunk | ||
| except ImportError: | ||
| raise DidNotEnable("OpenAI not installed") | ||
| try: | ||
| import tiktoken # type: ignore | ||
| enc = tiktoken.get_encoding("cl100k_base") | ||
| def count_tokens(s): | ||
| # type: (str) -> int | ||
| return len(enc.encode_ordinary(s)) | ||
| logger.debug("[OpenAI] using tiktoken to count tokens") | ||
| except ImportError: | ||
| logger.info( | ||
| "The Sentry Python SDK requires 'tiktoken' in order to measure token usage from some OpenAI APIs" | ||
| "Please install 'tiktoken' if you aren't receiving token usage in Sentry." | ||
| "See https://docs.sentry.io/platforms/python/integrations/openai/ for more information." | ||
| ) | ||
| def count_tokens(s): | ||
| # type: (str) -> int | ||
| return 0 | ||
| COMPLETION_TOKENS_USED = "ai.completion_tоkens.used" | ||
| PROMPT_TOKENS_USED = "ai.prompt_tоkens.used" | ||
| TOTAL_TOKENS_USED = "ai.total_tоkens.used" | ||
| class OpenAIIntegration(Integration): | ||
| identifier = "openai" | ||
| def __init__(self, include_prompts=True): | ||
| # type: (OpenAIIntegration, bool) -> None | ||
| self.include_prompts = include_prompts | ||
| @staticmethod | ||
| def setup_once(): | ||
| # type: () -> None | ||
| Completions.create = _wrap_chat_completion_create(Completions.create) | ||
| Embeddings.create = _wrap_embeddings_create(Embeddings.create) | ||
| def _capture_exception(hub, exc): | ||
| # type: (Hub, Any) -> None | ||
| if hub.client is not None: | ||
| event, hint = event_from_exception( | ||
| exc, | ||
| client_options=hub.client.options, | ||
| mechanism={"type": "openai", "handled": False}, | ||
| ) | ||
| hub.capture_event(event, hint=hint) | ||
| def _calculate_chat_completion_usage( | ||
| messages, response, span, streaming_message_responses=None | ||
| ): | ||
| # type: (Iterable[ChatCompletionMessageParam], Any, Span, Optional[List[str]]) -> None | ||
| completion_tokens = 0 | ||
| prompt_tokens = 0 | ||
| total_tokens = 0 | ||
| if hasattr(response, "usage"): | ||
| if hasattr(response.usage, "completion_tokens") and isinstance( | ||
| response.usage.completion_tokens, int | ||
| ): | ||
| completion_tokens = response.usage.completion_tokens | ||
| if hasattr(response.usage, "prompt_tokens") and isinstance( | ||
| response.usage.prompt_tokens, int | ||
| ): | ||
| prompt_tokens = response.usage.prompt_tokens | ||
| if hasattr(response.usage, "total_tokens") and isinstance( | ||
| response.usage.total_tokens, int | ||
| ): | ||
| total_tokens = response.usage.total_tokens | ||
| if prompt_tokens == 0: | ||
| for message in messages: | ||
| if "content" in message: | ||
| prompt_tokens += count_tokens(message["content"]) | ||
| if completion_tokens == 0: | ||
| if streaming_message_responses is not None: | ||
| for message in streaming_message_responses: | ||
| completion_tokens += count_tokens(message) | ||
| elif hasattr(response, "choices"): | ||
| for choice in response.choices: | ||
| if hasattr(choice, "message"): | ||
| completion_tokens += count_tokens(choice.message) | ||
| if total_tokens == 0: | ||
| total_tokens = prompt_tokens + completion_tokens | ||
| if completion_tokens != 0: | ||
| span.set_data(COMPLETION_TOKENS_USED, completion_tokens) | ||
| if prompt_tokens != 0: | ||
| span.set_data(PROMPT_TOKENS_USED, prompt_tokens) | ||
| if total_tokens != 0: | ||
| span.set_data(TOTAL_TOKENS_USED, total_tokens) | ||
| def _wrap_chat_completion_create(f): | ||
| # type: (Callable[..., Any]) -> Callable[..., Any] | ||
| @wraps(f) | ||
| def new_chat_completion(*args, **kwargs): | ||
| # type: (*Any, **Any) -> Any | ||
| hub = Hub.current | ||
| if not hub: | ||
| return f(*args, **kwargs) | ||
| integration = hub.get_integration(OpenAIIntegration) # type: OpenAIIntegration | ||
| if not integration: | ||
| 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) | ||
| kwargs["messages"] = list(kwargs["messages"]) | ||
| messages = kwargs["messages"] | ||
| model = kwargs.get("model") | ||
| streaming = kwargs.get("stream") | ||
| span = sentry_sdk.start_span( | ||
| op=consts.OP.OPENAI_CHAT_COMPLETIONS_CREATE, description="Chat Completion" | ||
| ) | ||
| span.__enter__() | ||
| try: | ||
| res = f(*args, **kwargs) | ||
| except Exception as e: | ||
| _capture_exception(Hub.current, e) | ||
| span.__exit__(None, None, None) | ||
| raise e from None | ||
| with capture_internal_exceptions(): | ||
| if _should_send_default_pii() and integration.include_prompts: | ||
| span.set_data("ai.input_messages", messages) | ||
| span.set_data("ai.model_id", model) | ||
| span.set_data("ai.streaming", streaming) | ||
| if hasattr(res, "choices"): | ||
| if _should_send_default_pii() and integration.include_prompts: | ||
| span.set_data( | ||
| "ai.responses", list(map(lambda x: x.message, res.choices)) | ||
| ) | ||
| _calculate_chat_completion_usage(messages, res, span) | ||
| span.__exit__(None, None, None) | ||
| elif hasattr(res, "_iterator"): | ||
| data_buf: list[list[str]] = [] # one for each choice | ||
| old_iterator = res._iterator # type: Iterator[ChatCompletionChunk] | ||
| def new_iterator(): | ||
| # type: () -> Iterator[ChatCompletionChunk] | ||
| with capture_internal_exceptions(): | ||
| for x in old_iterator: | ||
| if hasattr(x, "choices"): | ||
| choice_index = 0 | ||
| for choice in x.choices: | ||
| if hasattr(choice, "delta") and hasattr( | ||
| choice.delta, "content" | ||
| ): | ||
| content = choice.delta.content | ||
| if len(data_buf) <= choice_index: | ||
| data_buf.append([]) | ||
| data_buf[choice_index].append(content or "") | ||
| choice_index += 1 | ||
| yield x | ||
| if len(data_buf) > 0: | ||
| all_responses = list( | ||
| map(lambda chunk: "".join(chunk), data_buf) | ||
| ) | ||
| if ( | ||
| _should_send_default_pii() | ||
| and integration.include_prompts | ||
| ): | ||
| span.set_data("ai.responses", all_responses) | ||
| _calculate_chat_completion_usage( | ||
| messages, res, span, all_responses | ||
| ) | ||
| span.__exit__(None, None, None) | ||
| res._iterator = new_iterator() | ||
| else: | ||
| span.set_data("unknown_response", True) | ||
| span.__exit__(None, None, None) | ||
| return res | ||
| return new_chat_completion | ||
| def _wrap_embeddings_create(f): | ||
| # type: (Callable[..., Any]) -> Callable[..., Any] | ||
| @wraps(f) | ||
| def new_embeddings_create(*args, **kwargs): | ||
| # type: (*Any, **Any) -> Any | ||
| hub = Hub.current | ||
| if not hub: | ||
| return f(*args, **kwargs) | ||
| integration = hub.get_integration(OpenAIIntegration) # type: OpenAIIntegration | ||
| if not integration: | ||
| return f(*args, **kwargs) | ||
| with sentry_sdk.start_span( | ||
| op=consts.OP.OPENAI_EMBEDDINGS_CREATE, | ||
| description="OpenAI Embedding Creation", | ||
| ) as span: | ||
| if "input" in kwargs and ( | ||
| _should_send_default_pii() and integration.include_prompts | ||
| ): | ||
| if isinstance(kwargs["input"], str): | ||
| span.set_data("ai.input_messages", [kwargs["input"]]) | ||
| elif ( | ||
| isinstance(kwargs["input"], list) | ||
| and len(kwargs["input"]) > 0 | ||
| and isinstance(kwargs["input"][0], str) | ||
| ): | ||
| span.set_data("ai.input_messages", kwargs["input"]) | ||
| if "model" in kwargs: | ||
| span.set_data("ai.model_id", kwargs["model"]) | ||
| try: | ||
| response = f(*args, **kwargs) | ||
| except Exception as e: | ||
| _capture_exception(Hub.current, e) | ||
| raise e from None | ||
| prompt_tokens = 0 | ||
| total_tokens = 0 | ||
| if hasattr(response, "usage"): | ||
| if hasattr(response.usage, "prompt_tokens") and isinstance( | ||
| response.usage.prompt_tokens, int | ||
| ): | ||
| prompt_tokens = response.usage.prompt_tokens | ||
| if hasattr(response.usage, "total_tokens") and isinstance( | ||
| response.usage.total_tokens, int | ||
| ): | ||
| total_tokens = response.usage.total_tokens | ||
| if prompt_tokens == 0: | ||
| prompt_tokens = count_tokens(kwargs["input"] or "") | ||
| if total_tokens == 0: | ||
| total_tokens = prompt_tokens | ||
| span.set_data(PROMPT_TOKENS_USED, prompt_tokens) | ||
| span.set_data(TOTAL_TOKENS_USED, total_tokens) | ||
| return response | ||
| return new_embeddings_create | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,3 @@ | ||
| import pytest | ||
| pytest.importorskip("openai") |
Oops, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.