mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
test(otel): wait for each exported span before the next request in Arize OTel v2 audit tests
This commit is contained in:
parent
4b0d28f704
commit
e6f66c97b1
3 changed files with 76 additions and 34 deletions
|
|
@ -561,55 +561,65 @@ def _matching_marker_spans(requests: tuple[Request, ...], marker: str) -> Iterat
|
|||
yield attributes
|
||||
|
||||
|
||||
def _single_span(spans: tuple[dict[str, str], ...]) -> dict[str, str]:
|
||||
assert len(spans) == 1, spans
|
||||
return spans[0]
|
||||
|
||||
|
||||
def _matching_span(destination: Wire, response_id: str) -> dict[str, str]:
|
||||
return eventually(
|
||||
spans: Final = eventually(
|
||||
lambda: tuple(_matching_llm_spans(destination.drain(), response_id)),
|
||||
lambda spans: len(spans) == 1,
|
||||
bool,
|
||||
seconds=30,
|
||||
)[0]
|
||||
)
|
||||
return _single_span(spans)
|
||||
|
||||
|
||||
def _matching_marker_span(destination: Wire, marker: str) -> dict[str, str]:
|
||||
return eventually(
|
||||
spans: Final = eventually(
|
||||
lambda: tuple(_matching_marker_spans(destination.drain(), marker)),
|
||||
lambda spans: len(spans) == 1,
|
||||
bool,
|
||||
seconds=30,
|
||||
)[0]
|
||||
)
|
||||
return _single_span(spans)
|
||||
|
||||
|
||||
def _matching_output_value_span(destination: Wire, marker: str) -> dict[str, str]:
|
||||
return eventually(
|
||||
spans: Final = eventually(
|
||||
lambda: tuple(
|
||||
attributes
|
||||
for attributes in _spans(destination.drain())
|
||||
if attributes.get("openinference.span.kind") == "LLM" and marker in attributes.get("output.value", "")
|
||||
),
|
||||
lambda spans: len(spans) == 1,
|
||||
bool,
|
||||
seconds=30,
|
||||
)[0]
|
||||
)
|
||||
return _single_span(spans)
|
||||
|
||||
|
||||
def _matching_genai_marker_span(destination: Wire, marker: str) -> dict[str, str]:
|
||||
return eventually(
|
||||
spans: Final = eventually(
|
||||
lambda: tuple(
|
||||
attributes
|
||||
for attributes in _spans(destination.drain())
|
||||
if attributes.get("gen_ai.operation.name") == "chat" and marker in attributes.values()
|
||||
),
|
||||
lambda spans: len(spans) == 1,
|
||||
bool,
|
||||
seconds=30,
|
||||
)[0]
|
||||
)
|
||||
return _single_span(spans)
|
||||
|
||||
|
||||
def _matching_any_marker_span(destination: Wire, marker: str) -> dict[str, str]:
|
||||
def matches(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]:
|
||||
return tuple(attributes for attributes in _spans(requests) if marker in attributes.values())
|
||||
|
||||
return eventually(
|
||||
spans: Final = eventually(
|
||||
lambda: matches(destination.drain()),
|
||||
lambda spans: len(spans) == 1,
|
||||
bool,
|
||||
seconds=30,
|
||||
)[0]
|
||||
)
|
||||
return _single_span(spans)
|
||||
|
||||
|
||||
def _collect_marker_spans(
|
||||
|
|
@ -636,6 +646,20 @@ def _collect_marker_spans(
|
|||
return collect(())
|
||||
|
||||
|
||||
def _llm_spans_through_marker(destination: Wire, marker: str) -> tuple[dict[str, str], ...]:
|
||||
def llm_spans(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]:
|
||||
return tuple(
|
||||
attributes for attributes in _spans(requests) if attributes.get("openinference.span.kind") == "LLM"
|
||||
)
|
||||
|
||||
def collect(previous: tuple[dict[str, str], ...]) -> tuple[dict[str, str], ...]:
|
||||
current: Final = eventually(lambda: llm_spans(destination.drain()), bool, seconds=30)
|
||||
combined: Final = (*previous, *current)
|
||||
return combined if any(marker in attributes.values() for attributes in combined) else collect(combined)
|
||||
|
||||
return collect(())
|
||||
|
||||
|
||||
def _write_config(
|
||||
directory: Path,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from _openinference_support import (
|
|||
_chat_tool_call,
|
||||
_collect_marker_spans,
|
||||
_json_object,
|
||||
_matching_marker_span,
|
||||
_messages_caller_raw_stream,
|
||||
_messages_caller_response,
|
||||
_normalize_chat_caller_stream,
|
||||
|
|
@ -297,14 +298,7 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path
|
|||
timer.cancel()
|
||||
if timer.ident is not None:
|
||||
timer.join(timeout=5)
|
||||
requests: Final = rig.destination.drain()
|
||||
spans: Final = tuple(
|
||||
attributes
|
||||
for attributes in _spans(requests)
|
||||
if attributes.get("openinference.span.kind") == "LLM"
|
||||
and attributes.get("litellm.metadata.trace_marker") == marker
|
||||
)
|
||||
assert len(spans) == 1, spans
|
||||
_matching_marker_span(rig.destination, marker)
|
||||
|
||||
|
||||
def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None:
|
||||
|
|
|
|||
|
|
@ -24,11 +24,13 @@ from _openinference_support import (
|
|||
_chat_caller_response,
|
||||
_chat_caller_stream,
|
||||
_chat_plain_response,
|
||||
_chat_request_marker,
|
||||
_chat_response,
|
||||
_chat_stream_response,
|
||||
_chat_tool_call,
|
||||
_json_messages,
|
||||
_json_object,
|
||||
_llm_spans_through_marker,
|
||||
_matching_marker_span,
|
||||
_messages_caller_response,
|
||||
_messages_caller_stream,
|
||||
|
|
@ -43,7 +45,7 @@ from _openinference_support import (
|
|||
_responses_stream_response,
|
||||
_rig,
|
||||
)
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import Gateway, object_value, string_value
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -844,27 +846,34 @@ def _cache_call(
|
|||
@pytest.mark.parametrize("surface", ("chat", "chat-stream", "responses", "messages"))
|
||||
def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = f"a-cache-{surface}-" + uuid.uuid4().hex
|
||||
sentinel: Final = f"{marker}-sentinel"
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
body: Final = _json_object(request.body)
|
||||
request_marker: Final = (
|
||||
_chat_request_marker(request)
|
||||
if surface in ("chat", "chat-stream", "messages")
|
||||
else string_value(object_value(body["metadata"])["trace_marker"])
|
||||
)
|
||||
assert request_marker in (marker, sentinel), request
|
||||
if surface in ("chat", "chat-stream"):
|
||||
_assert_chat_request(
|
||||
request,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
messages=[{"role": "user", "content": request_marker}],
|
||||
stream=True if surface == "chat-stream" else None,
|
||||
stream_options={"include_usage": False} if surface == "chat-stream" else None,
|
||||
)
|
||||
return (
|
||||
_chat_stream_response(marker, (_chat_tool_call(marker),), include_usage=False)
|
||||
_chat_stream_response(request_marker, (_chat_tool_call(request_marker),), include_usage=False)
|
||||
if surface == "chat-stream"
|
||||
else _chat_response(marker)
|
||||
else _chat_response(request_marker)
|
||||
)
|
||||
if surface == "responses":
|
||||
assert body.get("metadata") == {"trace_marker": marker}, body
|
||||
_assert_responses_request(request, marker=marker, input_value=marker)
|
||||
return _responses_response(marker)
|
||||
_assert_messages_request(request, marker=marker, prompt=marker)
|
||||
return _messages_response(marker)
|
||||
assert body.get("metadata") == {"trace_marker": request_marker}, body
|
||||
_assert_responses_request(request, marker=request_marker, input_value=request_marker)
|
||||
return _responses_response(request_marker)
|
||||
_assert_messages_request(request, marker=request_marker, prompt=request_marker)
|
||||
return _messages_response(request_marker)
|
||||
|
||||
with _rig(
|
||||
gateway,
|
||||
|
|
@ -879,7 +888,8 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
|
|||
assert first[0] == (f"call_{marker}", "lookup_weather", expected_arguments), first
|
||||
assert first[1].startswith("resp_") if surface == "responses" else first[1] == marker, first
|
||||
assert not first[2].get("x-litellm-cache-key"), first[2]
|
||||
rig.destination.drain()
|
||||
first_span: Final = _matching_marker_span(rig.destination, marker)
|
||||
_assert_tool_span_for_marker(first_span, marker)
|
||||
second: Final = _cache_call(rig, surface, marker, cache_hit=True)
|
||||
assert second[0] == first[0], second
|
||||
assert second[1].startswith("resp_") if surface == "responses" else second[1] == first[1], second
|
||||
|
|
@ -891,7 +901,6 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
|
|||
assert not second[2].get("x-litellm-cache-key"), second[2]
|
||||
else:
|
||||
assert second[2].get("x-litellm-cache-key"), second[2]
|
||||
_assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker)
|
||||
forwarded_body: Final = _json_object(forwarded[0].body)
|
||||
if surface in ("chat", "chat-stream"):
|
||||
assert "metadata" not in forwarded_body, forwarded[0]
|
||||
|
|
@ -899,3 +908,18 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
|
|||
assert forwarded_body["metadata"] == {"trace_marker": marker}, forwarded[0]
|
||||
else:
|
||||
assert forwarded_body["metadata"] == {}, forwarded[0]
|
||||
sentinel_response: Final = _cache_call(rig, surface, sentinel)
|
||||
assert not sentinel_response[2].get("x-litellm-cache-key"), sentinel_response[2]
|
||||
sentinel_forwarded: Final = tuple(
|
||||
request
|
||||
for request in rig.provider.drain()
|
||||
if request.method == "POST" and sentinel.encode() in request.body
|
||||
)
|
||||
assert len(sentinel_forwarded) == 1, sentinel_forwarded
|
||||
spans: Final = _llm_spans_through_marker(rig.destination, sentinel)
|
||||
sentinel_spans: Final = tuple(
|
||||
attributes for attributes in spans if attributes.get("litellm.metadata.trace_marker") == sentinel
|
||||
)
|
||||
assert len(sentinel_spans) == 1, spans
|
||||
_assert_tool_span_for_marker(sentinel_spans[0], sentinel)
|
||||
assert not any(attributes.get("litellm.metadata.trace_marker") == marker for attributes in spans), spans
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue