mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(proxy)!: honor caller litellm_session_id in spend log session_id (#44199)
* fix(responses): scope session history lookup to the referenced key * Implement SQLite-backed DB for session tests Added a SQLite-backed database class for testing and a new test to ensure session isolation between different API keys. * fix(proxy): honor caller litellm_session_id in spend log session_id Updated docstring to clarify session ID resolution logic for spend log rows. * test(proxy): caller litellm_session_id wins over the trace id * style(proxy): fix blank lines around new tests * fix(proxy): read the caller session from the logging payload for spend logs * fix(responses): match user and team too when replaying session history * fix(proxy): keep bridged Responses chains whole and scope session history by owner * fix(proxy): keep a metadata session next to a trace header and let session headers win * fix(proxy): send the trace-header session only to spend logs and let session headers win beside a trace header * fix(proxy): drop a client-sent caller session and keep it out of requester metadata * fix(proxy): keep the trace id as the spend log session when the caller session contains a NUL byte * fix(responses): keep a continued bridged chain's session under request_correlation_in_logs With the flag on, the caller session outranked the chain session the bridge sets for a continued Responses chain, so a chain whose session changed or first appeared on a later turn split. The continuation now logs under the chain's session and drops the turn's caller session key, as on main. Adds the /audit integration cells for the spend log caller session. * fix(responses): log a continued bridged chain under its session through the spend-log key only The previous commit set the chain session as the turn's litellm_session_id, which also moved the OTel conversation id of a continued turn away from its trace header. The bridge now writes the chain session into the internal spend-log session key of the metadata the spend log reads, so only the spend log session changes, as on main. Adds an OTel cell for a continued chain. * test(spend): cover bridged chains that share one caller session * test(spend): cover the caller session on configured pass-through routes * test(spend): type the SQLite fake and split stacked comprehensions in the caller session tests --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Deepanshu Pal
GitHub
parent
4e2dd65549
commit
d5ce7dbb8c
10 changed files with 1208 additions and 27 deletions
|
|
@ -1629,6 +1629,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
|||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
|
||||
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
|
||||
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
|
||||
CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id"
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
|
||||
"Truncation is a DB storage safeguard. "
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
|
|||
from litellm._service_logger import ServiceLogging
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
CLIENT_OUTPUT_CEILING_METADATA_KEY,
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
|
|
@ -187,6 +188,30 @@ def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str,
|
|||
return requester_value if isinstance(requester_value, str) and requester_value else None
|
||||
|
||||
|
||||
def _caller_body_session_id(data: Mapping[str, object], metadata_variable_name: str) -> str | None:
|
||||
top_level: Final = data.get("litellm_session_id")
|
||||
if isinstance(top_level, str) and top_level:
|
||||
return top_level
|
||||
return _caller_trace_field(data, metadata_variable_name, "session_id")
|
||||
|
||||
|
||||
def _caller_session_beside_trace_header(
|
||||
normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str
|
||||
) -> str | None:
|
||||
session_headers: Final = {key: value for key, value in normalized_headers.items() if key != "x-litellm-trace-id"}
|
||||
return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name)
|
||||
|
||||
|
||||
def _drop_client_caller_session(request_body: Mapping[str, object]) -> None:
|
||||
for metadata_key in ("metadata", "litellm_metadata"):
|
||||
if isinstance(metadata := request_body.get(metadata_key), dict):
|
||||
metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None)
|
||||
|
||||
|
||||
def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY})
|
||||
|
||||
|
||||
def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
||||
"""Only proxy-validated keys are stamped, proven by the unforgeable
|
||||
via_virtual_key marker AND a known non-secret shape: the sha256 hex digest
|
||||
|
|
@ -1623,6 +1648,8 @@ class LiteLLMProxyRequestSetup:
|
|||
from litellm.proxy._types import LitellmMetadataFromRequestHeaders
|
||||
|
||||
metadata_from_headers: Final = LitellmMetadataFromRequestHeaders()
|
||||
request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object
|
||||
_drop_client_caller_session(request_body)
|
||||
spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers)
|
||||
if spend_logs_metadata is not None:
|
||||
metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata
|
||||
|
|
@ -1631,6 +1658,9 @@ class LiteLLMProxyRequestSetup:
|
|||
# Finally update the requests metadata with the `metadata_from_headers`
|
||||
#########################################################################################
|
||||
|
||||
normalized_headers: Final = MappingProxyType(
|
||||
{k.lower(): v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)}
|
||||
)
|
||||
agent_id_from_header: Final = headers.get("x-litellm-agent-id")
|
||||
# Explicit litellm headers take precedence; fall back to any x-*-session-id header.
|
||||
chain_id: Final = get_chain_id_from_headers(dict(headers))
|
||||
|
|
@ -1639,6 +1669,13 @@ class LiteLLMProxyRequestSetup:
|
|||
metadata_from_headers["agent_id"] = agent_id_from_header
|
||||
verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header)
|
||||
|
||||
spend_log_session_id: Final = (
|
||||
_caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name)
|
||||
if litellm.request_correlation_in_logs
|
||||
and chain_id is not None
|
||||
and chain_id == normalized_headers.get("x-litellm-trace-id")
|
||||
else None
|
||||
)
|
||||
if chain_id:
|
||||
metadata_from_headers["trace_id"] = chain_id
|
||||
metadata_from_headers["session_id"] = chain_id
|
||||
|
|
@ -1663,7 +1700,6 @@ class LiteLLMProxyRequestSetup:
|
|||
# - but lets a caller's existing traceparent/baggage headers (from
|
||||
# real OTel instrumentation) correlate with litellm's own logs instead
|
||||
# of generating an unrelated trace_id.
|
||||
normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)})
|
||||
if "litellm_trace_id" not in data:
|
||||
traceparent: Final = normalized_headers.get("traceparent")
|
||||
if isinstance(traceparent, str):
|
||||
|
|
@ -1689,11 +1725,7 @@ class LiteLLMProxyRequestSetup:
|
|||
baggage: Final = normalized_headers.get("baggage")
|
||||
if isinstance(baggage, str):
|
||||
session_id_from_baggage: Final = _session_id_from_baggage(baggage)
|
||||
caller_session_id: Final = _caller_trace_field(
|
||||
cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
|
||||
_metadata_variable_name,
|
||||
"session_id",
|
||||
)
|
||||
caller_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id")
|
||||
if session_id_from_baggage and (
|
||||
caller_session_id is None or caller_session_id == session_id_from_baggage
|
||||
):
|
||||
|
|
@ -1702,7 +1734,11 @@ class LiteLLMProxyRequestSetup:
|
|||
verbose_proxy_logger.debug("Extracted session_id from W3C baggage header")
|
||||
|
||||
if isinstance(data[_metadata_variable_name], dict):
|
||||
data[_metadata_variable_name].update(metadata_from_headers)
|
||||
data[_metadata_variable_name].update(
|
||||
{**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id}
|
||||
if spend_log_session_id
|
||||
else metadata_from_headers
|
||||
)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -2417,7 +2453,7 @@ async def add_litellm_data_to_request(
|
|||
# them — from leaking into requester_metadata where guardrails and audit
|
||||
# paths may read from it.
|
||||
if "metadata" in data and isinstance(data["metadata"], dict):
|
||||
data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"])
|
||||
data[_metadata_variable_name]["requester_metadata"] = _requester_metadata_snapshot(data["metadata"])
|
||||
if _metadata_variable_name == "litellm_metadata":
|
||||
data[_metadata_variable_name].update(
|
||||
_promoted_trace_control_fields(
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
CLI_SESSION_KEY_PREFIX,
|
||||
EMPTY_MAPPING,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
|
|
@ -33,6 +34,8 @@ from litellm.constants import (
|
|||
from litellm.constants import (
|
||||
MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
|
||||
)
|
||||
from litellm.integrations.otel.model.metadata import caller_session_id
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls
|
||||
from litellm.litellm_core_utils.classifier_logging import (
|
||||
CLASSIFIER_AUDIT_FIELDS,
|
||||
classifier_audit_fields,
|
||||
|
|
@ -930,6 +933,16 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -
|
|||
return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX)
|
||||
|
||||
|
||||
def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None:
|
||||
beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None
|
||||
caller_session: Final = (
|
||||
beside_trace_header
|
||||
if isinstance(beside_trace_header, str) and beside_trace_header
|
||||
else caller_session_id(kwargs, TraceControls())
|
||||
)
|
||||
return caller_session if caller_session and "\x00" not in caller_session else None
|
||||
|
||||
|
||||
def _get_session_id_for_spend_log(
|
||||
kwargs: Mapping[str, object],
|
||||
metadata: Mapping[str, object] | None,
|
||||
|
|
@ -937,9 +950,12 @@ def _get_session_id_for_spend_log(
|
|||
omit_when_missing: bool,
|
||||
batch_trace_session_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may
|
||||
be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks
|
||||
the per-request trace ids because those differ between the create call and the cost poller's row."""
|
||||
"""Resolve the session id for the spend log row: `omit` honors only metadata.session_id,
|
||||
batch sessions outrank everything, then the conversation id the caller sent (also one an x-litellm-trace-id header
|
||||
displaced from the request's session), then trace ids.
|
||||
|
||||
A session minted by `missing_session_id: generate` or copied from the request's trace id is not a conversation, so
|
||||
it yields to the trace id, which the Responses bridge points at the previous turn's session to keep a chain whole."""
|
||||
if omit_when_missing:
|
||||
session_id: Final = metadata.get("session_id") if metadata else None
|
||||
return str(session_id) if session_id else None
|
||||
|
|
@ -948,11 +964,14 @@ def _get_session_id_for_spend_log(
|
|||
|
||||
if batch_trace_session_id is not None:
|
||||
return batch_trace_session_id
|
||||
if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
|
||||
return str(standard_logging_payload.get("trace_id"))
|
||||
if kwargs.get("litellm_trace_id") is not None:
|
||||
return str(kwargs.get("litellm_trace_id"))
|
||||
return str(uuid.uuid4())
|
||||
if litellm.request_correlation_in_logs and (caller_session := _caller_session_for_spend_log(kwargs, metadata)):
|
||||
return caller_session
|
||||
match standard_logging_payload:
|
||||
case {"trace_id": trace_id} if trace_id is not None:
|
||||
return str(trace_id)
|
||||
case _:
|
||||
litellm_trace_id: Final = kwargs.get("litellm_trace_id")
|
||||
return str(litellm_trace_id) if litellm_trace_id is not None else str(uuid.uuid4())
|
||||
|
||||
|
||||
def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None:
|
||||
|
|
|
|||
|
|
@ -281,7 +281,13 @@ class ResponsesSessionHandler:
|
|||
|
||||
SQL query
|
||||
|
||||
SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id
|
||||
SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then
|
||||
SELECT * FROM spend_logs with that session_id and the same owner
|
||||
|
||||
Only rows with the referenced response's owner are returned: the same user and team, and for a key with
|
||||
neither, the same key. The session id can be caller supplied, so matching on it alone would mix owners. A
|
||||
user's or team's regenerated or second key keeps the history; a key with neither does not, and turns written
|
||||
by another owner (an admin continuing the chain) are not replayed.
|
||||
|
||||
A just-finished turn gets a short second chance: the worker that served it may
|
||||
still be writing its spend log when the follow-up arrives, and an empty result
|
||||
|
|
@ -303,13 +309,24 @@ class ResponsesSessionHandler:
|
|||
|
||||
query: Final = """
|
||||
WITH matching_session AS (
|
||||
SELECT session_id
|
||||
SELECT session_id, api_key, "user", team_id
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = $1
|
||||
)
|
||||
SELECT *
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE session_id IN (SELECT session_id FROM matching_session)
|
||||
FROM "LiteLLM_SpendLogs" AS logs
|
||||
WHERE EXISTS (
|
||||
SELECT 1
|
||||
FROM matching_session
|
||||
WHERE matching_session.session_id = logs.session_id
|
||||
AND matching_session."user" IS NOT DISTINCT FROM logs."user"
|
||||
AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id
|
||||
AND (
|
||||
COALESCE(matching_session."user", '') <> ''
|
||||
OR COALESCE(matching_session.team_id, '') <> ''
|
||||
OR matching_session.api_key = logs.api_key
|
||||
)
|
||||
)
|
||||
ORDER BY "endTime" ASC;
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -39,7 +39,11 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching import InMemoryCache
|
||||
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
REDACTED_BY_LITELLM,
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER,
|
||||
)
|
||||
from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook
|
||||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
|
|
@ -563,6 +567,14 @@ class ChatCompletionSession(TypedDict, total=False):
|
|||
########### End of Initialize Classes used for Responses API ###########
|
||||
|
||||
|
||||
def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata"
|
||||
metadata: Final = request_kwargs.get(key)
|
||||
if not isinstance(metadata, dict):
|
||||
return MappingProxyType({})
|
||||
return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}}
|
||||
|
||||
|
||||
class LiteLLMCompletionResponsesConfig:
|
||||
@staticmethod
|
||||
def get_supported_openai_params(model: str) -> list:
|
||||
|
|
@ -1078,8 +1090,6 @@ class LiteLLMCompletionResponsesConfig:
|
|||
else:
|
||||
# Both are empty - this likely means function_call_output had empty/invalid call_id
|
||||
# Provide a helpful error message
|
||||
import litellm
|
||||
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
f"Unable to create messages for completion request. "
|
||||
|
|
@ -1094,8 +1104,14 @@ class LiteLLMCompletionResponsesConfig:
|
|||
)
|
||||
|
||||
litellm_completion_request["messages"] = combined_messages
|
||||
litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id")
|
||||
return litellm_completion_request
|
||||
chain_session: Final = chat_completion_session.get("litellm_session_id")
|
||||
litellm_completion_request["litellm_trace_id"] = chain_session
|
||||
chain_session_metadata: Final[Mapping[str, object]] = (
|
||||
_chain_session_metadata(chain_session, request_kwargs)
|
||||
if chain_session and litellm.request_correlation_in_logs
|
||||
else MappingProxyType({})
|
||||
)
|
||||
return {**litellm_completion_request, **chain_session_metadata}
|
||||
|
||||
@staticmethod
|
||||
def _transform_response_input_param_to_chat_completion_message(
|
||||
|
|
|
|||
783
tests/integration/spend/test_spend_log_caller_session.py
Normal file
783
tests/integration/spend/test_spend_log_caller_session.py
Normal file
|
|
@ -0,0 +1,783 @@
|
|||
import itertools
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from urllib.parse import urlsplit
|
||||
from collections import deque
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
list_value,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.budget import readiness_seconds
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240)
|
||||
|
||||
MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}")
|
||||
CALLER_SESSION_KEY: Final = "litellm_caller_session_id"
|
||||
PATHS: Final = MappingProxyType(
|
||||
{"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"}
|
||||
)
|
||||
ENDPOINT_STREAMS: Final = tuple(itertools.product(PATHS, (False, True)))
|
||||
Headers = Mapping[str, str] | Sequence[tuple[str, str]]
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return "callersess-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _token(label: str) -> str:
|
||||
return f"{label}-{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _frame(payload: Mapping[str, JsonValue]) -> bytes:
|
||||
return b"data: " + json.dumps(payload).encode() + b"\n\n"
|
||||
|
||||
|
||||
def _chat_reply(text: str, stream: bool) -> Reply:
|
||||
identity: Final = "chatcmpl-" + uuid.uuid4().hex
|
||||
usage: Final[dict[str, JsonValue]] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": usage,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final[dict[str, JsonValue]] = {
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}}]}),
|
||||
_frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _upstream(request: Request) -> Reply:
|
||||
markers: Final = tuple(found.group(0).decode() for found in MARKER.finditer(request.body))
|
||||
if request.method != "POST" or not request.target.endswith("/chat/completions") or not markers:
|
||||
return Reply(status=404, body=b'{"error":"unexpected request"}')
|
||||
return _chat_reply(f"answer {markers[-1]}", JSON_OBJECT.validate_json(request.body).get("stream") is True)
|
||||
|
||||
|
||||
def _user_markers(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
user_contents: Final = tuple(
|
||||
json.dumps(object_value(message)["content"])
|
||||
for message in list_value(body.get("messages", []))
|
||||
if object_value(message)["role"] == "user"
|
||||
)
|
||||
return tuple(found.group(0).decode() for content in user_contents if (found := MARKER.search(content.encode())))
|
||||
|
||||
|
||||
def _spans(batch: Request) -> Iterator[dict[str, JsonValue]]:
|
||||
for resource in list_value(JSON_OBJECT.validate_json(batch.body)["resourceSpans"]):
|
||||
for scope in list_value(object_value(resource)["scopeSpans"]):
|
||||
yield from (object_value(span) for span in list_value(object_value(scope)["spans"]))
|
||||
|
||||
|
||||
def _attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {
|
||||
string_value(object_value(attribute)["key"]): object_value(attribute)["value"]
|
||||
for attribute in list_value(span.get("attributes", []))
|
||||
}
|
||||
|
||||
|
||||
def _batch_attributes(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]:
|
||||
return (_attributes(span) for span in itertools.chain.from_iterable(_spans(batch) for batch in batches))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Collector:
|
||||
wire: Wire
|
||||
accepted: deque[Request]
|
||||
guard: threading.Lock
|
||||
|
||||
def span_attributes(self, response_id: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
with self.guard:
|
||||
batches: Final = tuple(self.accepted)
|
||||
attributes: Final = tuple(_batch_attributes(batches))
|
||||
return tuple(
|
||||
found
|
||||
for found in attributes
|
||||
if object_value(found.get("gen_ai.response.id", {})).get("stringValue") == response_id
|
||||
)
|
||||
|
||||
def conversation_ids(self, conversations: frozenset[str]) -> tuple[str, ...]:
|
||||
with self.guard:
|
||||
batches: Final = tuple(self.accepted)
|
||||
responded: Final = (
|
||||
object_value(found.get("gen_ai.conversation.id", {})).get("stringValue")
|
||||
for found in _batch_attributes(batches)
|
||||
if "gen_ai.response.id" in found
|
||||
)
|
||||
return tuple(
|
||||
conversation
|
||||
for conversation in responded
|
||||
if isinstance(conversation, str) and conversation in conversations
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
scenario: Scenario
|
||||
model: str
|
||||
upstream: Wire
|
||||
process: OwnedProxy
|
||||
|
||||
def body(self, endpoint: str, marker: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
match endpoint:
|
||||
case "chat":
|
||||
return {"model": self.model, "messages": [{"role": "user", "content": marker}], **extra}
|
||||
case "messages":
|
||||
return {
|
||||
"model": self.model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
**extra,
|
||||
}
|
||||
case _:
|
||||
return {"model": self.model, "input": marker, **extra}
|
||||
|
||||
def post(
|
||||
self,
|
||||
endpoint: str,
|
||||
marker: str,
|
||||
fields: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
*,
|
||||
stream: bool = False,
|
||||
headers: Headers = (),
|
||||
key: str | None = None,
|
||||
) -> httpx.Response:
|
||||
payload: Final = self.body(endpoint, marker, fields) | ({"stream": True} if stream else {})
|
||||
request_headers: Final = (
|
||||
("Authorization", f"Bearer {self.proxy.key if key is None else key}"),
|
||||
*(headers.items() if isinstance(headers, Mapping) else headers),
|
||||
)
|
||||
with self.proxy.client.stream("POST", PATHS[endpoint], json=payload, headers=request_headers) as response:
|
||||
response.read()
|
||||
return response
|
||||
|
||||
def send(
|
||||
self,
|
||||
endpoint: str,
|
||||
marker: str,
|
||||
fields: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
*,
|
||||
stream: bool = False,
|
||||
headers: Headers = (),
|
||||
key: str | None = None,
|
||||
) -> httpx.Response:
|
||||
response: Final = self.post(endpoint, marker, fields, stream=stream, headers=headers, key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
return response
|
||||
|
||||
def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
bodies: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in self.upstream.drain())
|
||||
return tuple(body for body in bodies if _user_markers(body)[-1:] == (marker,))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rigs:
|
||||
flag_off: Rig
|
||||
flag_on: Rig
|
||||
sink: Collector
|
||||
|
||||
|
||||
def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT session_id, request_id, metadata::text AS metadata_text FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE litellm_call_id=%s",
|
||||
(call_id,),
|
||||
)
|
||||
|
||||
|
||||
def _spend_row(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
call_id: Final = response.headers["x-litellm-call-id"]
|
||||
rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) >= 1, seconds=70)
|
||||
assert len(rows) == 1, rows
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _session(response: httpx.Response) -> JsonValue:
|
||||
return _spend_row(response)["session_id"]
|
||||
|
||||
|
||||
def _model(scenario: Scenario, upstream: Wire) -> str:
|
||||
return scenario.model(
|
||||
model="hosted_vllm/gpt-4o-mini",
|
||||
api_base=f"{upstream.url}/v1",
|
||||
use_chat_completions_api=True,
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
)
|
||||
|
||||
|
||||
def _rig_config(
|
||||
directory: Path, name: str, litellm_settings: Mapping[str, JsonValue], extra: Mapping[str, JsonValue]
|
||||
) -> Path:
|
||||
shared: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = {
|
||||
**shared,
|
||||
"general_settings": {**object_value(shared["general_settings"]), "store_prompts_in_spend_logs": True},
|
||||
"litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings},
|
||||
**extra,
|
||||
}
|
||||
path: Final = directory / f"{name}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rigs(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rigs]:
|
||||
accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted OTLP batch
|
||||
guard: Final = threading.Lock()
|
||||
|
||||
def sink(request: Request) -> Reply:
|
||||
with guard:
|
||||
accepted.append(request)
|
||||
return Reply()
|
||||
|
||||
directory: Final = tmp_path_factory.mktemp("caller-session")
|
||||
with ExitStack() as stack:
|
||||
upstream: Final = stack.enter_context(wire_server(_upstream))
|
||||
collector: Final = Collector(stack.enter_context(wire_server(sink)), accepted, guard)
|
||||
shared: Final = stack.enter_context(gateway_from_environment())
|
||||
flag_off: Final = stack.enter_context(
|
||||
owned_proxy_process(shared, directory, {}, config=_rig_config(directory, "flag-off", {}, {}), workers=2)
|
||||
)
|
||||
flag_on: Final = stack.enter_context(
|
||||
owned_proxy_process(
|
||||
shared,
|
||||
directory,
|
||||
{"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"},
|
||||
config=_rig_config(
|
||||
directory,
|
||||
"flag-on",
|
||||
{"request_correlation_in_logs": True, "callbacks": ["otel"]},
|
||||
{
|
||||
"callback_settings": {
|
||||
"otel": {"exporter": "http/json", "endpoint": collector.wire.url, "mapper_names": ["genai"]}
|
||||
}
|
||||
},
|
||||
),
|
||||
workers=2,
|
||||
)
|
||||
)
|
||||
off_scenario: Final = stack.enter_context(flag_off.gateway.scenario())
|
||||
on_scenario: Final = stack.enter_context(flag_on.gateway.scenario())
|
||||
yield Rigs(
|
||||
Rig(flag_off.gateway, off_scenario, _model(off_scenario, upstream), upstream, flag_off),
|
||||
Rig(flag_on.gateway, on_scenario, _model(on_scenario, upstream), upstream, flag_on),
|
||||
collector,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("endpoint", "stream"), ENDPOINT_STREAMS)
|
||||
def test_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str, stream: bool) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = _token("convo")
|
||||
response: Final = rigs.flag_on.send(
|
||||
endpoint,
|
||||
marker,
|
||||
stream=stream,
|
||||
headers={"x-litellm-trace-id": _token("trace")},
|
||||
fields={"litellm_session_id": session},
|
||||
)
|
||||
|
||||
row: Final = _spend_row(response)
|
||||
upstream: Final = rigs.flag_on.upstream_bodies(marker)
|
||||
|
||||
assert row["session_id"] == session, row
|
||||
assert CALLER_SESSION_KEY not in string_value(row["metadata_text"]), row
|
||||
assert len(upstream) == 1, upstream
|
||||
assert CALLER_SESSION_KEY not in json.dumps(upstream[0]), upstream
|
||||
|
||||
|
||||
def test_sdk_clients_record_the_body_session_beside_a_trace_header(rigs: Rigs) -> None:
|
||||
rig: Final = rigs.flag_on
|
||||
base_url: Final = str(rig.proxy.client.base_url)
|
||||
traces: Final = tuple(_token("trace") for _ in range(3))
|
||||
sessions: Final = tuple(_token("convo") for _ in range(3))
|
||||
markers: Final = tuple(_marker() for _ in range(3))
|
||||
chat_client: Final = openai.OpenAI(base_url=base_url + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=rig.proxy.key, max_retries=0)
|
||||
|
||||
with chat_client.chat.completions.with_streaming_response.create(
|
||||
model=rig.model,
|
||||
messages=[{"role": "user", "content": markers[0]}],
|
||||
stream=True,
|
||||
extra_headers={"x-litellm-trace-id": traces[0]},
|
||||
extra_body={"litellm_session_id": sessions[0]},
|
||||
) as chat_stream:
|
||||
chat_call: Final = chat_stream.headers["x-litellm-call-id"]
|
||||
chat_text: Final = "".join(
|
||||
chunk.choices[0].delta.content or "" for chunk in chat_stream.parse() if chunk.choices
|
||||
)
|
||||
message: Final = messages_client.messages.with_raw_response.create(
|
||||
model=rig.model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": markers[1]}],
|
||||
extra_headers={"x-litellm-trace-id": traces[1]},
|
||||
extra_body={"litellm_session_id": sessions[1]},
|
||||
)
|
||||
with chat_client.responses.with_streaming_response.create(
|
||||
model=rig.model,
|
||||
input=markers[2],
|
||||
stream=True,
|
||||
extra_headers={"x-litellm-trace-id": traces[2]},
|
||||
extra_body={"litellm_session_id": sessions[2]},
|
||||
) as responses_stream:
|
||||
responses_call: Final = responses_stream.headers["x-litellm-call-id"]
|
||||
completed: Final = tuple(event for event in responses_stream.parse() if event.type == "response.completed")
|
||||
|
||||
first_block: Final = message.parse().content[0]
|
||||
assert chat_text == f"answer {markers[0]}", chat_text
|
||||
assert isinstance(first_block, anthropic.types.TextBlock) and first_block.text == f"answer {markers[1]}", message
|
||||
assert len(completed) == 1, completed
|
||||
calls: Final = (chat_call, message.headers["x-litellm-call-id"], responses_call)
|
||||
landed: Final = tuple(
|
||||
eventually(lambda call=call: _spend_rows(call), lambda values: len(values) == 1, seconds=70)[0]["session_id"]
|
||||
for call in calls
|
||||
)
|
||||
assert landed == sessions, (landed, calls)
|
||||
|
||||
|
||||
def test_conversation_span_keeps_the_trace_header_and_no_internal_key(rigs: Rigs) -> None:
|
||||
trace: Final = _token("trace")
|
||||
response: Final = rigs.flag_on.send(
|
||||
"chat", _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")}
|
||||
)
|
||||
response_id: Final = string_value(JSON_OBJECT.validate_json(response.content)["id"])
|
||||
|
||||
spans: Final = eventually(lambda: rigs.sink.span_attributes(response_id), lambda found: len(found) == 1, seconds=30)
|
||||
|
||||
assert spans[0]["gen_ai.conversation.id"] == {"stringValue": trace}, spans
|
||||
assert CALLER_SESSION_KEY not in json.dumps(spans[0]), spans
|
||||
|
||||
|
||||
def test_a_continued_chains_conversation_spans_keep_each_turns_trace_header(rigs: Rigs) -> None:
|
||||
traces: Final = (_token("trace"), _token("trace"))
|
||||
first: Final = _turn(rigs.flag_on, _marker(), headers={"x-litellm-trace-id": traces[0]})
|
||||
_turn(
|
||||
rigs.flag_on,
|
||||
_marker(),
|
||||
first,
|
||||
headers={"x-litellm-trace-id": traces[1]},
|
||||
fields={"litellm_session_id": _token("convo")},
|
||||
)
|
||||
|
||||
conversations: Final = eventually(
|
||||
lambda: rigs.sink.conversation_ids(frozenset(traces)), lambda found: len(found) == 2, seconds=30
|
||||
)
|
||||
|
||||
assert sorted(conversations) == sorted(traces), conversations
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", list(PATHS))
|
||||
def test_flag_off_keeps_the_trace_id_beside_a_body_session(rigs: Rigs, endpoint: str) -> None:
|
||||
trace: Final = _token("trace")
|
||||
|
||||
response: Final = rigs.flag_off.send(
|
||||
endpoint, _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")}
|
||||
)
|
||||
|
||||
assert _session(response) == trace
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", list(PATHS))
|
||||
def test_response_cache_hit_records_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str) -> None:
|
||||
marker: Final = _marker()
|
||||
sessions: Final = (_token("convo"), _token("convo"))
|
||||
responses: Final = tuple(
|
||||
rigs.flag_on.send(
|
||||
endpoint, marker, headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session}
|
||||
)
|
||||
for session in sessions
|
||||
)
|
||||
|
||||
assert tuple(_session(response) for response in responses) == sessions
|
||||
assert len(rigs.flag_on.upstream_bodies(marker)) == 1
|
||||
|
||||
|
||||
def test_metadata_session_beside_a_new_trace_per_turn_groups_both_turns(rigs: Rigs) -> None:
|
||||
session: Final = _token("convo")
|
||||
responses: Final = tuple(
|
||||
rigs.flag_on.send(
|
||||
"chat",
|
||||
_marker(),
|
||||
headers={"x-litellm-trace-id": _token("trace")},
|
||||
fields={"metadata": {"session_id": session}},
|
||||
)
|
||||
for _ in range(2)
|
||||
)
|
||||
|
||||
assert tuple(_session(response) for response in responses) == (session, session)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
page: Final = rigs.flag_on.proxy.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"session_id": session,
|
||||
"start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
},
|
||||
)
|
||||
listed: Final = sorted(string_value(object_value(row)["request_id"]) for row in list_value(page["data"]))
|
||||
assert listed == sorted(
|
||||
string_value(JSON_OBJECT.validate_json(response.content)["id"]) for response in responses
|
||||
), page
|
||||
|
||||
|
||||
def _shape(shape: str, trace: str, session: str, header: str) -> tuple[Headers, dict[str, JsonValue], str]:
|
||||
traced: Final = {"x-litellm-trace-id": trace}
|
||||
match shape:
|
||||
case "session-header-wins":
|
||||
return {"x-litellm-session-id": header}, {"litellm_session_id": session}, header
|
||||
case "trace-header-only":
|
||||
return traced, {}, trace
|
||||
case "body-session-only":
|
||||
return {}, {"litellm_session_id": session}, session
|
||||
case "int-session":
|
||||
return traced, {"litellm_session_id": 7}, trace
|
||||
case "list-session":
|
||||
return traced, {"litellm_session_id": [session]}, trace
|
||||
case "empty-session":
|
||||
return traced, {"litellm_session_id": ""}, trace
|
||||
case "null-session":
|
||||
return traced, {"litellm_session_id": None}, trace
|
||||
case "nul-session":
|
||||
return traced, {"litellm_session_id": session + "\x00x"}, trace
|
||||
case "two-trace-headers":
|
||||
return (
|
||||
(("x-litellm-trace-id", trace), ("x-litellm-trace-id", header)),
|
||||
{"litellm_session_id": session},
|
||||
session,
|
||||
)
|
||||
case _:
|
||||
return traced, {"litellm_session_id": session + "s" * 5120}, session + "s" * 5120
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
"session-header-wins",
|
||||
"trace-header-only",
|
||||
"body-session-only",
|
||||
"int-session",
|
||||
"list-session",
|
||||
"empty-session",
|
||||
"null-session",
|
||||
"nul-session",
|
||||
"two-trace-headers",
|
||||
"5kb-session",
|
||||
],
|
||||
)
|
||||
def test_spend_log_session_for_each_session_shape(rigs: Rigs, shape: str) -> None:
|
||||
headers, fields, expected = _shape(shape, _token("trace"), _token("convo"), _token("header"))
|
||||
neighbor_trace: Final = _token("neighbor")
|
||||
|
||||
response: Final = rigs.flag_on.send("chat", _marker(), fields, headers=headers)
|
||||
neighbor: Final = rigs.flag_on.send("chat", _marker(), headers={"x-litellm-trace-id": neighbor_trace})
|
||||
|
||||
assert _session(response) == expected
|
||||
assert _session(neighbor) == neighbor_trace
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ["chat", "messages"])
|
||||
def test_client_sent_internal_session_key_is_ignored(rigs: Rigs, endpoint: str) -> None:
|
||||
forged: Final = _token("forged")
|
||||
trace: Final = _token("trace")
|
||||
forged_metadata: Final[dict[str, JsonValue]] = {
|
||||
"metadata" if endpoint == "chat" else "litellm_metadata": {CALLER_SESSION_KEY: forged}
|
||||
}
|
||||
|
||||
untraced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata)
|
||||
traced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata, headers={"x-litellm-trace-id": trace})
|
||||
|
||||
assert _session(untraced) not in (forged, None)
|
||||
assert _session(traced) == trace
|
||||
|
||||
|
||||
def test_identical_requests_each_log_once_with_the_body_session(rigs: Rigs) -> None:
|
||||
session: Final = _token("convo")
|
||||
trace: Final = _token("trace")
|
||||
marker: Final = _marker()
|
||||
|
||||
responses: Final = tuple(
|
||||
rigs.flag_on.send("chat", marker, headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": session})
|
||||
for _ in range(3)
|
||||
)
|
||||
|
||||
assert len({response.headers["x-litellm-call-id"] for response in responses}) == 3
|
||||
assert tuple(_session(response) for response in responses) == (session, session, session)
|
||||
|
||||
|
||||
def _pass_through_route(rig: Rig, target: str) -> str:
|
||||
path: Final = f"/caller-session-{uuid.uuid4().hex}"
|
||||
created: Final = rig.proxy.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": True})
|
||||
endpoint_id: Final = string_value(object_value(list_value(created["endpoints"])[0])["id"])
|
||||
rig.scenario.cleanups.callback(
|
||||
lambda: rig.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id})
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["litellm_session_id", "metadata"])
|
||||
def test_pass_through_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, field: str) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = _token("convo")
|
||||
rig: Final = rigs.flag_on
|
||||
path: Final = _pass_through_route(rig, f"{rig.upstream.url}/v1/chat/completions")
|
||||
key: Final = rig.scenario.key(allowed_passthrough_routes=[path])
|
||||
fields: Final[dict[str, JsonValue]] = (
|
||||
{"litellm_session_id": session} if field == "litellm_session_id" else {"metadata": {"session_id": session}}
|
||||
)
|
||||
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
path,
|
||||
rig.body("chat", marker, fields),
|
||||
key=key,
|
||||
headers={"x-litellm-trace-id": _token("trace")},
|
||||
)
|
||||
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT session_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)
|
||||
),
|
||||
lambda found: len(found) >= 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert f"answer {marker}" in response.text
|
||||
assert len(rig.upstream_bodies(marker)) == 1
|
||||
assert rows == [{"session_id": session}]
|
||||
|
||||
|
||||
def _turn(
|
||||
rig: Rig,
|
||||
marker: str,
|
||||
previous: str | None = None,
|
||||
*,
|
||||
key: str | None = None,
|
||||
headers: Headers = (),
|
||||
fields: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
) -> str:
|
||||
chained: Final[dict[str, JsonValue]] = {"previous_response_id": previous} if previous else {}
|
||||
response: Final = rig.send("responses", marker, {**chained, **fields}, headers=headers, key=key)
|
||||
_spend_row(response)
|
||||
return string_value(JSON_OBJECT.validate_json(response.content)["id"])
|
||||
|
||||
|
||||
def _turn_fields(shape: str, session: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]:
|
||||
match shape:
|
||||
case "session-every-turn":
|
||||
return {"litellm_session_id": session}, {"litellm_session_id": session}
|
||||
case "session-from-turn-2":
|
||||
return {}, {"litellm_session_id": session}
|
||||
case _:
|
||||
return {"litellm_session_id": _token("convo")}, {"litellm_session_id": _token("convo")}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("correlated", "shape", "replays_turn_1"),
|
||||
[
|
||||
(True, "session-every-turn", True),
|
||||
(True, "session-from-turn-2", True),
|
||||
(True, "new-session-per-turn", True),
|
||||
(False, "session-from-turn-2", False),
|
||||
],
|
||||
)
|
||||
def test_bridged_chain_history(rigs: Rigs, correlated: bool, shape: str, replays_turn_1: bool) -> None:
|
||||
rig: Final = rigs.flag_on if correlated else rigs.flag_off
|
||||
first_fields, second_fields = _turn_fields(shape, _token("convo"))
|
||||
markers: Final = (_marker(), _marker(), _marker())
|
||||
first: Final = _turn(rig, markers[0], headers={"x-litellm-trace-id": _token("trace")}, fields=first_fields)
|
||||
second: Final = _turn(rig, markers[1], first, headers={"x-litellm-trace-id": _token("trace")}, fields=second_fields)
|
||||
|
||||
_turn(rig, markers[2], second)
|
||||
|
||||
replayed: Final = rig.upstream_bodies(markers[2])
|
||||
assert len(replayed) == 1, replayed
|
||||
assert _user_markers(replayed[0]) == (markers if replays_turn_1 else markers[1:])
|
||||
|
||||
|
||||
def test_bridged_chains_sharing_a_metadata_session_replay_each_others_turns(rigs: Rigs) -> None:
|
||||
shared: Final[dict[str, JsonValue]] = {"metadata": {"session_id": _token("convo")}}
|
||||
markers: Final = (_marker(), _marker(), _marker())
|
||||
first: Final = _turn(rigs.flag_on, markers[0], fields=shared)
|
||||
_turn(rigs.flag_on, markers[1], fields=shared)
|
||||
|
||||
_turn(rigs.flag_on, markers[2], first)
|
||||
|
||||
replayed: Final = rigs.flag_on.upstream_bodies(markers[2])
|
||||
assert len(replayed) == 1, replayed
|
||||
assert _user_markers(replayed[0]) == markers
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Owners:
|
||||
first: str
|
||||
second: str
|
||||
|
||||
|
||||
def _owners(scenario: Scenario, model: str, shape: str) -> Owners:
|
||||
user: Final = _token("user")
|
||||
match shape:
|
||||
case "other-user":
|
||||
return Owners(
|
||||
scenario.key(user_id=user, models=[model]), scenario.key(user_id=_token("user"), models=[model])
|
||||
)
|
||||
case "same-user":
|
||||
return Owners(scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]))
|
||||
case "same-team":
|
||||
team: Final = scenario.team()
|
||||
return Owners(scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]))
|
||||
case _:
|
||||
return Owners(scenario.key(models=[model]), scenario.key(models=[model]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("shape", "replays_first_owner"),
|
||||
[("other-user", False), ("same-user", True), ("same-team", True), ("ownerless-keys", False)],
|
||||
)
|
||||
def test_bridged_history_follows_the_previous_responses_owner(
|
||||
rigs: Rigs, shape: str, replays_first_owner: bool
|
||||
) -> None:
|
||||
rig: Final = rigs.flag_off
|
||||
owners: Final = _owners(rig.scenario, rig.model, shape)
|
||||
session: Final = _token("convo")
|
||||
markers: Final = (_marker(), _marker(), _marker())
|
||||
_turn(rig, markers[0], key=owners.first, fields={"litellm_session_id": session})
|
||||
second: Final = _turn(rig, markers[1], key=owners.second, fields={"litellm_session_id": session})
|
||||
|
||||
_turn(rig, markers[2], second, key=owners.second)
|
||||
|
||||
replayed: Final = rig.upstream_bodies(markers[2])
|
||||
assert len(replayed) == 1, replayed
|
||||
assert _user_markers(replayed[0]) == (markers if replays_first_owner else markers[1:])
|
||||
|
||||
|
||||
BURST: Final = ENDPOINT_STREAMS * 4
|
||||
|
||||
|
||||
def test_concurrent_burst_across_endpoints_logs_each_session_once(rigs: Rigs) -> None:
|
||||
sessions: Final = tuple(_token("convo") for _ in BURST)
|
||||
|
||||
def one(index: int) -> httpx.Response:
|
||||
endpoint, stream = BURST[index]
|
||||
return rigs.flag_on.send(
|
||||
endpoint,
|
||||
_marker(),
|
||||
stream=stream,
|
||||
headers={"x-litellm-trace-id": _token("trace")},
|
||||
fields={"litellm_session_id": sessions[index]},
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=12) as pool:
|
||||
responses: Final = tuple(pool.map(one, range(len(BURST))))
|
||||
|
||||
assert tuple(_session(response) for response in responses) == sessions
|
||||
|
||||
|
||||
def _open_upstream_connections(pid: int, upstream: Wire) -> int:
|
||||
port: Final = urlsplit(upstream.url).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
||||
|
||||
def _post_or_none(rig: Rig, session: str) -> httpx.Response | None:
|
||||
try:
|
||||
return rig.post(
|
||||
"chat", _marker(), headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session}
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
return None
|
||||
|
||||
|
||||
def test_killing_one_worker_mid_burst_keeps_logging_the_body_session(rigs: Rigs) -> None:
|
||||
release: Final = threading.Event()
|
||||
arrived: Final[SimpleQueue[bytes]] = SimpleQueue()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
arrived.put(request.body)
|
||||
assert release.wait(timeout=60), "the burst was never released"
|
||||
return _upstream(request)
|
||||
|
||||
process: Final = rigs.flag_on.process
|
||||
sessions: Final = tuple(_token("convo") for _ in range(20))
|
||||
after_session: Final = _token("convo")
|
||||
with wire_server(held) as wire:
|
||||
rig: Final = replace(rigs.flag_on, model=_model(rigs.flag_on.scenario, wire), upstream=wire)
|
||||
workers: Final = eventually(process.worker_pids, lambda found: len(found) == 2, seconds=30)
|
||||
with ThreadPoolExecutor(max_workers=len(sessions)) as pool:
|
||||
pending: Final = tuple(pool.submit(_post_or_none, rig, session) for session in sessions)
|
||||
eventually(arrived.qsize, lambda size: size == len(sessions), seconds=60)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire) for pid in workers})
|
||||
victim, survivor = sorted(workers, key=held_by.__getitem__)
|
||||
with process.replacing_worker(victim) as kill:
|
||||
kill.freeze()
|
||||
kill()
|
||||
release.set()
|
||||
results: Final = tuple(future.result() for future in pending)
|
||||
after: Final = rig.send(
|
||||
"chat",
|
||||
_marker(),
|
||||
headers={"x-litellm-trace-id": _token("trace")},
|
||||
fields={"litellm_session_id": after_session},
|
||||
)
|
||||
|
||||
served: Final = tuple(
|
||||
(response, session)
|
||||
for response, session in zip(results, sessions, strict=True)
|
||||
if response is not None and response.status_code == 200
|
||||
)
|
||||
assert sum(held_by.values()) == len(sessions), held_by
|
||||
assert len(served) == held_by[survivor] >= len(sessions) // 2, (held_by, len(served))
|
||||
assert _session(after) == after_session
|
||||
assert tuple(_session(response) for response, _ in served) == tuple(session for _, session in served)
|
||||
|
|
@ -14,8 +14,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
import litellm
|
||||
import litellm.constants as litellm_constants
|
||||
import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils
|
||||
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
|
||||
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
|
@ -243,6 +243,82 @@ def test_batch_session_outranks_the_per_request_trace_id():
|
|||
assert session_id == "batch-uid-1"
|
||||
|
||||
|
||||
def _chat_session_id(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
litellm_params: Mapping[str, object],
|
||||
correlation_in_logs: bool = True,
|
||||
batch_trace_session_id: str | None = None,
|
||||
metadata: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> str | None:
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs)
|
||||
return _get_session_id_for_spend_log(
|
||||
kwargs={"litellm_trace_id": "trace-abc", "litellm_params": litellm_params},
|
||||
metadata=metadata,
|
||||
standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
|
||||
omit_when_missing=False,
|
||||
batch_trace_session_id=batch_trace_session_id,
|
||||
)
|
||||
|
||||
|
||||
_CALLER_SESSION_PARAMS: Final = MappingProxyType({"litellm_session_id": "sess-1", "metadata": {"trace_id": "otel-1"}})
|
||||
|
||||
|
||||
def test_caller_session_wins_over_the_per_request_trace_id(monkeypatch: pytest.MonkeyPatch):
|
||||
assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS) == "sess-1"
|
||||
|
||||
|
||||
def test_batch_session_outranks_a_caller_session(monkeypatch: pytest.MonkeyPatch):
|
||||
assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, batch_trace_session_id="batch-uid-1") == "batch-uid-1"
|
||||
|
||||
|
||||
def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(monkeypatch: pytest.MonkeyPatch):
|
||||
assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, correlation_in_logs=False) == "trace-abc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("correlation_in_logs", "expected_session_id"), [(True, "convo-1"), (False, "trace-abc")])
|
||||
def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging(
|
||||
monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str
|
||||
):
|
||||
header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}}
|
||||
session_id: Final = _chat_session_id(
|
||||
monkeypatch,
|
||||
header_chain,
|
||||
correlation_in_logs=correlation_in_logs,
|
||||
metadata={CALLER_SESSION_ID_METADATA_KEY: "convo-1"},
|
||||
)
|
||||
assert session_id == expected_session_id
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("litellm_params", "metadata"),
|
||||
[
|
||||
pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"),
|
||||
pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"),
|
||||
],
|
||||
)
|
||||
def test_caller_session_postgres_cannot_store_keeps_the_trace_id(
|
||||
monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], metadata: Mapping[str, object]
|
||||
):
|
||||
assert _chat_session_id(monkeypatch, litellm_params, metadata=metadata) == "trace-abc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params",
|
||||
[
|
||||
pytest.param({"metadata": {}}, id="no-session"),
|
||||
pytest.param(
|
||||
{"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}},
|
||||
id="minted-by-generate-policy",
|
||||
),
|
||||
pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"),
|
||||
],
|
||||
)
|
||||
def test_session_that_is_not_a_caller_conversation_keeps_the_trace_id(
|
||||
monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object]
|
||||
):
|
||||
assert _chat_session_id(monkeypatch, litellm_params) == "trace-abc"
|
||||
|
||||
|
||||
def test_omit_policy_still_suppresses_batch_sessions():
|
||||
session_id: Final = _get_session_id_for_spend_log(
|
||||
kwargs={},
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from starlette.datastructures import Headers
|
|||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
|
||||
SERVER_STREAMING_CLASSIFICATION_KEY,
|
||||
SERVER_STREAMING_CLASSIFICATION_MARKER,
|
||||
|
|
@ -3308,6 +3309,93 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain
|
|||
assert data["litellm_trace_id"] == "foo"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)])
|
||||
@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"])
|
||||
def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
correlation_in_logs: bool,
|
||||
expected_spend_log_session: str | None,
|
||||
session_in_metadata: bool,
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs)
|
||||
metadata: Final[dict[str, str]] = {"session_id": "convo-1"} if session_in_metadata else {}
|
||||
data: Final[dict[str, object]] = (
|
||||
{"metadata": metadata} if session_in_metadata else {"metadata": metadata, "litellm_session_id": "convo-1"}
|
||||
)
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session
|
||||
assert data["litellm_session_id"] == "req-trace-1"
|
||||
assert metadata["session_id"] == "req-trace-1"
|
||||
assert data["litellm_trace_id"] == "req-trace-1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"]
|
||||
)
|
||||
@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"])
|
||||
def test_session_headers_still_beat_the_body_session_with_correlation_logging(
|
||||
monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
|
||||
metadata: Final[dict[str, str]] = {}
|
||||
data: Final[dict[str, object]] = {"metadata": metadata, "litellm_session_id": "convo-1"}
|
||||
trace_header: Final = {"x-litellm-trace-id": "req-trace-1"} if with_trace_header else {}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={**trace_header, session_header: "hdr-session-1"}, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
spend_log_session: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) or data["litellm_session_id"]
|
||||
assert spend_log_session == "hdr-session-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("forged_in", ["metadata", "litellm_metadata"])
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[{"x-litellm-session-id": "hdr-session-1"}, {}],
|
||||
ids=["session-header", "no-headers"],
|
||||
)
|
||||
async def test_client_metadata_cannot_set_the_spend_log_caller_session(
|
||||
monkeypatch: pytest.MonkeyPatch, forged_in: str, headers: dict[str, str]
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-5.4-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
forged_in: {CALLER_SESSION_ID_METADATA_KEY: "forged-session"},
|
||||
},
|
||||
request=_make_request_mock("/v1/chat/completions", headers),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-5.4-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"litellm_session_id": "convo-1",
|
||||
"metadata": {"team_label": "kept"},
|
||||
},
|
||||
request=_make_request_mock("/v1/chat/completions", {"x-litellm-trace-id": "req-trace-1"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
assert updated["metadata"][CALLER_SESSION_ID_METADATA_KEY] == "convo-1"
|
||||
assert updated["metadata"]["requester_metadata"]["team_label"] == "kept"
|
||||
assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"]
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id():
|
||||
"""x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining."""
|
||||
headers = {"x-litellm-session-id": "bar"}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from openai.types.responses.response_function_web_search import (
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import CALLER_SESSION_ID_METADATA_KEY
|
||||
from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -6499,3 +6500,69 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No
|
|||
)
|
||||
|
||||
assert result["service_tier"] == "priority"
|
||||
|
||||
|
||||
def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
async def session_lookup(previous_response_id: str) -> ChatCompletionSession:
|
||||
assert previous_response_id == "resp_turn_1"
|
||||
return ChatCompletionSession(
|
||||
messages=[{"role": "user", "content": "turn 1"}], litellm_session_id="chain-session"
|
||||
)
|
||||
|
||||
return asyncio.run(
|
||||
LiteLLMCompletionResponsesConfig.async_responses_api_session_handler(
|
||||
previous_response_id="resp_turn_1",
|
||||
litellm_completion_request={
|
||||
"model": _CLAUDE_ON_BEDROCK,
|
||||
"custom_llm_provider": "bedrock",
|
||||
"messages": [{"role": "user", "content": "turn 2"}],
|
||||
"metadata": request_kwargs.get("metadata"),
|
||||
},
|
||||
instructions=None,
|
||||
session_lookup=session_lookup,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_kwargs", "logged_metadata_key"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"metadata": {"tag": "client"},
|
||||
"litellm_metadata": {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"},
|
||||
},
|
||||
"litellm_metadata",
|
||||
),
|
||||
({"metadata": {"user_api_key": "hashed"}}, "metadata"),
|
||||
],
|
||||
)
|
||||
def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_correlation(
|
||||
monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
|
||||
sent_kwargs: Final = deepcopy(request_kwargs)
|
||||
|
||||
request: Final = _continued_chain_request(request_kwargs)
|
||||
|
||||
assert request["litellm_trace_id"] == "chain-session"
|
||||
assert "litellm_session_id" not in request
|
||||
assert request[logged_metadata_key] == {
|
||||
**TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]),
|
||||
CALLER_SESSION_ID_METADATA_KEY: "chain-session",
|
||||
}
|
||||
assert request_kwargs == sent_kwargs
|
||||
|
||||
|
||||
def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", False)
|
||||
caller_metadata: Final = {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
|
||||
|
||||
request: Final = _continued_chain_request({"metadata": caller_metadata, "litellm_metadata": caller_metadata})
|
||||
|
||||
assert request["litellm_trace_id"] == "chain-session"
|
||||
assert "litellm_metadata" not in request
|
||||
assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
import sqlite3
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -787,3 +788,80 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None:
|
|||
tool_message: Final = result["messages"][2]
|
||||
assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments
|
||||
assert json.loads(tool_message["content"]) == function_output
|
||||
|
||||
|
||||
class _SqliteBackedPrismaDB:
|
||||
def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None:
|
||||
self._connection = sqlite3.connect(":memory:")
|
||||
self._connection.row_factory = sqlite3.Row
|
||||
self._connection.execute(
|
||||
'CREATE TABLE "LiteLLM_SpendLogs" '
|
||||
'(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT)'
|
||||
)
|
||||
self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows)
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]:
|
||||
return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id():
|
||||
"""A caller-supplied litellm_session_id can collide with another key's session; history must stay per key."""
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01", "", ""),
|
||||
("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02", "", ""),
|
||||
("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03", "", ""),
|
||||
("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04", "", ""),
|
||||
]
|
||||
)
|
||||
fake_prisma_client = _FakePrismaClient(results=[])
|
||||
fake_prisma_client.db = sqlite_db
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
|
||||
spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1")
|
||||
victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1")
|
||||
|
||||
assert [row["request_id"] for row in spend_logs] == ["attacker-1"]
|
||||
assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"]
|
||||
|
||||
|
||||
async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: str) -> list[str]:
|
||||
fake_prisma_client: Final = _FakePrismaClient(results=[])
|
||||
fake_prisma_client.db = sqlite_db
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
|
||||
spend_logs: Final = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(response_id)
|
||||
return [row["request_id"] for row in spend_logs]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key():
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("alice-1", "", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"),
|
||||
("bob-1", "", "shared-session", "2026-01-01T00:00:02", "bob", "team-a"),
|
||||
("service-a-1", "", "shared-session", "2026-01-01T00:00:03", "", "team-a"),
|
||||
("service-b-1", "", "shared-session", "2026-01-01T00:00:04", "", "team-b"),
|
||||
("alice-2", "", "shared-session", "2026-01-01T00:00:05", "alice", "team-a"),
|
||||
]
|
||||
)
|
||||
|
||||
assert await _history_request_ids(sqlite_db, "alice-1") == ["alice-1", "alice-2"]
|
||||
assert await _history_request_ids(sqlite_db, "bob-1") == ["bob-1"]
|
||||
assert await _history_request_ids(sqlite_db, "service-a-1") == ["service-a-1"]
|
||||
assert await _history_request_ids(sqlite_db, "service-b-1") == ["service-b-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_keys_apart():
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("alice-old-key", "hashed-old", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"),
|
||||
("ownerless-1", "hashed-other", "shared-session", "2026-01-01T00:00:02", "", ""),
|
||||
("alice-new-key", "hashed-new", "shared-session", "2026-01-01T00:00:03", "alice", "team-a"),
|
||||
("ownerless-2", "hashed-another", "shared-session", "2026-01-01T00:00:04", "", ""),
|
||||
]
|
||||
)
|
||||
|
||||
assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"]
|
||||
assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue