mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(responses): store each bridged turn's path to the root so a fork replays only its ancestors
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo
parent
e708343309
commit
61358572b7
10 changed files with 341 additions and 59 deletions
|
|
@ -1698,6 +1698,7 @@ 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"
|
||||
PREVIOUS_RESPONSE_IDS_METADATA_KEY: Final = "litellm_previous_response_ids"
|
||||
COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: Final = "_complexity_router_child_session_id"
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
|
||||
|
|
|
|||
|
|
@ -2429,7 +2429,9 @@ class ScheduledJobStaggerSettings(LiteLLMPydanticObjectBase):
|
|||
)
|
||||
|
||||
|
||||
SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset({"status", "cold_storage_object_key"})
|
||||
SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset(
|
||||
{"status", "cold_storage_object_key", "litellm_previous_response_ids"}
|
||||
)
|
||||
|
||||
|
||||
def _known_spend_logs_metadata_field(name: str) -> str:
|
||||
|
|
@ -3771,6 +3773,8 @@ class SpendLogsMetadata(TypedDict):
|
|||
requester_ip_address: str | None
|
||||
user_agent: ReadOnly[str | None]
|
||||
litellm_call_id: str | None
|
||||
litellm_previous_response_ids: ReadOnly[NotRequired[list[str] | None]]
|
||||
"""Spend log request ids of the bridged /v1/responses turns this turn continues, root first."""
|
||||
applied_guardrails: list[str] | None
|
||||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None
|
||||
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo
|
|||
from litellm.responses.litellm_completion_transformation.tool_search import has_client_tool_search
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
previous_response_ids_metadata,
|
||||
with_compaction_items_as_params,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
|
@ -225,19 +226,23 @@ class LiteLLMCompletionTransformationHandler:
|
|||
**kwargs,
|
||||
) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator:
|
||||
previous_response_id: Final[str | None] = responses_api_request.get("previous_response_id")
|
||||
if previous_response_id:
|
||||
litellm_completion_request = await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler(
|
||||
request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs)
|
||||
completion_request: Final = (
|
||||
await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler(
|
||||
previous_response_id=previous_response_id,
|
||||
litellm_completion_request=litellm_completion_request,
|
||||
instructions=responses_api_request.get("instructions"),
|
||||
request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs),
|
||||
request_kwargs=request_kwargs,
|
||||
previous_response_client_ids=previous_response_client_ids,
|
||||
stored_turns=stored_turns,
|
||||
input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)),
|
||||
)
|
||||
if previous_response_id
|
||||
else {**litellm_completion_request, **previous_response_ids_metadata((), request_kwargs)}
|
||||
)
|
||||
|
||||
litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion(
|
||||
**{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True},
|
||||
**{**kwargs, **completion_request, "_skip_responses_api_bridge": True},
|
||||
)
|
||||
if isinstance(litellm_completion_response, ModelResponse):
|
||||
responses_api_response: Final[ResponsesAPIResponse] = (
|
||||
|
|
@ -252,11 +257,11 @@ class LiteLLMCompletionTransformationHandler:
|
|||
|
||||
elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper):
|
||||
return LiteLLMCompletionStreamingIterator(
|
||||
model=litellm_completion_request.get("model") or "",
|
||||
model=completion_request.get("model") or "",
|
||||
litellm_custom_stream_wrapper=litellm_completion_response,
|
||||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
custom_llm_provider=litellm_completion_request.get("custom_llm_provider"),
|
||||
custom_llm_provider=completion_request.get("custom_llm_provider"),
|
||||
litellm_metadata=kwargs.get("litellm_metadata", {}),
|
||||
tool_call_provider_fields=not embeds_thought_signature_in_tool_call_id(
|
||||
_REQUEST_KWARGS_ADAPTER.validate_python(kwargs)
|
||||
|
|
|
|||
|
|
@ -5,9 +5,15 @@ from dataclasses import dataclass
|
|||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER
|
||||
from litellm.constants import (
|
||||
PREVIOUS_RESPONSE_IDS_METADATA_KEY,
|
||||
REDACTED_BY_LITELLM,
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER,
|
||||
)
|
||||
from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler
|
||||
from litellm.responses.litellm_completion_transformation.client_tool_search import (
|
||||
bridge_client_tool_search,
|
||||
|
|
@ -30,6 +36,7 @@ from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, Messag
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
ChatCompletionSession,
|
||||
SessionMessage,
|
||||
|
|
@ -136,6 +143,40 @@ def _referenced_request_ids(previous_response_id: str, client_response_ids: tupl
|
|||
return tuple(dict.fromkeys((response_id, previous_response_id, *client_response_ids)))
|
||||
|
||||
|
||||
def _stored_previous_response_ids(spend_log: Mapping[str, object]) -> tuple[str, ...] | None:
|
||||
"""The path to the root a turn stored on its row, None for a row written before paths were stored."""
|
||||
metadata: Final = spend_log.get("metadata")
|
||||
try:
|
||||
parsed: Final = (
|
||||
_STORED_METADATA_ADAPTER.validate_json(metadata)
|
||||
if isinstance(metadata, str)
|
||||
else _STORED_METADATA_ADAPTER.validate_python(metadata)
|
||||
)
|
||||
except ValidationError:
|
||||
return None
|
||||
try:
|
||||
return tuple(_REQUEST_IDS_ADAPTER.validate_python(parsed.get(PREVIOUS_RESPONSE_IDS_METADATA_KEY)))
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _path_through_referenced_rows(
|
||||
referenced_rows: Sequence[Mapping[str, object]], request_ids: Sequence[str]
|
||||
) -> tuple[str, ...] | None:
|
||||
"""The request ids to replay, root first, ending at the referenced response, or None without a stored path."""
|
||||
stored_path: Final = next(
|
||||
(path for path in map(_stored_previous_response_ids, referenced_rows) if path is not None), None
|
||||
)
|
||||
if stored_path is None:
|
||||
return None
|
||||
return tuple(dict.fromkeys((*stored_path, *request_ids)))
|
||||
|
||||
|
||||
def _in_path_order(spend_logs: Sequence["SpendLogsPayload"], path: Sequence[str]) -> list["SpendLogsPayload"]:
|
||||
position: Final = {request_id: index for index, request_id in enumerate(path)}
|
||||
return sorted(spend_logs, key=lambda spend_log: position.get(spend_log["request_id"], len(path)))
|
||||
|
||||
|
||||
def _turns_through_referenced_response(
|
||||
spend_logs: Sequence["SpendLogsPayload"], referenced_request_ids: tuple[str, ...]
|
||||
) -> int:
|
||||
|
|
@ -194,6 +235,55 @@ def _latest_replayable_instructions(proxy_server_requests: Sequence[Mapping[str,
|
|||
return None if latest == REDACTED_BY_LITELLM else latest
|
||||
|
||||
|
||||
_STORED_METADATA_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_REQUEST_IDS_ADAPTER: Final = TypeAdapter(list[str])
|
||||
|
||||
_REFERENCED_ROWS_QUERY: Final = """
|
||||
SELECT *
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = ANY($1::text[]);
|
||||
"""
|
||||
|
||||
_SAME_OWNER_AS_REFERENCED: Final = """
|
||||
EXISTS (
|
||||
SELECT 1
|
||||
FROM matching_session
|
||||
WHERE {same_conversation}
|
||||
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
|
||||
)
|
||||
)
|
||||
"""
|
||||
|
||||
_OWNER_SCOPED_QUERY: Final = """
|
||||
WITH matching_session AS (
|
||||
SELECT session_id, api_key, "user", team_id
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = ANY($1::text[])
|
||||
)
|
||||
SELECT *
|
||||
FROM "LiteLLM_SpendLogs" AS logs
|
||||
WHERE {same_owner}
|
||||
ORDER BY "endTime" ASC;
|
||||
"""
|
||||
|
||||
_SESSION_QUERY: Final = _OWNER_SCOPED_QUERY.format(
|
||||
same_owner=_SAME_OWNER_AS_REFERENCED.format(same_conversation="matching_session.session_id = logs.session_id")
|
||||
)
|
||||
|
||||
_PATH_QUERY: Final = _OWNER_SCOPED_QUERY.format(
|
||||
same_owner=_SAME_OWNER_AS_REFERENCED.format(same_conversation="logs.request_id = ANY($2::text[])")
|
||||
)
|
||||
|
||||
|
||||
async def _read_spend_logs(prisma_client: "PrismaClient", query: str, *params: object) -> list["SpendLogsPayload"]:
|
||||
return cast("list[SpendLogsPayload]", await prisma_client.db.query_raw(query, *params))
|
||||
|
||||
|
||||
class ResponsesSessionHandler:
|
||||
@staticmethod
|
||||
async def get_chat_completion_message_history_for_previous_response_id(
|
||||
|
|
@ -264,6 +354,7 @@ class ResponsesSessionHandler:
|
|||
messages=chat_completion_message_history,
|
||||
litellm_session_id=litellm_session_id,
|
||||
instructions=_latest_replayable_instructions(proxy_server_requests),
|
||||
previous_response_ids=tuple(spend_log["request_id"] for spend_log in all_spend_logs[:replayable_turns]),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -510,8 +601,11 @@ class ResponsesSessionHandler:
|
|||
"""
|
||||
Get all spend logs for a previous response id
|
||||
|
||||
|
||||
SQL query
|
||||
The referenced response's row is read first. A row written with a path to the root
|
||||
(``litellm_previous_response_ids`` in its metadata) names every earlier turn it continues, so
|
||||
only those rows plus the referenced one are read, which keeps a fork from replaying its
|
||||
siblings and the turns sent after the referenced response. A row written before paths were
|
||||
stored falls back to its whole session:
|
||||
|
||||
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
|
||||
|
|
@ -539,40 +633,25 @@ class ResponsesSessionHandler:
|
|||
if prisma_client is None:
|
||||
return []
|
||||
|
||||
query: Final = """
|
||||
WITH matching_session AS (
|
||||
SELECT session_id, api_key, "user", team_id
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = ANY($1::text[])
|
||||
)
|
||||
SELECT *
|
||||
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;
|
||||
"""
|
||||
|
||||
max_attempts: Final = 1 if disable_spend_logs else RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS
|
||||
for attempt in range(max_attempts):
|
||||
if attempt:
|
||||
await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL)
|
||||
if spend_logs := await prisma_client.db.query_raw(query, request_ids):
|
||||
verbose_proxy_logger.debug(
|
||||
"Found the following spend logs for previous response id %s: %s",
|
||||
response_id,
|
||||
json.dumps(spend_logs, indent=4, default=str),
|
||||
)
|
||||
return spend_logs
|
||||
referenced_rows = await _read_spend_logs(prisma_client, _REFERENCED_ROWS_QUERY, request_ids)
|
||||
if not referenced_rows:
|
||||
continue
|
||||
path = _path_through_referenced_rows(referenced_rows, request_ids)
|
||||
spend_logs = (
|
||||
_in_path_order(await _read_spend_logs(prisma_client, _PATH_QUERY, request_ids, list(path)), path)
|
||||
if path is not None
|
||||
else await _read_spend_logs(prisma_client, _SESSION_QUERY, request_ids)
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"Found the following spend logs for previous response id %s: %s",
|
||||
response_id,
|
||||
json.dumps(spend_logs, indent=4, default=str),
|
||||
)
|
||||
return spend_logs
|
||||
|
||||
verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id)
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.caching import InMemoryCache
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
PREVIOUS_RESPONSE_IDS_METADATA_KEY,
|
||||
REDACTED_BY_LITELLM,
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER,
|
||||
)
|
||||
|
|
@ -917,6 +918,9 @@ class ChatCompletionSession(TypedDict, total=False):
|
|||
messages: list[SessionMessage]
|
||||
litellm_session_id: str | None
|
||||
instructions: ReadOnly[str | None]
|
||||
previous_response_ids: ReadOnly[tuple[str, ...]]
|
||||
"""Spend log request ids of the replayed turns, root first, stored on the new turn so a follow-up or a fork
|
||||
rebuilds exactly this path instead of the whole session."""
|
||||
|
||||
|
||||
########### End of Initialize Classes used for Responses API ###########
|
||||
|
|
@ -924,10 +928,23 @@ class ChatCompletionSession(TypedDict, total=False):
|
|||
|
||||
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"
|
||||
return _with_metadata(key, {CALLER_SESSION_ID_METADATA_KEY: chain_session}, request_kwargs)
|
||||
|
||||
|
||||
def previous_response_ids_metadata(
|
||||
previous_response_ids: Sequence[str], request_kwargs: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
"""The proxy's ``litellm_metadata`` carrying the path to the root, so the spend log row of this turn stores it."""
|
||||
return _with_metadata(
|
||||
"litellm_metadata", {PREVIOUS_RESPONSE_IDS_METADATA_KEY: list(previous_response_ids)}, request_kwargs
|
||||
)
|
||||
|
||||
|
||||
def _with_metadata(key: str, extra: Mapping[str, object], request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
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}}
|
||||
return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), **extra}}
|
||||
|
||||
|
||||
class LiteLLMCompletionResponsesConfig:
|
||||
|
|
@ -1562,7 +1579,13 @@ class LiteLLMCompletionResponsesConfig:
|
|||
if chain_session and litellm.request_correlation_in_logs
|
||||
else MappingProxyType({})
|
||||
)
|
||||
return {**litellm_completion_request, **chain_session_metadata}
|
||||
replayed_path: Final = chat_completion_session.get("previous_response_ids")
|
||||
path_metadata: Final[Mapping[str, object]] = (
|
||||
previous_response_ids_metadata(replayed_path, {**request_kwargs, **chain_session_metadata})
|
||||
if replayed_path is not None
|
||||
else MappingProxyType({})
|
||||
)
|
||||
return {**litellm_completion_request, **chain_session_metadata, **path_metadata}
|
||||
|
||||
@staticmethod
|
||||
def _transform_response_input_param_to_chat_completion_message(
|
||||
|
|
|
|||
|
|
@ -1595,7 +1595,7 @@ async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search
|
|||
[("tool_use", "toolu_standup")],
|
||||
[("tool_result", "toolu_standup"), ("text", None)],
|
||||
]
|
||||
assert prisma_client.db.queries == 1
|
||||
assert prisma_client.db.queries == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import httpx
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import PREVIOUS_RESPONSE_IDS_METADATA_KEY
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.responses.litellm_completion_transformation.handler import (
|
||||
LiteLLMCompletionTransformationHandler,
|
||||
|
|
@ -56,6 +57,28 @@ def test_sync_fallback_tags_skip_responses_api_bridge():
|
|||
assert captured.get("_skip_responses_api_bridge") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_first_turn_stores_an_empty_path_to_the_root_so_a_fork_off_it_replays_only_it():
|
||||
handler = LiteLLMCompletionTransformationHandler()
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_acompletion(**kwargs):
|
||||
captured.update(kwargs)
|
||||
raise _StopForwarding()
|
||||
|
||||
with patch("litellm.acompletion", fake_acompletion), pytest.raises(_StopForwarding):
|
||||
await handler.response_api_handler(
|
||||
model="gpt-4o",
|
||||
input="hello",
|
||||
responses_api_request={},
|
||||
custom_llm_provider="openai",
|
||||
_is_async=True,
|
||||
litellm_metadata={"user_api_key": "hashed"},
|
||||
)
|
||||
|
||||
assert captured["litellm_metadata"] == {"user_api_key": "hashed", PREVIOUS_RESPONSE_IDS_METADATA_KEY: []}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fallback_tags_skip_responses_api_bridge():
|
||||
handler = LiteLLMCompletionTransformationHandler()
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from pydantic import TypeAdapter, ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field
|
||||
from litellm.constants import CALLER_SESSION_ID_METADATA_KEY
|
||||
from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, PREVIOUS_RESPONSE_IDS_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 (
|
||||
|
|
@ -7109,6 +7109,42 @@ def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_corre
|
|||
assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_correlation_in_logs", [True, False])
|
||||
def test_a_continued_chain_stores_the_replayed_turns_as_its_path_to_the_root(
|
||||
monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool
|
||||
):
|
||||
"""The proxy's spend log row for this turn carries the ids it replayed, so a fork off it replays only them."""
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs)
|
||||
|
||||
async def session_lookup(previous_response_id: str) -> ChatCompletionSession:
|
||||
return ChatCompletionSession(
|
||||
messages=[{"role": "user", "content": "turn 1"}],
|
||||
litellm_session_id="chain-session",
|
||||
previous_response_ids=("chatcmpl-root", "chatcmpl-turn-1"),
|
||||
)
|
||||
|
||||
request: Final = 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"}],
|
||||
},
|
||||
instructions=None,
|
||||
session_lookup=session_lookup,
|
||||
request_kwargs={"litellm_metadata": {"user_api_key": "hashed"}, "metadata": {"tag": "client"}},
|
||||
)
|
||||
)
|
||||
|
||||
assert request["litellm_metadata"] == {
|
||||
"user_api_key": "hashed",
|
||||
**({CALLER_SESSION_ID_METADATA_KEY: "chain-session"} if request_correlation_in_logs else {}),
|
||||
PREVIOUS_RESPONSE_IDS_METADATA_KEY: ["chatcmpl-root", "chatcmpl-turn-1"],
|
||||
}
|
||||
assert "metadata" not in request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chunks():
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
import json
|
||||
import sqlite3
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import chain
|
||||
from types import SimpleNamespace
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -13,6 +13,7 @@ from fastapi import HTTPException
|
|||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm.constants import PREVIOUS_RESPONSE_IDS_METADATA_KEY
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
_get_proxy_server_request_for_spend_logs_payload,
|
||||
)
|
||||
|
|
@ -445,15 +446,21 @@ def _chat_completion_response(request_id: str, content: str) -> dict:
|
|||
|
||||
|
||||
class _FakePrismaDB:
|
||||
"""Serves the next queued result per history read; the referenced-row read peeks at it, dropping a miss."""
|
||||
|
||||
def __init__(self, results):
|
||||
self._results = list(results)
|
||||
self.calls = []
|
||||
|
||||
async def query_raw(self, query, request_ids):
|
||||
async def query_raw(self, query, request_ids, *path):
|
||||
self.calls.append(tuple(request_ids))
|
||||
if not self._results:
|
||||
return []
|
||||
return list(self._results.pop(0))
|
||||
if "matching_session" in query:
|
||||
return list(self._results.pop(0)) if self._results else []
|
||||
queued = self._results[0] if self._results else []
|
||||
referenced = [row for row in queued if row.get("request_id") in request_ids]
|
||||
if not referenced and self._results:
|
||||
self._results.pop(0)
|
||||
return referenced
|
||||
|
||||
|
||||
class _FakePrismaClient:
|
||||
|
|
@ -539,7 +546,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus
|
|||
("assistant", "OK"),
|
||||
]
|
||||
assert result["litellm_session_id"] == session_id
|
||||
assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id]
|
||||
assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id, request_id]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -630,8 +637,8 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id():
|
|||
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
|
||||
await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id)
|
||||
|
||||
assert len(fake_prisma_client.db.calls) == 1
|
||||
assert request_id in fake_prisma_client.db.calls[0]
|
||||
assert len(fake_prisma_client.db.calls) == 2
|
||||
assert all(request_id in call for call in fake_prisma_client.db.calls)
|
||||
|
||||
|
||||
class _StoredSpendLogs:
|
||||
|
|
@ -1169,18 +1176,27 @@ async def test_a_compacted_session_replays_no_turn_sent_after_the_referenced_res
|
|||
|
||||
|
||||
class _SqliteBackedPrismaDB:
|
||||
def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
rows: list[tuple[str, str, str, str, str | None, str | None]],
|
||||
metadata: Mapping[str, Mapping[str, object]] = MappingProxyType({}),
|
||||
) -> 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)'
|
||||
'(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT, metadata TEXT)'
|
||||
)
|
||||
self._connection.executemany(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?, ?)',
|
||||
[(*row, json.dumps(metadata[row[0]]) if row[0] in metadata else None) for row in rows],
|
||||
)
|
||||
self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows)
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]:
|
||||
sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace(
|
||||
"$1", "?"
|
||||
sqlite_query: Final = (
|
||||
query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?1))")
|
||||
.replace("= ANY($2::text[])", "IN (SELECT value FROM json_each(?2))")
|
||||
.replace("IS NOT DISTINCT FROM", "IS")
|
||||
)
|
||||
sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args)
|
||||
return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)]
|
||||
|
|
@ -1216,6 +1232,101 @@ async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: st
|
|||
return [row["request_id"] for row in spend_logs]
|
||||
|
||||
|
||||
def _stored_path(*previous_response_ids: str) -> Mapping[str, object]:
|
||||
return {PREVIOUS_RESPONSE_IDS_METADATA_KEY: list(previous_response_ids)}
|
||||
|
||||
|
||||
def _forked_session_db() -> _SqliteBackedPrismaDB:
|
||||
"""One session: root -> a2 -> a3 on branch a, root -> b2 -> b3 on branch b, b2 sent between a2 and a3."""
|
||||
return _SqliteBackedPrismaDB(
|
||||
[
|
||||
("root", "key", "session", "2026-01-01T00:00:01", "", ""),
|
||||
("a2", "key", "session", "2026-01-01T00:00:02", "", ""),
|
||||
("b2", "key", "session", "2026-01-01T00:00:03", "", ""),
|
||||
("a3", "key", "session", "2026-01-01T00:00:04", "", ""),
|
||||
("b3", "key", "session", "2026-01-01T00:00:05", "", ""),
|
||||
],
|
||||
metadata={
|
||||
"root": _stored_path(),
|
||||
"a2": _stored_path("root"),
|
||||
"b2": _stored_path("root"),
|
||||
"a3": _stored_path("root", "a2"),
|
||||
"b3": _stored_path("root", "b2"),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("previous_response_id", "replayed"),
|
||||
[
|
||||
("root", ["root"]),
|
||||
("a2", ["root", "a2"]),
|
||||
("b2", ["root", "b2"]),
|
||||
("a3", ["root", "a2", "a3"]),
|
||||
("b3", ["root", "b2", "b3"]),
|
||||
],
|
||||
)
|
||||
async def test_a_turn_with_a_stored_path_replays_only_the_path_to_the_root(
|
||||
previous_response_id: str, replayed: list[str]
|
||||
):
|
||||
"""A fork shares its session with the other branch, so the stored path, not the session, picks the turns."""
|
||||
assert await _history_request_ids(_forked_session_db(), previous_response_id) == replayed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stored_path_is_replayed_in_path_order_whatever_the_clock_says():
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("late-root", "key", "session", "2026-01-01T00:00:09", "", ""),
|
||||
("child", "key", "session", "2026-01-01T00:00:01", "", ""),
|
||||
],
|
||||
metadata={"late-root": _stored_path(), "child": _stored_path("late-root")},
|
||||
)
|
||||
|
||||
assert await _history_request_ids(sqlite_db, "child") == ["late-root", "child"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_turn_stored_before_paths_existed_still_replays_its_whole_session():
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("old-1", "key", "session", "2026-01-01T00:00:01", "", ""),
|
||||
("old-2", "key", "session", "2026-01-01T00:00:02", "", ""),
|
||||
("old-3", "key", "session", "2026-01-01T00:00:03", "", ""),
|
||||
]
|
||||
)
|
||||
|
||||
assert await _history_request_ids(sqlite_db, "old-2") == ["old-1", "old-2", "old-3"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stored_path_never_replays_another_owners_turns():
|
||||
"""A path naming another key's row (a forged or reused id) replays nothing of that key's history."""
|
||||
sqlite_db: Final = _SqliteBackedPrismaDB(
|
||||
[
|
||||
("victim-1", "victim-key", "session", "2026-01-01T00:00:01", "", ""),
|
||||
("attacker-1", "attacker-key", "session", "2026-01-01T00:00:02", "", ""),
|
||||
],
|
||||
metadata={"victim-1": _stored_path(), "attacker-1": _stored_path("victim-1")},
|
||||
)
|
||||
|
||||
assert await _history_request_ids(sqlite_db, "attacker-1") == ["attacker-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_replayed_turns_become_the_next_turns_path_to_the_root():
|
||||
fake_prisma_client: Final = _FakePrismaClient(results=[])
|
||||
fake_prisma_client.db = _forked_session_db()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
|
||||
session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
|
||||
"b2"
|
||||
)
|
||||
|
||||
assert session["previous_response_ids"] == ("root", "b2")
|
||||
|
||||
|
||||
@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(
|
||||
|
|
|
|||
|
|
@ -2366,7 +2366,7 @@ async def test_bridged_auto_execute_follow_up_continues_the_client_conversation_
|
|||
await _settle(lambda: len(recorder.logged) > 0)
|
||||
|
||||
follow_up: Final = json.loads(provider.calls[1].request.content)["messages"]
|
||||
assert spend_logs.asked == (frozenset({"chatcmpl-prior-turn"}),) * 2
|
||||
assert spend_logs.asked and set(spend_logs.asked) == {frozenset({"chatcmpl-prior-turn"})}
|
||||
assert [message["role"] for message in follow_up] == ["user", "assistant", "user", "assistant", "user"]
|
||||
assert json.dumps(follow_up).count("My favorite color is teal.") == 1
|
||||
assert [block["type"] for block in follow_up[3]["content"]] == ["thinking", "tool_use"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue