diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index 835acc7279..5db9870210 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -1,3 +1,4 @@ +import contextvars import inspect import sys from functools import wraps @@ -32,6 +33,11 @@ raise DidNotEnable("Huggingface not installed") +_active_huggingface_task = contextvars.ContextVar( + "active_huggingface_task", default=False +) + + class HuggingfaceHubIntegration(Integration): identifier = "huggingface_hub" origin = f"auto.ai.{identifier}" @@ -56,6 +62,29 @@ def setup_once() -> None: OP.GEN_AI_CHAT, ) ) + _patch_huggingface_chat_completion_alias() + + +def _patch_huggingface_chat_completion_alias() -> None: + proxy_class = getattr( + huggingface_hub.inference._client, "ProxyClientChatCompletions", None + ) + if proxy_class is None: + inference_client = huggingface_hub.inference._client.InferenceClient + try: + chat = getattr(inference_client(), "chat", None) + completions = getattr(chat, "completions", None) + except Exception: + return + if completions is None: + return + proxy_class = completions.__class__ + + create = getattr(proxy_class, "create", None) + if create is None or isinstance(create, property): + return + + proxy_class.create = _wrap_huggingface_task(create, OP.GEN_AI_CHAT) def _capture_exception(exc: "Any") -> None: @@ -87,7 +116,10 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": # invalid call, dont instrument, let it return error return f(*args, **kwargs) - client = args[0] + client = getattr(args[0], "_client", args[0]) + if _active_huggingface_task.get(): + return f(*args, **kwargs) + model = client.model or kwargs.get("model") or "" operation_name = op.split(".")[-1] @@ -139,6 +171,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": set_data_normalized(span, span_attribute, value, unpack=False) # LLM Execution + task_token = _active_huggingface_task.set(True) try: res = f(*args, **kwargs) except Exception as e: @@ -147,6 +180,8 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": _capture_exception(e) span.__exit__(*exc_info) reraise(*exc_info) + finally: + _active_huggingface_task.reset(task_token) # Output attributes finish_reason = None diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index 91e0909731..75deca0824 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -892,6 +892,114 @@ def test_chat_completion( assert span["data"] == expected_data +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_chat_completion_openai_compatible_alias( + sentry_init: "Any", + capture_events: "Any", + capture_items: "Any", + mock_hf_chat_completion_api: "Any", + stream_gen_ai_spans: "Any", +) -> None: + client = get_hf_provider_inference_client() + chat = getattr(client, "chat", None) + completions = getattr(chat, "completions", None) if chat is not None else None + if completions is None or not hasattr(completions, "create"): + pytest.skip("OpenAI-compatible chat completion alias is unavailable") + + sentry_init( + traces_sample_rate=1.0, + send_default_pii=True, + integrations=[HuggingfaceHubIntegration()], + stream_gen_ai_spans=stream_gen_ai_spans, + ) + + messages = [{"role": "user", "content": "Hello!"}] + + if stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with sentry_sdk.start_transaction(name="test"): + client.chat.completions.create(messages=messages, stream=False) + + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if sp["attributes"]["sentry.op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.messages": safe_serialize(messages), + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "stop", + "gen_ai.response.model": "test-model-123", + "gen_ai.response.streaming": False, + "gen_ai.response.text": "[mocked] Hello! How can I help you today?", + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 18, + "process.runtime.name": mock.ANY, + "process.runtime.version": mock.ANY, + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + "sentry.release": mock.ANY, + "sentry.sdk.name": "sentry.python", + "sentry.sdk.version": mock.ANY, + "sentry.segment.id": mock.ANY, + "sentry.segment.name": "test", + "server.address": mock.ANY, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } + assert span["attributes"] == expected_data + else: + events = capture_events() + + with sentry_sdk.start_transaction(name="test"): + client.chat.completions.create(messages=messages, stream=False) + + (transaction,) = events + + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + assert sp["op"] == "http.client" + + assert span is not None + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.messages": safe_serialize(messages), + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "stop", + "gen_ai.response.model": "test-model-123", + "gen_ai.response.streaming": False, + "gen_ai.response.text": "[mocked] Hello! How can I help you today?", + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 18, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } + assert span["data"] == expected_data + + @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False])