mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py
This commit is contained in:
Ishaan Jaffer
parent
4b7e8efe6c
commit
7d0aca21bc
3 changed files with 69 additions and 61 deletions
|
|
@ -29,7 +29,7 @@ from litellm.constants import (
|
|||
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
|
||||
STREAM_SSE_DATA_PREFIX,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, tracer
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
||||
get_response_headers,
|
||||
|
|
@ -41,6 +41,7 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
)
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -245,61 +246,6 @@ async def create_response(
|
|||
)
|
||||
|
||||
|
||||
class DDSpanTagger:
|
||||
"""Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata."""
|
||||
|
||||
@staticmethod
|
||||
def tag_call_id(litellm_call_id: Optional[str]) -> None:
|
||||
"""
|
||||
Attach LiteLLM call id to the active Datadog APM span.
|
||||
|
||||
This enables searching APM traces by LiteLLM call id returned in
|
||||
`x-litellm-call-id`.
|
||||
"""
|
||||
if not litellm_call_id:
|
||||
return
|
||||
try:
|
||||
set_active_span_tag("litellm.call_id", str(litellm_call_id))
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Failed to tag active ddtrace span with litellm.call_id",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def tag_request(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: Optional[str],
|
||||
) -> None:
|
||||
"""
|
||||
Attach key and model tags to the active Datadog APM span.
|
||||
|
||||
Tags set (all best-effort, skipped when value is absent):
|
||||
- ``litellm.key_alias`` — human-readable alias for the API key
|
||||
- ``litellm.key_hash`` — hashed API key (safe to log; never the raw secret)
|
||||
- ``litellm.requested_model``— model name as sent by the client
|
||||
|
||||
Use cases:
|
||||
- Trace all requests from a specific user/key: filter by ``litellm.key_alias`` or
|
||||
``litellm.key_hash``.
|
||||
- Trace all requests for a specific model: filter by ``litellm.requested_model``.
|
||||
|
||||
Note: key_alias / key_hash are not available for unauthenticated (e.g. 401) requests.
|
||||
"""
|
||||
try:
|
||||
if user_api_key_dict.key_alias:
|
||||
set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias))
|
||||
if user_api_key_dict.token:
|
||||
set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token))
|
||||
if requested_model:
|
||||
set_active_span_tag("litellm.requested_model", str(requested_model))
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Failed to tag active ddtrace span with key/model tags",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
def _override_openai_response_model(
|
||||
*,
|
||||
response_obj: Any,
|
||||
|
|
|
|||
60
litellm/proxy/dd_span_tagger.py
Normal file
60
litellm/proxy/dd_span_tagger.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
from typing import Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.dd_tracing import set_active_span_tag
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
class DDSpanTagger:
|
||||
"""Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata."""
|
||||
|
||||
@staticmethod
|
||||
def tag_call_id(litellm_call_id: Optional[str]) -> None:
|
||||
"""
|
||||
Attach LiteLLM call id to the active Datadog APM span.
|
||||
|
||||
This enables searching APM traces by LiteLLM call id returned in
|
||||
`x-litellm-call-id`.
|
||||
"""
|
||||
if not litellm_call_id:
|
||||
return
|
||||
try:
|
||||
set_active_span_tag("litellm.call_id", str(litellm_call_id))
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Failed to tag active ddtrace span with litellm.call_id",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def tag_request(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: Optional[str],
|
||||
) -> None:
|
||||
"""
|
||||
Attach key and model tags to the active Datadog APM span.
|
||||
|
||||
Tags set (all best-effort, skipped when value is absent):
|
||||
- ``litellm.key_alias`` — human-readable alias for the API key
|
||||
- ``litellm.key_hash`` — hashed API key (safe to log; never the raw secret)
|
||||
- ``litellm.requested_model``— model name as sent by the client
|
||||
|
||||
Use cases:
|
||||
- Trace all requests from a specific user/key: filter by ``litellm.key_alias`` or
|
||||
``litellm.key_hash``.
|
||||
- Trace all requests for a specific model: filter by ``litellm.requested_model``.
|
||||
|
||||
Note: key_alias / key_hash are not available for unauthenticated (e.g. 401) requests.
|
||||
"""
|
||||
try:
|
||||
if user_api_key_dict.key_alias:
|
||||
set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias))
|
||||
if user_api_key_dict.token:
|
||||
set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token))
|
||||
if requested_model:
|
||||
set_active_span_tag("litellm.requested_model", str(requested_model))
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Failed to tag active ddtrace span with key/model tags",
|
||||
exc_info=True,
|
||||
)
|
||||
|
|
@ -11,7 +11,6 @@ import litellm
|
|||
from litellm._uuid import uuid
|
||||
from litellm.integrations.opentelemetry import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import (
|
||||
DDSpanTagger,
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
ProxyConfig,
|
||||
_extract_error_from_sse_chunk,
|
||||
|
|
@ -20,6 +19,7 @@ from litellm.proxy.common_request_processing import (
|
|||
_parse_event_data_for_error,
|
||||
create_response,
|
||||
)
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
|
||||
|
|
@ -82,8 +82,10 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
|
||||
def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch):
|
||||
mock_set_active_span_tag = MagicMock(return_value=True)
|
||||
import litellm.proxy.dd_span_tagger
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing,
|
||||
litellm.proxy.dd_span_tagger,
|
||||
"set_active_span_tag",
|
||||
mock_set_active_span_tag,
|
||||
)
|
||||
|
|
@ -1582,7 +1584,7 @@ class TestDDSpanTaggerTagRequest:
|
|||
user_key = self._make_user_api_key_dict(key_alias="my-prod-key", token="hashed123")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_request_processing.set_active_span_tag"
|
||||
"litellm.proxy.dd_span_tagger.set_active_span_tag"
|
||||
) as mock_set_tag:
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_key,
|
||||
|
|
@ -1598,7 +1600,7 @@ class TestDDSpanTaggerTagRequest:
|
|||
user_key = self._make_user_api_key_dict(key_alias=None, token=None)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_request_processing.set_active_span_tag"
|
||||
"litellm.proxy.dd_span_tagger.set_active_span_tag"
|
||||
) as mock_set_tag:
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_key,
|
||||
|
|
@ -1612,7 +1614,7 @@ class TestDDSpanTaggerTagRequest:
|
|||
user_key = self._make_user_api_key_dict(key_alias=None, token=None)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_request_processing.set_active_span_tag"
|
||||
"litellm.proxy.dd_span_tagger.set_active_span_tag"
|
||||
) as mock_set_tag:
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_key,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue