scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix: Datadog span kind fallback when parent_id missing

Browse source
This commit is contained in:
Yuta Saito 2025-12-24 14:12:17 +09:00
parent 97230e5179
commit 61e1e42e97
2 changed files with 48 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -217,8 +217,14 @@ class DataDogLLMObsLogger(CustomBatchLogger):
error_info = self._assemble_error_info(standard_logging_payload)
metadata_parent_id: Optional[str] = None
if isinstance(metadata, dict):
metadata_parent_id = metadata.get("parent_id")
meta = Meta(
kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type")),
kind=self._get_datadog_span_kind(
standard_logging_payload.get("call_type"), metadata_parent_id
),
input=input_meta,
output=output_meta,
metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
@ -237,7 +243,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
)
payload: LLMObsPayload = LLMObsPayload(
parent_id=metadata.get("parent_id", "undefined"),
parent_id=metadata_parent_id if metadata_parent_id else "undefined",
trace_id=standard_logging_payload.get("trace_id", str(uuid.uuid4())),
span_id=metadata.get("span_id", str(uuid.uuid4())),
name=metadata.get("name", "litellm_llm_call"),
@ -367,14 +373,16 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return []
def _get_datadog_span_kind(
self, call_type: Optional[str]
self, call_type: Optional[str], parent_id: Optional[str] = None
) -> Literal["llm", "tool", "task", "embedding", "retrieval"]:
"""
Map liteLLM call_type to appropriate DataDog LLM Observability span kind.
Available DataDog span kinds: "llm", "tool", "task", "embedding", "retrieval"
see: https://docs.datadoghq.com/ja/llm_observability/terms/
"""
if call_type is None:
# Non llm/workflow/agent kinds cannot be root spans, so fallback to "llm" when parent metadata is missing
if call_type is None or parent_id is None:
return "llm"
# Embedding operations
@ -392,6 +400,8 @@ class DataDogLLMObsLogger(CustomBatchLogger):
CallTypes.generate_content_stream.value,
CallTypes.agenerate_content_stream.value,
CallTypes.anthropic_messages.value,
CallTypes.responses.value,
CallTypes.aresponses.value,
]:
return "llm"
@ -417,8 +427,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
CallTypes.aretrieve_batch.value,
CallTypes.retrieve_fine_tuning_job.value,
CallTypes.aretrieve_fine_tuning_job.value,
CallTypes.responses.value,
CallTypes.aresponses.value,
CallTypes.alist_input_items.value,
]:
return "retrieval"

52
tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py
View file

@ -257,41 +257,57 @@ class TestDataDogLLMObsLogger:
logger = DataDogLLMObsLogger()
# Test embedding operations
assert logger._get_datadog_span_kind(CallTypes.embedding.value) == "embedding"
assert logger._get_datadog_span_kind(CallTypes.aembedding.value) == "embedding"
assert logger._get_datadog_span_kind(CallTypes.embedding.value, "123") == "embedding"
assert logger._get_datadog_span_kind(CallTypes.aembedding.value, "123") == "embedding"
# Test LLM completion operations
assert logger._get_datadog_span_kind(CallTypes.completion.value) == "llm"
assert logger._get_datadog_span_kind(CallTypes.acompletion.value) == "llm"
assert logger._get_datadog_span_kind(CallTypes.text_completion.value) == "llm"
assert logger._get_datadog_span_kind(CallTypes.generate_content.value) == "llm"
assert logger._get_datadog_span_kind(CallTypes.completion.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.acompletion.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.text_completion.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.generate_content.value, None) == "llm"
assert (
logger._get_datadog_span_kind(CallTypes.anthropic_messages.value) == "llm"
logger._get_datadog_span_kind(CallTypes.anthropic_messages.value, None) == "llm"
)
assert logger._get_datadog_span_kind(CallTypes.responses.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.aresponses.value, None) == "llm"
# Test tool operations
assert logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value) == "tool"
assert logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") == "tool"
# Test retrieval operations
assert (
logger._get_datadog_span_kind(CallTypes.get_assistants.value) == "retrieval"
logger._get_datadog_span_kind(CallTypes.get_assistants.value, "123") == "retrieval"
)
assert (
logger._get_datadog_span_kind(CallTypes.file_retrieve.value) == "retrieval"
logger._get_datadog_span_kind(CallTypes.file_retrieve.value, "123") == "retrieval"
)
assert (
logger._get_datadog_span_kind(CallTypes.retrieve_batch.value) == "retrieval"
logger._get_datadog_span_kind(CallTypes.retrieve_batch.value, "123") == "retrieval"
)
# Test task operations
assert logger._get_datadog_span_kind(CallTypes.create_batch.value) == "task"
assert logger._get_datadog_span_kind(CallTypes.image_generation.value) == "task"
assert logger._get_datadog_span_kind(CallTypes.moderation.value) == "task"
assert logger._get_datadog_span_kind(CallTypes.transcription.value) == "task"
assert logger._get_datadog_span_kind(CallTypes.create_batch.value, "123") == "task"
assert logger._get_datadog_span_kind(CallTypes.image_generation.value, "123") == "task"
assert logger._get_datadog_span_kind(CallTypes.moderation.value, "123") == "task"
assert logger._get_datadog_span_kind(CallTypes.transcription.value, "123") == "task"
# Test default fallback
assert logger._get_datadog_span_kind("unknown_call_type") == "llm"
assert logger._get_datadog_span_kind(None) == "llm"
assert logger._get_datadog_span_kind("unknown_call_type", None) == "llm"
assert logger._get_datadog_span_kind(None, None) == "llm"
def test_datadog_span_kind_defaults_without_parent(self, mock_env_vars):
"""Test that non-llm kinds fallback to llm when no parent span is provided"""
from litellm.types.utils import CallTypes
with patch(
"litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"
), patch("asyncio.create_task"):
logger = DataDogLLMObsLogger()
# Tool/task/retrieval span kinds should fallback to llm when parent_id missing
assert logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.create_batch.value, None) == "llm"
assert logger._get_datadog_span_kind(CallTypes.get_assistants.value, None) == "llm"
@pytest.mark.asyncio
async def test_async_log_failure_event(self, mock_env_vars):
@ -796,7 +812,7 @@ class TestDataDogLLMObsLoggerToolCalls:
from litellm.types.utils import CallTypes
assert (
logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value) == "tool"
logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") == "tool"
)
def test_tool_call_payload_creation(self, mock_env_vars):