Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 18 additions & 9 deletions sentry_sdk/integrations/cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,17 +253,26 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
old_iterator = res

def new_iterator() -> "Iterator[StreamedChatResponse]":
with capture_internal_exceptions():
try:
for x in old_iterator:
if isinstance(x, ChatStreamEndEvent) or isinstance(
x, StreamEndStreamedChatResponse
):
collect_chat_response_fields(
span,
x.response,
include_pii=_should_record(integration, "outputs"),
)
with capture_internal_exceptions():
if isinstance(x, ChatStreamEndEvent) or isinstance(
x, StreamEndStreamedChatResponse
):
collect_chat_response_fields(
span,
x.response,
include_pii=_should_record(
integration, "outputs"
),
)
yield x
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
span.__exit__(*exc_info)
reraise(*exc_info)

_end_span(span)

return new_iterator()
Expand Down
115 changes: 66 additions & 49 deletions sentry_sdk/integrations/huggingface_hub.py
Original file line number Diff line number Diff line change
Expand Up @@ -297,29 +297,36 @@ def new_details_iterator() -> "Iterable[Any]":
response_text_buffer: "list[str]" = []
tokens_used = 0

with capture_internal_exceptions():
try:
for chunk in res:
if (
hasattr(chunk, "token")
and hasattr(chunk.token, "text")
and chunk.token.text is not None
):
response_text_buffer.append(chunk.token.text)
with capture_internal_exceptions():
if (
hasattr(chunk, "token")
and hasattr(chunk.token, "text")
and chunk.token.text is not None
):
response_text_buffer.append(chunk.token.text)

if hasattr(chunk, "details") and hasattr(
chunk.details, "finish_reason"
):
finish_reason = chunk.details.finish_reason
if hasattr(chunk, "details") and hasattr(
chunk.details, "finish_reason"
):
finish_reason = chunk.details.finish_reason

if (
hasattr(chunk, "details")
and hasattr(chunk.details, "generated_tokens")
and chunk.details.generated_tokens is not None
):
tokens_used = chunk.details.generated_tokens
if (
hasattr(chunk, "details")
and hasattr(chunk.details, "generated_tokens")
and chunk.details.generated_tokens is not None
):
tokens_used = chunk.details.generated_tokens

yield chunk
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
span.__exit__(*exc_info)
reraise(*exc_info)

with capture_internal_exceptions():
if finish_reason is not None:
set_data_normalized(
span,
Expand Down Expand Up @@ -360,44 +367,54 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]":
tool_calls = None
usage = None

with capture_internal_exceptions():
try:
for chunk in cast("Iterable[ChatCompletionStreamOutput]", res):
if hasattr(chunk, "model") and chunk.model is not None:
response_model = chunk.model

if hasattr(chunk, "usage") and chunk.usage is not None:
usage = chunk.usage

if isinstance(chunk, str):
if chunk is not None:
response_text_buffer.append(chunk)

if hasattr(chunk, "choices") and chunk.choices is not None:
for choice in chunk.choices:
if (
hasattr(choice, "delta")
and hasattr(choice.delta, "content")
and choice.delta.content is not None
):
response_text_buffer.append(
choice.delta.content
)
with capture_internal_exceptions():
if hasattr(chunk, "model") and chunk.model is not None:
response_model = chunk.model

if (
hasattr(choice, "finish_reason")
and choice.finish_reason is not None
):
finish_reason = choice.finish_reason
if hasattr(chunk, "usage") and chunk.usage is not None:
usage = chunk.usage

if (
hasattr(choice, "delta")
and hasattr(choice.delta, "tool_calls")
and choice.delta.tool_calls is not None
):
tool_calls = choice.delta.tool_calls
if isinstance(chunk, str):
if chunk is not None:
response_text_buffer.append(chunk)

if (
hasattr(chunk, "choices")
and chunk.choices is not None
):
for choice in chunk.choices:
if (
hasattr(choice, "delta")
and hasattr(choice.delta, "content")
and choice.delta.content is not None
):
response_text_buffer.append(
choice.delta.content
)

if (
hasattr(choice, "finish_reason")
and choice.finish_reason is not None
):
finish_reason = choice.finish_reason

if (
hasattr(choice, "delta")
and hasattr(choice.delta, "tool_calls")
and choice.delta.tool_calls is not None
):
tool_calls = choice.delta.tool_calls

yield chunk
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
span.__exit__(*exc_info)
reraise(*exc_info)

with capture_internal_exceptions():
if response_model is not None:
_set_span_data_attribute(
span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
Expand Down
39 changes: 39 additions & 0 deletions tests/integrations/cohere/test_cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -866,3 +866,42 @@ def test_embed_data_collection(
assert "hello" in str(attributes[SPANDATA.AI_INPUT_MESSAGES])
else:
assert SPANDATA.AI_INPUT_MESSAGES not in attributes


@pytest.mark.parametrize("span_streaming", [True, False])
def test_streaming_chat_iteration_error_propagates(
sentry_init, capture_events, capture_items, span_streaming
):
"""Mid-stream provider errors must reach the caller, not be swallowed.

Regression test for https://github.com/getsentry/sentry-python/issues/7889
"""
from sentry_sdk.integrations.cohere import _wrap_chat

sentry_init(
integrations=[CohereIntegration()],
traces_sample_rate=1.0,
trace_lifecycle="stream" if span_streaming else "static",
)
events = capture_events()
items = capture_items("span") if span_streaming else None

def failing_stream(*args, **kwargs):
yield object()
raise httpx.RemoteProtocolError("connection dropped", request=None)

wrapped = _wrap_chat(failing_stream, streaming=True)

with start_transaction(name="test"):
with pytest.raises(httpx.RemoteProtocolError):
list(wrapped(message="hello"))

sentry_sdk.flush()

if span_streaming:
assert len(items) == 1
assert items[0].payload["status"] == "error"
else:
transactions = [e for e in events if e.get("type") == "transaction"]
assert len(transactions) == 1
assert transactions[0]["spans"][0]["status"] == "internal_error"
69 changes: 69 additions & 0 deletions tests/integrations/huggingface_hub/test_huggingface_hub.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import re
from types import SimpleNamespace
from typing import TYPE_CHECKING

import pytest
Expand Down Expand Up @@ -2712,3 +2713,71 @@ def test_chat_completion_streaming_data_collection_tools(
assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 183
assert span_data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 14
assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 197


def test_text_generation_streaming_iteration_error_propagates(
sentry_init: "Any", capture_events: "Any"
) -> None:
"""Mid-stream provider errors must reach the caller, not be swallowed.

Regression test for https://github.com/getsentry/sentry-python/issues/7889
"""
from sentry_sdk.integrations.huggingface_hub import _wrap_huggingface_task

sentry_init(
integrations=[HuggingfaceHubIntegration()],
traces_sample_rate=1.0,
)

def failing_stream(*args: "Any", **kwargs: "Any") -> "Any":
yield SimpleNamespace(
token=SimpleNamespace(text="the mocked "),
details=SimpleNamespace(finish_reason=None, generated_tokens=None),
)
raise RuntimeError("mid-stream failure")

wrapped = _wrap_huggingface_task(failing_stream, "gen_ai.text_completion")

with sentry_sdk.start_transaction(name="test"):
with pytest.raises(RuntimeError, match="mid-stream failure"):
list(
wrapped(
SimpleNamespace(model="test-model"),
prompt="Hello",
details=True,
)
)

sentry_sdk.flush()


def test_chat_completion_streaming_iteration_error_propagates(
sentry_init: "Any", capture_events: "Any"
) -> None:
"""Mid-stream provider errors must reach the caller, not be swallowed.

Regression test for https://github.com/getsentry/sentry-python/issues/7889
"""
from sentry_sdk.integrations.huggingface_hub import _wrap_huggingface_task

sentry_init(
integrations=[HuggingfaceHubIntegration()],
traces_sample_rate=1.0,
)

def failing_stream(*args: "Any", **kwargs: "Any") -> "Any":
yield SimpleNamespace(model="test-model", usage=None, choices=None)
raise RuntimeError("mid-stream failure")

wrapped = _wrap_huggingface_task(failing_stream, "gen_ai.chat_completions")

with sentry_sdk.start_transaction(name="test"):
with pytest.raises(RuntimeError, match="mid-stream failure"):
list(
wrapped(
SimpleNamespace(model="test-model"),
messages=[{"role": "user", "content": "Hello"}],
)
)

sentry_sdk.flush()