mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): flush the spend-log queue only for Responses turns the bridge reads back (#47614)
Every /v1/responses spend log asked its worker to flush the queue right away so a
chained previous_response_id could find the row. Only the chat-completions bridge
reads that row. Native Responses providers hold the conversation upstream, so their
rows now wait for the batch write like chat rows do
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Backport of #47614 to stable/1.105.x.
Cherry-picked from 423362bab2 (main).
This commit is contained in:
parent
7ffaaab9a0
commit
4ef6f611e0
3 changed files with 100 additions and 16 deletions
|
|
@ -89,6 +89,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
|
||||
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
else:
|
||||
PrismaClient = Any
|
||||
ProxyLogging = Any
|
||||
|
|
@ -98,6 +99,34 @@ RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, Call
|
|||
_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _turn_chains_through_the_db(payload: Mapping[str, object], llm_router: "Router | None") -> bool:
|
||||
"""A Responses turn the chat-completions bridge replays from its spend log on the next `previous_response_id`.
|
||||
|
||||
Native Responses providers hold the conversation themselves, so their rows can wait for the batch write.
|
||||
A row with no router deployment behind it is treated as bridged, which keeps today's immediate flush.
|
||||
"""
|
||||
if payload.get("call_type") not in RESPONSES_SESSION_CALL_TYPES:
|
||||
return False
|
||||
model_id: Final = payload.get("model_id")
|
||||
deployment: Final = (
|
||||
llm_router.get_deployment(model_id=model_id) if llm_router is not None and isinstance(model_id, str) else None
|
||||
)
|
||||
if deployment is None:
|
||||
return True
|
||||
from litellm.responses.main import provider_stores_the_conversation
|
||||
|
||||
litellm_params: Final = deployment.litellm_params
|
||||
return not provider_stores_the_conversation(
|
||||
litellm_params.model,
|
||||
litellm_params.custom_llm_provider,
|
||||
{
|
||||
"use_chat_completions_api": litellm_params.use_chat_completions_api,
|
||||
"model_info": deployment.model_info.model_dump(),
|
||||
"api_base": litellm_params.api_base,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
|
||||
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
|
||||
|
||||
|
|
@ -526,12 +555,13 @@ class DBSpendUpdateWriter:
|
|||
get_request_model_access_groups,
|
||||
)
|
||||
|
||||
llm_router: Final = get_llm_router()
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
llm_router=get_llm_router(),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
payload["spend"] = response_cost or 0.0
|
||||
if isinstance(payload["startTime"], datetime):
|
||||
|
|
@ -546,7 +576,10 @@ class DBSpendUpdateWriter:
|
|||
payload["team_id"] = team_id
|
||||
|
||||
if not await self._record_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
payload=payload,
|
||||
prisma_client=prisma_client,
|
||||
disable_spend_logs=disable_spend_logs,
|
||||
llm_router=llm_router,
|
||||
):
|
||||
return False
|
||||
|
||||
|
|
@ -607,14 +640,18 @@ class DBSpendUpdateWriter:
|
|||
return True
|
||||
|
||||
async def _record_spend_log(
|
||||
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
|
||||
self,
|
||||
payload: SpendLogsPayload,
|
||||
prisma_client: "PrismaClient | None",
|
||||
disable_spend_logs: bool,
|
||||
llm_router: "Router | None" = None,
|
||||
) -> bool:
|
||||
if prisma_client is not None and _is_batch_cost_row(payload):
|
||||
return await self._claim_batch_cost_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
)
|
||||
if disable_spend_logs is False:
|
||||
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
|
||||
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router)
|
||||
return True
|
||||
|
||||
async def _claim_batch_cost_spend_log(
|
||||
|
|
@ -1475,6 +1512,7 @@ class DBSpendUpdateWriter:
|
|||
payload: dict | SpendLogsPayload,
|
||||
prisma_client: PrismaClient | None = None,
|
||||
spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"),
|
||||
llm_router: "Router | None" = None,
|
||||
) -> PrismaClient | None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Writing spend log to db - request_id: {}, spend: {}".format(
|
||||
|
|
@ -1485,7 +1523,7 @@ class DBSpendUpdateWriter:
|
|||
from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush
|
||||
|
||||
await enqueue_spend_logs(prisma_client, (payload,))
|
||||
if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES:
|
||||
if _turn_chains_through_the_db(_SPEND_METADATA_ADAPTER.validate_python(payload), llm_router):
|
||||
request_spend_log_flush(prisma_client)
|
||||
else:
|
||||
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
|
||||
|
|
|
|||
|
|
@ -547,6 +547,22 @@ def _will_bridge_to_chat_completions(
|
|||
)
|
||||
|
||||
|
||||
def provider_stores_the_conversation(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
|
||||
"""Whether the provider, not a LiteLLM spend log, holds this turn for a `previous_response_id` follow-up.
|
||||
|
||||
A bridged turn is replayed from its spend log, so the bridge needs that row in the DB, and a bridged
|
||||
first MCP round runs as an internal call that writes no spend log at all.
|
||||
"""
|
||||
api_base: Final = _api_base_kwarg(kwargs)
|
||||
return not _will_bridge_to_chat_completions(
|
||||
model,
|
||||
custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1],
|
||||
bool(kwargs.get("use_chat_completions_api")),
|
||||
kwargs.get("model_info"),
|
||||
api_base,
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs: dict[str, object], # mutable-ok: the signal is read and popped out of the caller's own kwargs
|
||||
|
|
|
|||
|
|
@ -3882,28 +3882,58 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
|
|||
proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
def _responses_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "native",
|
||||
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test"},
|
||||
"model_info": {"id": "dep-native"},
|
||||
},
|
||||
{
|
||||
"model_name": "bridged",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"},
|
||||
"model_info": {"id": "dep-bridged"},
|
||||
},
|
||||
{
|
||||
"model_name": "forced-bridge",
|
||||
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test", "use_chat_completions_api": True},
|
||||
"model_info": {"id": "dep-forced-bridge"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type, expects_flush",
|
||||
[("aresponses", True), ("responses", True), ("acompletion", False)],
|
||||
"call_type, model_id, expects_flush",
|
||||
[
|
||||
("aresponses", "dep-bridged", True),
|
||||
("responses", "dep-bridged", True),
|
||||
("aresponses", "dep-forced-bridge", True),
|
||||
("aresponses", None, True),
|
||||
("aresponses", "dep-native", False),
|
||||
("responses", "dep-native", False),
|
||||
("acompletion", "dep-bridged", False),
|
||||
],
|
||||
)
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back(
|
||||
call_type: str, expects_flush: bool
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bridge_reads_back(
|
||||
call_type: str, model_id: str | None, expects_flush: bool
|
||||
):
|
||||
"""
|
||||
A `previous_response_id` chained straight off the previous turn reads the DB, so a
|
||||
Responses row cannot sit in this worker's queue until the monitor's next poll.
|
||||
A `previous_response_id` chained straight off a bridged turn reads that turn's row from the
|
||||
DB, so it cannot sit in this worker's queue until the monitor's next poll. A native Responses
|
||||
provider holds the conversation itself, so its row waits for the batch write like a chat row.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
prisma = _tool_usage_prisma()
|
||||
prisma.spend_log_flush_requested = asyncio.Event()
|
||||
payload = {"request_id": "req-1", "call_type": call_type, "model_id": model_id}
|
||||
|
||||
await db_writer._insert_spend_log_to_db(
|
||||
payload={"request_id": "req-1", "call_type": call_type},
|
||||
prisma_client=prisma,
|
||||
)
|
||||
await db_writer._insert_spend_log_to_db(payload=payload, prisma_client=prisma, llm_router=_responses_router())
|
||||
|
||||
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
|
||||
assert prisma.spend_log_transactions == [payload]
|
||||
assert prisma.spend_log_flush_requested.is_set() is expects_flush
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue