This commit is contained in:
Quinn Xu 2026-11-04 06:19:33 +00:00 • committed by GitHub
commit bd2b9f958f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 51 additions and 0 deletions

View file

@ -663,6 +663,15 @@ class DataDogLLMObsLogger(CustomBatchLogger):
redact_prompt_text=redact_payload,
)
# Datadog prices llm/embedding spans from top-level meta.model_name /
# model_provider. Those values already live in metadata; lift them so
# embedding spans are not flagged "Partial cost — Unsupported model provider".
# Narrow via isinstance so reportUnknownVariableType stays flat vs
# indexing a payload that basedpyright still treats as partially unknown.
_meta_model: Final[object] = payload_metadata.get("model_name", "unknown")
_meta_provider: Final[object] = payload_metadata.get("model_provider", "unknown")
model_name: Final[str] = _meta_model if isinstance(_meta_model, str) else "unknown"
model_provider: Final[str] = _meta_provider if isinstance(_meta_provider, str) else "unknown"
meta: Final[Meta] = {
"kind": span_kind,
"input": input_meta,
@ -670,6 +679,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
"metadata": payload_metadata,
"error": error_info,
**({"tool_definitions": tool_definitions} if tool_definitions else {}),
**(
{"model_name": model_name, "model_provider": model_provider}
if span_kind in ("llm", "embedding")
else {}
),
}
metrics: Final = self._assemble_metrics(standard_logging_payload, tool_output_tokens)

View file

@ -71,6 +71,10 @@ class DDLLMObsError(TypedDict, total=False):
class Meta(TypedDict, total=False):
# The span kind: "agent", "workflow", "llm", "tool", "task", "embedding", or "retrieval".
kind: Literal["llm", "tool", "task", "embedding", "retrieval"]
# Top-level model fields required for Datadog cost attribution on llm/embedding spans.
# Values are also kept under metadata for backwards compatibility.
model_name: ReadOnly[str]
model_provider: ReadOnly[str]
input: InputMeta # The span's input information.
output: OutputMeta # The span's output information.
metadata: ReadOnly[dict[str, object]]

View file

@ -1141,6 +1141,39 @@ def test_reasoning_content_survives_the_mapping(logger: DataDogLLMObsLogger) ->
assert payload["meta"]["output"]["messages"][0]["reasoning_content"] == "thinking"
def test_embedding_span_sets_top_level_model_name(logger: DataDogLLMObsLogger) -> None:
"""Datadog prices embedding spans from meta.model_name, not meta.metadata (issue #41601)."""
kwargs = build_payload(messages=None)
kwargs["standard_logging_object"]["call_type"] = "aembedding"
kwargs["standard_logging_object"]["model"] = "openai/text-embedding-3-small"
kwargs["standard_logging_object"]["custom_llm_provider"] = "openai"
kwargs["standard_logging_object"]["response"] = {"data": [{"embedding": [0.1, 0.2]}]}
kwargs["litellm_params"]["metadata"]["parent_id"] = "parent-span-1"
start = datetime(2026, 9, 1, 12, 0, 0)
span = json.loads(safe_dumps(logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=1))))
assert span["meta"]["kind"] == "embedding"
assert span["meta"]["model_name"] == "openai/text-embedding-3-small"
assert span["meta"]["model_provider"] == "openai"
# Still present under metadata for backwards compatibility.
assert span["meta"]["metadata"]["model_name"] == "openai/text-embedding-3-small"
assert span["meta"]["metadata"]["model_provider"] == "openai"
def test_llm_span_sets_top_level_model_name(logger: DataDogLLMObsLogger) -> None:
"""llm spans must keep top-level model_name/model_provider for Datadog cost attribution."""
kwargs = build_payload()
kwargs["standard_logging_object"]["model"] = "openai/gpt-4.1-mini"
kwargs["standard_logging_object"]["custom_llm_provider"] = "openai"
start = datetime(2026, 9, 1, 12, 0, 0)
span = json.loads(safe_dumps(logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=1))))
assert span["meta"]["kind"] == "llm"
assert span["meta"]["model_name"] == "openai/gpt-4.1-mini"
assert span["meta"]["model_provider"] == "openai"
@pytest.mark.parametrize(
("raw_arguments", "shipped"),
[