wehub-resource-sync
1111 lines
40 KiB
Python
1111 lines
40 KiB
Python
import asyncio
|
|
import json
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities import SpanType
|
|
from mlflow.gateway.constants import MLFLOW_GATEWAY_CALLER_HEADER
|
|
from mlflow.gateway.schemas.chat import StreamResponsePayload
|
|
from mlflow.gateway.tracing_utils import (
|
|
_extract_caller,
|
|
_get_model_span_info,
|
|
aggregate_anthropic_messages_stream_chunks,
|
|
aggregate_chat_stream_chunks,
|
|
aggregate_gemini_stream_generate_content_chunks,
|
|
aggregate_openai_responses_stream_chunks,
|
|
maybe_traced_gateway_call,
|
|
)
|
|
from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig
|
|
from mlflow.tracing.client import TracingClient
|
|
from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey
|
|
from mlflow.tracing.distributed import get_tracing_context_headers_for_http_request
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
from mlflow.types.chat import ChatChoiceDelta, ChatChunkChoice, ChatUsage, Function, ToolCallDelta
|
|
|
|
|
|
@pytest.fixture
|
|
def gateway_experiment_id():
|
|
experiment_name = "gateway-test-endpoint"
|
|
experiment = mlflow.get_experiment_by_name(experiment_name)
|
|
if experiment is not None:
|
|
return experiment.experiment_id
|
|
return mlflow.create_experiment(experiment_name)
|
|
|
|
|
|
def get_traces():
|
|
return TracingClient().search_traces(locations=[_get_experiment_id()])
|
|
|
|
|
|
@pytest.fixture
|
|
def endpoint_config():
|
|
return GatewayEndpointConfig(
|
|
endpoint_id="test-endpoint-id",
|
|
endpoint_name="test-endpoint",
|
|
experiment_id=_get_experiment_id(),
|
|
usage_tracking=True,
|
|
models=[],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def endpoint_config_no_experiment():
|
|
return GatewayEndpointConfig(
|
|
endpoint_id="test-endpoint-id",
|
|
endpoint_name="test-endpoint",
|
|
experiment_id=None,
|
|
models=[],
|
|
)
|
|
|
|
|
|
async def mock_async_func(payload):
|
|
return {"result": "success", "payload": payload}
|
|
|
|
|
|
def _make_chunk(
|
|
content=None,
|
|
finish_reason=None,
|
|
id="chunk-1",
|
|
model="test-model",
|
|
created=1700000000,
|
|
usage=None,
|
|
tool_calls=None,
|
|
role="assistant",
|
|
choice_index=0,
|
|
):
|
|
delta = ChatChoiceDelta(role=role, content=content, tool_calls=tool_calls)
|
|
choice = ChatChunkChoice(index=choice_index, finish_reason=finish_reason, delta=delta)
|
|
return StreamResponsePayload(
|
|
id=id,
|
|
created=created,
|
|
model=model,
|
|
choices=[choice],
|
|
usage=usage,
|
|
)
|
|
|
|
|
|
def test_aggregate_chat_stream_chunks_aggregates_content():
|
|
chunks = [
|
|
_make_chunk(content="Hello"),
|
|
_make_chunk(content=" "),
|
|
_make_chunk(content="world"),
|
|
_make_chunk(content=None, finish_reason="stop"),
|
|
]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert result["object"] == "chat.completion"
|
|
assert result["model"] == "test-model"
|
|
assert result["choices"][0]["message"]["role"] == "assistant"
|
|
assert result["choices"][0]["message"]["content"] == "Hello world"
|
|
assert result["choices"][0]["finish_reason"] == "stop"
|
|
|
|
|
|
def test_aggregate_chat_stream_chunks_with_usage():
|
|
usage = ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
|
|
chunks = [
|
|
_make_chunk(content="Hi"),
|
|
_make_chunk(content=None, finish_reason="stop", usage=usage),
|
|
]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert result["choices"][0]["message"]["content"] == "Hi"
|
|
assert result["usage"] == {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 15,
|
|
}
|
|
|
|
|
|
def test_aggregate_chat_stream_chunks_empty():
|
|
assert aggregate_chat_stream_chunks([]) is None
|
|
|
|
|
|
def test_aggregate_chat_stream_chunks_defaults_finish_reason():
|
|
chunks = [_make_chunk(content="Hi")]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert result["choices"][0]["finish_reason"] == "stop"
|
|
|
|
|
|
def test_reduce_chat_stream_chunks_aggregates_tool_calls():
|
|
chunks = [
|
|
# First chunk: tool call id, type, and function name
|
|
_make_chunk(
|
|
tool_calls=[
|
|
ToolCallDelta(
|
|
index=0,
|
|
id="call_abc",
|
|
type="function",
|
|
function=Function(name="get_weather", arguments=""),
|
|
),
|
|
],
|
|
),
|
|
# Subsequent chunks: argument fragments
|
|
_make_chunk(
|
|
tool_calls=[
|
|
ToolCallDelta(index=0, function=Function(arguments='{"loc')),
|
|
],
|
|
),
|
|
_make_chunk(
|
|
tool_calls=[
|
|
ToolCallDelta(index=0, function=Function(arguments='ation": "SF"}')),
|
|
],
|
|
),
|
|
_make_chunk(finish_reason="tool_calls"),
|
|
]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert result["choices"][0]["message"]["content"] is None
|
|
assert result["choices"][0]["finish_reason"] == "tool_calls"
|
|
|
|
tool_calls = result["choices"][0]["message"]["tool_calls"]
|
|
assert len(tool_calls) == 1
|
|
assert tool_calls[0]["id"] == "call_abc"
|
|
assert tool_calls[0]["type"] == "function"
|
|
assert tool_calls[0]["function"]["name"] == "get_weather"
|
|
assert tool_calls[0]["function"]["arguments"] == '{"location": "SF"}'
|
|
|
|
|
|
def test_reduce_chat_stream_chunks_derives_role_from_delta():
|
|
chunks = [
|
|
_make_chunk(role="developer", content="Hello"),
|
|
_make_chunk(role=None, content=" world"),
|
|
_make_chunk(role=None, finish_reason="stop"),
|
|
]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert result["choices"][0]["message"]["role"] == "developer"
|
|
|
|
|
|
def test_reduce_chat_stream_chunks_multiple_choice_indices():
|
|
chunks = [
|
|
_make_chunk(content="Hi", choice_index=0),
|
|
_make_chunk(content="Hey", choice_index=1),
|
|
_make_chunk(content=" there", choice_index=0),
|
|
_make_chunk(content=" you", choice_index=1),
|
|
_make_chunk(finish_reason="stop", choice_index=0),
|
|
_make_chunk(finish_reason="stop", choice_index=1),
|
|
]
|
|
result = aggregate_chat_stream_chunks(chunks)
|
|
|
|
assert len(result["choices"]) == 2
|
|
assert result["choices"][0]["index"] == 0
|
|
assert result["choices"][0]["message"]["content"] == "Hi there"
|
|
assert result["choices"][1]["index"] == 1
|
|
assert result["choices"][1]["message"]["content"] == "Hey you"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_basic(endpoint_config):
|
|
traced_func = maybe_traced_gateway_call(mock_async_func, endpoint_config)
|
|
result = await traced_func({"input": "test"})
|
|
|
|
assert result == {"result": "success", "payload": {"input": "test"}}
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
# Find the gateway span
|
|
span_name_to_span = {span.name: span for span in trace.data.spans}
|
|
assert f"gateway/{endpoint_config.endpoint_name}" in span_name_to_span
|
|
|
|
gateway_span = span_name_to_span[f"gateway/{endpoint_config.endpoint_name}"]
|
|
assert gateway_span.attributes.get("endpoint_id") == "test-endpoint-id"
|
|
assert gateway_span.attributes.get("endpoint_name") == "test-endpoint"
|
|
# Input should be unwrapped (not nested under "payload" key)
|
|
assert gateway_span.inputs == {"input": "test"}
|
|
# No user metadata should be present in trace
|
|
assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USERNAME) is None
|
|
assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USER_ID) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_user_metadata(endpoint_config):
|
|
traced_func = maybe_traced_gateway_call(
|
|
mock_async_func,
|
|
endpoint_config,
|
|
metadata={
|
|
TraceMetadataKey.AUTH_USERNAME: "alice",
|
|
TraceMetadataKey.AUTH_USER_ID: "123",
|
|
},
|
|
)
|
|
result = await traced_func({"input": "test"})
|
|
|
|
assert result == {"result": "success", "payload": {"input": "test"}}
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
span_name_to_span = {span.name: span for span in trace.data.spans}
|
|
gateway_span = span_name_to_span[f"gateway/{endpoint_config.endpoint_name}"]
|
|
|
|
assert gateway_span.attributes.get("endpoint_id") == "test-endpoint-id"
|
|
assert gateway_span.attributes.get("endpoint_name") == "test-endpoint"
|
|
# Input should be unwrapped (not nested under "payload" key)
|
|
assert gateway_span.inputs == {"input": "test"}
|
|
# User metadata should be in trace info, not span attributes
|
|
assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USERNAME) == "alice"
|
|
assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USER_ID) == "123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_without_usage_tracking(endpoint_config_no_experiment):
|
|
traced_func = maybe_traced_gateway_call(
|
|
mock_async_func,
|
|
endpoint_config_no_experiment,
|
|
metadata={
|
|
TraceMetadataKey.AUTH_USERNAME: "alice",
|
|
TraceMetadataKey.AUTH_USER_ID: "123",
|
|
},
|
|
)
|
|
|
|
# When usage_tracking is False, maybe_traced_gateway_call returns the original function
|
|
assert traced_func is mock_async_func
|
|
|
|
result = await traced_func({"input": "test"})
|
|
assert result == {"result": "success", "payload": {"input": "test"}}
|
|
|
|
# No traces should be created
|
|
traces = get_traces()
|
|
assert len(traces) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_output_reducer(endpoint_config):
|
|
async def mock_async_stream(payload):
|
|
yield _make_chunk(content="Hello")
|
|
yield _make_chunk(content=" world")
|
|
yield _make_chunk(
|
|
content=None,
|
|
finish_reason="stop",
|
|
usage=ChatUsage(prompt_tokens=5, completion_tokens=2, total_tokens=7),
|
|
)
|
|
|
|
traced_func = maybe_traced_gateway_call(
|
|
mock_async_stream,
|
|
endpoint_config,
|
|
output_reducer=aggregate_chat_stream_chunks,
|
|
)
|
|
|
|
# Consume the stream
|
|
chunks = [
|
|
chunk async for chunk in traced_func({"messages": [{"role": "user", "content": "hi"}]})
|
|
]
|
|
assert len(chunks) == 3
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
span_name_to_span = {span.name: span for span in trace.data.spans}
|
|
gateway_span = span_name_to_span[f"gateway/{endpoint_config.endpoint_name}"]
|
|
|
|
# Input should be unwrapped (not nested under "payload" key)
|
|
assert gateway_span.inputs == {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
# The output should be the reduced aggregated response, not raw chunks
|
|
output = gateway_span.outputs
|
|
assert output["object"] == "chat.completion"
|
|
assert output["choices"][0]["message"]["content"] == "Hello world"
|
|
assert output["choices"][0]["finish_reason"] == "stop"
|
|
assert output["usage"]["total_tokens"] == 7
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_message_format(endpoint_config):
|
|
async def mock_async_func(payload):
|
|
return {"id": "msg_1", "type": "message", "role": "assistant", "content": []}
|
|
|
|
traced_func = maybe_traced_gateway_call(
|
|
mock_async_func,
|
|
endpoint_config,
|
|
message_format="anthropic",
|
|
)
|
|
await traced_func({"messages": [{"role": "user", "content": "hi"}]})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
span_name_to_span = {span.name: span for span in trace.data.spans}
|
|
gateway_span = span_name_to_span[f"gateway/{endpoint_config.endpoint_name}"]
|
|
assert gateway_span.get_attribute("mlflow.message.format") == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_payload_kwarg(endpoint_config):
|
|
async def mock_passthrough_func(action, payload, headers=None):
|
|
return {"result": "success", "action": action, "payload": payload}
|
|
|
|
traced_func = maybe_traced_gateway_call(mock_passthrough_func, endpoint_config)
|
|
result = await traced_func(
|
|
action="test_action", payload={"messages": [{"role": "user", "content": "hi"}]}, headers={}
|
|
)
|
|
|
|
assert result["result"] == "success"
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
span_name_to_span = {span.name: span for span in trace.data.spans}
|
|
gateway_span = span_name_to_span[f"gateway/{endpoint_config.endpoint_name}"]
|
|
|
|
# Input should be unwrapped to just the payload dict
|
|
assert gateway_span.inputs == {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for distributed tracing helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_model_span_info_reads_child_span(endpoint_config):
|
|
async def func_with_child_span(payload):
|
|
with mlflow.start_span("provider/openai/gpt-4", span_type=SpanType.LLM) as child:
|
|
child.set_attributes({
|
|
SpanAttributeKey.CHAT_USAGE: {
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
SpanAttributeKey.MODEL: "gpt-4",
|
|
SpanAttributeKey.MODEL_PROVIDER: "openai",
|
|
})
|
|
return {"result": "ok"}
|
|
|
|
traced = maybe_traced_gateway_call(func_with_child_span, endpoint_config)
|
|
await traced({"input": "test"})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
gateway_trace_id = traces[0].info.trace_id
|
|
|
|
# After the trace is exported, spans are removed from InMemoryTraceManager,
|
|
# so we expect empty here. The actual reading happens inside the wrapper
|
|
# while the trace is still in memory.
|
|
assert _get_model_span_info(gateway_trace_id) == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration tests for distributed tracing via traceparent
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_traceparent(gateway_experiment_id):
|
|
ep_config = GatewayEndpointConfig(
|
|
endpoint_id="test-endpoint-id",
|
|
endpoint_name="test-endpoint",
|
|
experiment_id=gateway_experiment_id,
|
|
usage_tracking=True,
|
|
models=[],
|
|
)
|
|
|
|
async def func_with_usage(payload):
|
|
with mlflow.start_span("provider/openai/gpt-4", span_type=SpanType.LLM) as child:
|
|
child.set_attributes({
|
|
SpanAttributeKey.CHAT_USAGE: {
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
SpanAttributeKey.MODEL: "gpt-4",
|
|
SpanAttributeKey.MODEL_PROVIDER: "openai",
|
|
})
|
|
return {"result": "ok"}
|
|
|
|
# Step 1: Agent creates span and generates traceparent headers
|
|
with mlflow.start_span("agent-root") as agent_span:
|
|
headers = get_tracing_context_headers_for_http_request()
|
|
agent_trace_id = agent_span.trace_id
|
|
agent_span_id = agent_span.span_id
|
|
|
|
# Step 2: Gateway processes request (no active agent span, simulating separate server)
|
|
traced = maybe_traced_gateway_call(func_with_usage, ep_config, request_headers=headers)
|
|
result = await traced({"input": "test"})
|
|
|
|
assert result == {"result": "ok"}
|
|
|
|
# Flush to ensure all spans are written (batch processor may be active)
|
|
mlflow.flush_trace_async_logging(terminate=True)
|
|
|
|
# Gateway trace should exist in the gateway experiment
|
|
gateway_traces = TracingClient().search_traces(locations=[gateway_experiment_id])
|
|
assert len(gateway_traces) == 1
|
|
gateway_trace_id = gateway_traces[0].info.trace_id
|
|
|
|
# The gateway trace should be separate from the agent trace
|
|
assert gateway_trace_id != agent_trace_id
|
|
|
|
# Agent trace should contain two distributed spans (gateway + provider)
|
|
agent_trace = mlflow.get_trace(agent_trace_id)
|
|
assert agent_trace is not None
|
|
|
|
spans_by_name = {s.name: s for s in agent_trace.data.spans}
|
|
assert "agent-root" in spans_by_name
|
|
assert f"gateway/{ep_config.endpoint_name}" in spans_by_name
|
|
assert "provider/openai/gpt-4" in spans_by_name
|
|
|
|
# Gateway span: child of agent root, has endpoint attrs + link
|
|
gw_span = spans_by_name[f"gateway/{ep_config.endpoint_name}"]
|
|
assert gw_span.parent_id == agent_span_id
|
|
assert gw_span.attributes.get("endpoint_id") == ep_config.endpoint_id
|
|
assert gw_span.attributes.get("endpoint_name") == ep_config.endpoint_name
|
|
assert gw_span.attributes.get(SpanAttributeKey.LINKED_GATEWAY_TRACE_ID) == gateway_trace_id
|
|
|
|
# Provider span: child of gateway span, has provider attrs
|
|
provider_span = spans_by_name["provider/openai/gpt-4"]
|
|
assert provider_span.parent_id == gw_span.span_id
|
|
assert provider_span.attributes.get(SpanAttributeKey.CHAT_USAGE) == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
}
|
|
assert provider_span.attributes.get(SpanAttributeKey.MODEL) == "gpt-4"
|
|
assert provider_span.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "openai"
|
|
|
|
# Provider span should preserve timing from the gateway trace
|
|
gateway_provider_span = next(
|
|
s for s in gateway_traces[0].data.spans if s.name == "provider/openai/gpt-4"
|
|
)
|
|
assert provider_span.start_time_ns == gateway_provider_span.start_time_ns
|
|
assert provider_span.end_time_ns == gateway_provider_span.end_time_ns
|
|
|
|
# Neither span should have request/response payloads
|
|
assert gw_span.inputs is None
|
|
assert gw_span.outputs is None
|
|
assert provider_span.inputs is None
|
|
assert provider_span.outputs is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_streaming_with_traceparent(gateway_experiment_id):
|
|
ep_config = GatewayEndpointConfig(
|
|
endpoint_id="test-endpoint-id",
|
|
endpoint_name="test-endpoint",
|
|
experiment_id=gateway_experiment_id,
|
|
usage_tracking=True,
|
|
models=[],
|
|
)
|
|
|
|
async def mock_stream_with_usage(payload):
|
|
with mlflow.start_span("provider/openai/gpt-4", span_type=SpanType.LLM) as child:
|
|
child.set_attributes({
|
|
SpanAttributeKey.CHAT_USAGE: {
|
|
"input_tokens": 20,
|
|
"output_tokens": 10,
|
|
"total_tokens": 30,
|
|
},
|
|
SpanAttributeKey.MODEL: "gpt-4",
|
|
SpanAttributeKey.MODEL_PROVIDER: "openai",
|
|
})
|
|
yield _make_chunk(content="Hello")
|
|
yield _make_chunk(content=" world", finish_reason="stop")
|
|
|
|
# Agent creates headers
|
|
with mlflow.start_span("agent-root") as agent_span:
|
|
headers = get_tracing_context_headers_for_http_request()
|
|
agent_trace_id = agent_span.trace_id
|
|
agent_span_id = agent_span.span_id
|
|
|
|
# Gateway processes request (separate context)
|
|
traced = maybe_traced_gateway_call(
|
|
mock_stream_with_usage,
|
|
ep_config,
|
|
output_reducer=aggregate_chat_stream_chunks,
|
|
request_headers=headers,
|
|
)
|
|
chunks = [chunk async for chunk in traced({"input": "test"})]
|
|
|
|
assert len(chunks) == 2
|
|
|
|
# Flush to ensure all spans are written (batch processor may be active)
|
|
mlflow.flush_trace_async_logging(terminate=True)
|
|
|
|
# Gateway trace should exist
|
|
gateway_traces = TracingClient().search_traces(locations=[gateway_experiment_id])
|
|
assert len(gateway_traces) == 1
|
|
gateway_trace_id = gateway_traces[0].info.trace_id
|
|
assert gateway_trace_id != agent_trace_id
|
|
|
|
# Agent trace should contain two distributed spans (gateway + provider)
|
|
agent_trace = mlflow.get_trace(agent_trace_id)
|
|
assert agent_trace is not None
|
|
|
|
spans_by_name = {s.name: s for s in agent_trace.data.spans}
|
|
assert "agent-root" in spans_by_name
|
|
assert f"gateway/{ep_config.endpoint_name}" in spans_by_name
|
|
assert "provider/openai/gpt-4" in spans_by_name
|
|
|
|
# Gateway span: child of agent root, has endpoint attrs + link
|
|
gw_span = spans_by_name[f"gateway/{ep_config.endpoint_name}"]
|
|
assert gw_span.parent_id == agent_span_id
|
|
assert gw_span.attributes.get("endpoint_id") == ep_config.endpoint_id
|
|
assert gw_span.attributes.get("endpoint_name") == ep_config.endpoint_name
|
|
assert gw_span.attributes.get(SpanAttributeKey.LINKED_GATEWAY_TRACE_ID) == gateway_trace_id
|
|
|
|
# Provider span: child of gateway span, has provider attrs
|
|
provider_span = spans_by_name["provider/openai/gpt-4"]
|
|
assert provider_span.parent_id == gw_span.span_id
|
|
assert provider_span.attributes.get(SpanAttributeKey.CHAT_USAGE) == {
|
|
"input_tokens": 20,
|
|
"output_tokens": 10,
|
|
"total_tokens": 30,
|
|
}
|
|
assert provider_span.attributes.get(SpanAttributeKey.MODEL) == "gpt-4"
|
|
assert provider_span.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "openai"
|
|
|
|
# Provider span should preserve timing from the gateway trace
|
|
gateway_provider_span = next(
|
|
s for s in gateway_traces[0].data.spans if s.name == "provider/openai/gpt-4"
|
|
)
|
|
assert provider_span.start_time_ns == gateway_provider_span.start_time_ns
|
|
assert provider_span.end_time_ns == gateway_provider_span.end_time_ns
|
|
|
|
# Neither span should have request/response payloads
|
|
assert gw_span.inputs is None
|
|
assert gw_span.outputs is None
|
|
assert provider_span.inputs is None
|
|
assert provider_span.outputs is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_traced_gateway_call_with_traceparent_multiple_providers(gateway_experiment_id):
|
|
ep_config = GatewayEndpointConfig(
|
|
endpoint_id="test-endpoint-id",
|
|
endpoint_name="test-endpoint",
|
|
experiment_id=gateway_experiment_id,
|
|
usage_tracking=True,
|
|
models=[],
|
|
)
|
|
|
|
async def func_with_multiple_providers(payload):
|
|
with mlflow.start_span("provider/openai/gpt-4", span_type=SpanType.LLM) as child:
|
|
child.set_attributes({
|
|
SpanAttributeKey.CHAT_USAGE: {
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
SpanAttributeKey.MODEL: "gpt-4",
|
|
SpanAttributeKey.MODEL_PROVIDER: "openai",
|
|
})
|
|
with mlflow.start_span("provider/anthropic/claude-3", span_type=SpanType.LLM) as child:
|
|
child.set_attributes({
|
|
SpanAttributeKey.CHAT_USAGE: {
|
|
"input_tokens": 20,
|
|
"output_tokens": 10,
|
|
"total_tokens": 30,
|
|
},
|
|
SpanAttributeKey.MODEL: "claude-3",
|
|
SpanAttributeKey.MODEL_PROVIDER: "anthropic",
|
|
})
|
|
return {"result": "ok"}
|
|
|
|
with mlflow.start_span("agent-root") as agent_span:
|
|
headers = get_tracing_context_headers_for_http_request()
|
|
agent_trace_id = agent_span.trace_id
|
|
|
|
traced = maybe_traced_gateway_call(
|
|
func_with_multiple_providers, ep_config, request_headers=headers
|
|
)
|
|
await traced({"input": "test"})
|
|
|
|
mlflow.flush_trace_async_logging()
|
|
agent_trace = mlflow.get_trace(agent_trace_id)
|
|
assert agent_trace is not None
|
|
|
|
spans_by_name = {s.name: s for s in agent_trace.data.spans}
|
|
gw_span = spans_by_name[f"gateway/{ep_config.endpoint_name}"]
|
|
|
|
# Both provider spans should be children of the gateway span
|
|
provider_openai = spans_by_name["provider/openai/gpt-4"]
|
|
assert provider_openai.parent_id == gw_span.span_id
|
|
assert provider_openai.attributes.get(SpanAttributeKey.MODEL) == "gpt-4"
|
|
assert provider_openai.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "openai"
|
|
assert provider_openai.attributes.get(SpanAttributeKey.CHAT_USAGE) == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
}
|
|
|
|
provider_anthropic = spans_by_name["provider/anthropic/claude-3"]
|
|
assert provider_anthropic.parent_id == gw_span.span_id
|
|
assert provider_anthropic.attributes.get(SpanAttributeKey.MODEL) == "claude-3"
|
|
assert provider_anthropic.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "anthropic"
|
|
assert provider_anthropic.attributes.get(SpanAttributeKey.CHAT_USAGE) == {
|
|
"input_tokens": 20,
|
|
"output_tokens": 10,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
# Provider spans should preserve timing from the gateway trace
|
|
gateway_traces = TracingClient().search_traces(locations=[gateway_experiment_id])
|
|
assert len(gateway_traces) == 1
|
|
gw_spans_by_name = {s.name: s for s in gateway_traces[0].data.spans}
|
|
|
|
gw_openai = gw_spans_by_name["provider/openai/gpt-4"]
|
|
assert provider_openai.start_time_ns == gw_openai.start_time_ns
|
|
assert provider_openai.end_time_ns == gw_openai.end_time_ns
|
|
|
|
gw_anthropic = gw_spans_by_name["provider/anthropic/claude-3"]
|
|
assert provider_anthropic.start_time_ns == gw_anthropic.start_time_ns
|
|
assert provider_anthropic.end_time_ns == gw_anthropic.end_time_ns
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for aggregate_anthropic_messages_stream_chunks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _sse(event: dict[str, Any]) -> bytes:
|
|
"""Encode a single event dict as an SSE data line."""
|
|
return f"data: {json.dumps(event)}\n".encode()
|
|
|
|
|
|
def _msg_start(msg_id: str, model: str, input_tokens: int | None = None) -> bytes:
|
|
usage = {"input_tokens": input_tokens} if input_tokens is not None else {}
|
|
return _sse({
|
|
"type": "message_start",
|
|
"message": {"id": msg_id, "model": model, "role": "assistant", "usage": usage},
|
|
})
|
|
|
|
|
|
def _text_block_start(index: int) -> bytes:
|
|
return _sse({
|
|
"type": "content_block_start",
|
|
"index": index,
|
|
"content_block": {"type": "text", "text": ""},
|
|
})
|
|
|
|
|
|
def _text_delta(index: int, text: str) -> bytes:
|
|
return _sse({
|
|
"type": "content_block_delta",
|
|
"index": index,
|
|
"delta": {"type": "text_delta", "text": text},
|
|
})
|
|
|
|
|
|
def _tool_block_start(index: int, tool_id: str, name: str) -> bytes:
|
|
return _sse({
|
|
"type": "content_block_start",
|
|
"index": index,
|
|
"content_block": {"type": "tool_use", "id": tool_id, "name": name, "input": {}},
|
|
})
|
|
|
|
|
|
def _tool_delta(index: int, partial_json: str) -> bytes:
|
|
return _sse({
|
|
"type": "content_block_delta",
|
|
"index": index,
|
|
"delta": {"type": "input_json_delta", "partial_json": partial_json},
|
|
})
|
|
|
|
|
|
def _msg_delta(stop_reason: str, output_tokens: int, stop_sequence: str | None = None) -> bytes:
|
|
return _sse({
|
|
"type": "message_delta",
|
|
"delta": {"stop_reason": stop_reason, "stop_sequence": stop_sequence},
|
|
"usage": {"output_tokens": output_tokens},
|
|
})
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_empty():
|
|
assert aggregate_anthropic_messages_stream_chunks([]) is None
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_no_parseable_events():
|
|
chunks = [b"event: ping\n", b"data: [DONE]\n"]
|
|
assert aggregate_anthropic_messages_stream_chunks(chunks) is None
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_text():
|
|
chunks = [
|
|
_msg_start("msg_1", "claude-3-5-sonnet-20241022", input_tokens=10),
|
|
_text_block_start(0),
|
|
_text_delta(0, "Hello"),
|
|
_text_delta(0, " world"),
|
|
_sse({"type": "content_block_stop", "index": 0}),
|
|
_msg_delta("end_turn", output_tokens=5, stop_sequence=None),
|
|
_sse({"type": "message_stop"}),
|
|
]
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
|
|
assert result["id"] == "msg_1"
|
|
assert result["type"] == "message"
|
|
assert result["role"] == "assistant"
|
|
assert result["model"] == "claude-3-5-sonnet-20241022"
|
|
assert result["stop_reason"] == "end_turn"
|
|
assert result["stop_sequence"] is None
|
|
assert result["content"] == [{"type": "text", "text": "Hello world"}]
|
|
assert result["usage"] == {"input_tokens": 10, "output_tokens": 5}
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_tool_use():
|
|
chunks = [
|
|
_msg_start("msg_2", "claude-3-5-sonnet-20241022", input_tokens=20),
|
|
_tool_block_start(0, "toolu_abc", "get_weather"),
|
|
_tool_delta(0, '{"city"'),
|
|
_tool_delta(0, ': "Paris"}'),
|
|
_sse({"type": "content_block_stop", "index": 0}),
|
|
_msg_delta("tool_use", output_tokens=15, stop_sequence=None),
|
|
]
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
|
|
assert result["stop_reason"] == "tool_use"
|
|
assert len(result["content"]) == 1
|
|
block = result["content"][0]
|
|
assert block["type"] == "tool_use"
|
|
assert block["id"] == "toolu_abc"
|
|
assert block["name"] == "get_weather"
|
|
assert block["input"] == {"city": "Paris"}
|
|
assert result["usage"] == {"input_tokens": 20, "output_tokens": 15}
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_mixed_content():
|
|
chunks = [
|
|
_msg_start("msg_3", "claude-3-5-sonnet-20241022", input_tokens=30),
|
|
_text_block_start(0),
|
|
_text_delta(0, "Let me check that."),
|
|
_sse({"type": "content_block_stop", "index": 0}),
|
|
_tool_block_start(1, "toolu_xyz", "search"),
|
|
_tool_delta(1, '{"q": "mlflow"}'),
|
|
_sse({"type": "content_block_stop", "index": 1}),
|
|
_msg_delta("tool_use", output_tokens=25, stop_sequence=None),
|
|
]
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
|
|
assert len(result["content"]) == 2
|
|
assert result["content"][0] == {"type": "text", "text": "Let me check that."}
|
|
assert result["content"][1] == {
|
|
"type": "tool_use",
|
|
"id": "toolu_xyz",
|
|
"name": "search",
|
|
"input": {"q": "mlflow"},
|
|
}
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_multiple_chunks_per_sse():
|
|
# Multiple SSE events packed into one bytes chunk (newline-separated)
|
|
combined = (
|
|
_msg_start("msg_4", "claude-3-5-sonnet-20241022", input_tokens=5)
|
|
+ _text_block_start(0)
|
|
+ _text_delta(0, "Hi")
|
|
+ _msg_delta("end_turn", output_tokens=2, stop_sequence=None)
|
|
)
|
|
result = aggregate_anthropic_messages_stream_chunks([combined])
|
|
|
|
assert result["id"] == "msg_4"
|
|
assert result["content"] == [{"type": "text", "text": "Hi"}]
|
|
assert result["usage"] == {"input_tokens": 5, "output_tokens": 2}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("raw_json", "expected_input"),
|
|
[
|
|
('{"key": "val"}', {"key": "val"}),
|
|
("", {}),
|
|
("not-valid-json", {}),
|
|
],
|
|
)
|
|
def test_aggregate_anthropic_messages_stream_chunks_tool_input_edge_cases(raw_json, expected_input):
|
|
chunks = [
|
|
_msg_start("msg_5", "claude-3-5-sonnet-20241022"),
|
|
_tool_block_start(0, "t1", "fn"),
|
|
]
|
|
if raw_json:
|
|
chunks.append(_tool_delta(0, raw_json))
|
|
chunks.append(_msg_delta("tool_use", output_tokens=1))
|
|
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
assert result["content"][0]["input"] == expected_input
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_split_sse_lines():
|
|
# Simulate an aiohttp byte chunk that splits a "data:" SSE line mid-way.
|
|
# All events should still be parsed correctly after concatenation.
|
|
msg_start_bytes = _msg_start("msg_split", "claude-3-5-sonnet-20241022", input_tokens=3)
|
|
mid = len(msg_start_bytes) // 2
|
|
chunks = [
|
|
msg_start_bytes[:mid],
|
|
msg_start_bytes[mid:] + _msg_delta("end_turn", output_tokens=1),
|
|
]
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
|
|
assert result is not None
|
|
assert result["id"] == "msg_split"
|
|
assert result["usage"] == {"input_tokens": 3, "output_tokens": 1}
|
|
|
|
|
|
def test_aggregate_anthropic_messages_stream_chunks_cache_tokens():
|
|
chunks = [
|
|
_sse({
|
|
"type": "message_start",
|
|
"message": {
|
|
"id": "msg_cache",
|
|
"model": "claude-3-5-sonnet-20241022",
|
|
"role": "assistant",
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"cache_read_input_tokens": 5,
|
|
"cache_creation_input_tokens": 2,
|
|
},
|
|
},
|
|
}),
|
|
_msg_delta("end_turn", output_tokens=8),
|
|
]
|
|
result = aggregate_anthropic_messages_stream_chunks(chunks)
|
|
|
|
assert result["usage"] == {
|
|
"input_tokens": 10,
|
|
"cache_read_input_tokens": 5,
|
|
"cache_creation_input_tokens": 2,
|
|
"output_tokens": 8,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for aggregate_openai_responses_stream_chunks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_RESPONSES_CREATED = (
|
|
b'data: {"type":"response.created","response":{"id":"resp_1","object":"response",'
|
|
b'"created_at":1741290958,"status":"in_progress","output":[],"usage":null}}\n'
|
|
)
|
|
_RESPONSES_TEXT_DELTA = (
|
|
b'data: {"type":"response.output_text.delta","item_id":"msg_1",'
|
|
b'"output_index":0,"content_index":0,"delta":"Hi"}\n'
|
|
)
|
|
_RESPONSES_TEXT_DONE = (
|
|
b'data: {"type":"response.output_text.done","item_id":"msg_1",'
|
|
b'"output_index":0,"content_index":0,"text":"Hi there!"}\n'
|
|
)
|
|
_RESPONSES_COMPLETED = (
|
|
b'data: {"type":"response.completed","response":{"id":"resp_1","object":"response",'
|
|
b'"created_at":1741290958,"status":"completed",'
|
|
b'"output":[{"id":"msg_1","type":"message","status":"completed","role":"assistant",'
|
|
b'"content":[{"type":"output_text","text":"Hi there!","annotations":[]}]}],'
|
|
b'"usage":{"input_tokens":37,"output_tokens":11,"total_tokens":48}}}\n'
|
|
)
|
|
|
|
|
|
def test_aggregate_openai_responses_stream_chunks_empty():
|
|
assert aggregate_openai_responses_stream_chunks([]) is None
|
|
|
|
|
|
def test_aggregate_openai_responses_stream_chunks_no_completed_event():
|
|
chunks = [_RESPONSES_CREATED, _RESPONSES_TEXT_DELTA]
|
|
assert aggregate_openai_responses_stream_chunks(chunks) is None
|
|
|
|
|
|
def test_aggregate_openai_responses_stream_chunks_basic():
|
|
chunks = [
|
|
_RESPONSES_CREATED,
|
|
_RESPONSES_TEXT_DELTA,
|
|
_RESPONSES_TEXT_DONE,
|
|
_RESPONSES_COMPLETED,
|
|
]
|
|
result = aggregate_openai_responses_stream_chunks(chunks)
|
|
|
|
assert result["id"] == "resp_1"
|
|
assert result["object"] == "response"
|
|
assert result["status"] == "completed"
|
|
assert len(result["output"]) == 1
|
|
assert result["output"][0]["role"] == "assistant"
|
|
assert result["output"][0]["content"][0]["text"] == "Hi there!"
|
|
assert result["usage"] == {"input_tokens": 37, "output_tokens": 11, "total_tokens": 48}
|
|
|
|
|
|
def test_aggregate_openai_responses_stream_chunks_split_sse_lines():
|
|
# Simulate aiohttp yielding a chunk that splits the data: line mid-way.
|
|
mid = len(_RESPONSES_COMPLETED) // 2
|
|
chunks = [
|
|
_RESPONSES_CREATED,
|
|
_RESPONSES_COMPLETED[:mid],
|
|
_RESPONSES_COMPLETED[mid:],
|
|
]
|
|
result = aggregate_openai_responses_stream_chunks(chunks)
|
|
|
|
assert result is not None
|
|
assert result["id"] == "resp_1"
|
|
assert result["status"] == "completed"
|
|
|
|
|
|
def test_aggregate_openai_responses_stream_chunks_returns_completed_response():
|
|
# When multiple events are packed into a single bytes chunk, the
|
|
# completed response is still extracted correctly.
|
|
combined = _RESPONSES_CREATED + _RESPONSES_TEXT_DELTA + _RESPONSES_COMPLETED
|
|
result = aggregate_openai_responses_stream_chunks([combined])
|
|
|
|
assert result["status"] == "completed"
|
|
assert result["usage"]["total_tokens"] == 48
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for aggregate_gemini_stream_generate_content_chunks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _gemini_sse(event: dict[str, Any]) -> bytes:
|
|
return f"data: {json.dumps(event)}\n".encode()
|
|
|
|
|
|
def _gemini_text_chunk(text: str, finish_reason: str | None = None) -> bytes:
|
|
candidate: dict[str, Any] = {"content": {"parts": [{"text": text}], "role": "model"}}
|
|
if finish_reason:
|
|
candidate["finishReason"] = finish_reason
|
|
return _gemini_sse({"candidates": [candidate]})
|
|
|
|
|
|
def test_aggregate_gemini_stream_chunks_empty():
|
|
assert aggregate_gemini_stream_generate_content_chunks([]) is None
|
|
|
|
|
|
def test_aggregate_gemini_stream_chunks_no_parseable_events():
|
|
chunks = [b"event: ping\n", b"data: [DONE]\n"]
|
|
assert aggregate_gemini_stream_generate_content_chunks(chunks) is None
|
|
|
|
|
|
def test_aggregate_gemini_stream_chunks_text():
|
|
chunks = [
|
|
_gemini_text_chunk("Hello"),
|
|
_gemini_text_chunk(" world", finish_reason="STOP"),
|
|
_gemini_sse({
|
|
"usageMetadata": {
|
|
"promptTokenCount": 10,
|
|
"candidatesTokenCount": 5,
|
|
"totalTokenCount": 15,
|
|
}
|
|
}),
|
|
]
|
|
result = aggregate_gemini_stream_generate_content_chunks(chunks)
|
|
|
|
assert len(result["candidates"]) == 1
|
|
cand = result["candidates"][0]
|
|
assert cand["content"]["parts"] == [{"text": "Hello world"}]
|
|
assert cand["content"]["role"] == "model"
|
|
assert cand["finishReason"] == "STOP"
|
|
assert result["usageMetadata"] == {
|
|
"promptTokenCount": 10,
|
|
"candidatesTokenCount": 5,
|
|
"totalTokenCount": 15,
|
|
}
|
|
|
|
|
|
def test_aggregate_gemini_stream_chunks_tool_call():
|
|
chunks = [
|
|
_gemini_sse({
|
|
"candidates": [
|
|
{
|
|
"content": {
|
|
"parts": [
|
|
{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}
|
|
],
|
|
"role": "model",
|
|
},
|
|
"finishReason": "STOP",
|
|
"index": 0,
|
|
}
|
|
]
|
|
}),
|
|
_gemini_sse({
|
|
"usageMetadata": {
|
|
"promptTokenCount": 8,
|
|
"candidatesTokenCount": 12,
|
|
"totalTokenCount": 20,
|
|
}
|
|
}),
|
|
]
|
|
result = aggregate_gemini_stream_generate_content_chunks(chunks)
|
|
|
|
cand = result["candidates"][0]
|
|
assert cand["content"]["parts"] == [
|
|
{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}
|
|
]
|
|
assert cand["finishReason"] == "STOP"
|
|
|
|
|
|
def test_aggregate_gemini_stream_chunks_split_sse_lines():
|
|
chunk_bytes = _gemini_text_chunk("Hi", finish_reason="STOP")
|
|
mid = len(chunk_bytes) // 2
|
|
result = aggregate_gemini_stream_generate_content_chunks([chunk_bytes[:mid], chunk_bytes[mid:]])
|
|
|
|
assert result is not None
|
|
assert result["candidates"][0]["content"]["parts"] == [{"text": "Hi"}]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("finish_reasons", "expected"),
|
|
[
|
|
([None, None, "STOP"], "STOP"),
|
|
([None, "stop", None], "stop"),
|
|
([None, None, None], None),
|
|
],
|
|
)
|
|
def test_aggregate_gemini_stream_chunks_finish_reason(finish_reasons, expected):
|
|
chunks = [_gemini_text_chunk(f"t{i}", finish_reason=fr) for i, fr in enumerate(finish_reasons)]
|
|
result = aggregate_gemini_stream_generate_content_chunks(chunks)
|
|
assert result["candidates"][0]["finishReason"] == expected
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _extract_caller tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("headers", "expected"),
|
|
[
|
|
(None, None),
|
|
({}, None),
|
|
({"User-Agent": "openai-python/1.0.0"}, "openai-python"),
|
|
({"user-agent": "claude-cli/1.2.3"}, "claude-cli"),
|
|
({"User-Agent": "GeminiCLI/1.0 (Linux; x64)"}, "GeminiCLI"),
|
|
({"User-Agent": "anthropic/0.20.0 CPython/3.11.0 Darwin/23.0.0"}, "anthropic"),
|
|
(
|
|
{MLFLOW_GATEWAY_CALLER_HEADER: "judge", "User-Agent": "openai-python/1.0.0"},
|
|
"judge",
|
|
),
|
|
({MLFLOW_GATEWAY_CALLER_HEADER: "judge"}, "judge"),
|
|
({"User-Agent": " "}, None),
|
|
],
|
|
)
|
|
def test_extract_caller(headers, expected):
|
|
assert _extract_caller(headers) == expected
|
|
|
|
|
|
def test_maybe_traced_gateway_call_records_caller(endpoint_config):
|
|
async def fake_func(payload):
|
|
return {"ok": True}
|
|
|
|
traced = maybe_traced_gateway_call(
|
|
fake_func,
|
|
endpoint_config,
|
|
request_headers={"User-Agent": "openai-python/1.0.0"},
|
|
)
|
|
|
|
asyncio.run(traced({"prompt": "hi"}))
|
|
|
|
traces = get_traces()
|
|
assert traces
|
|
assert traces[0].info.request_metadata.get(TraceMetadataKey.GATEWAY_CALLER) == "openai-python"
|
|
|
|
|
|
def test_maybe_traced_gateway_call_no_caller_when_no_headers(endpoint_config):
|
|
async def fake_func(payload):
|
|
return {"ok": True}
|
|
|
|
traced = maybe_traced_gateway_call(fake_func, endpoint_config, request_headers=None)
|
|
|
|
asyncio.run(traced({"prompt": "hi"}))
|
|
|
|
traces = get_traces()
|
|
assert traces
|
|
assert TraceMetadataKey.GATEWAY_CALLER not in (traces[0].info.request_metadata or {})
|