mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(spend)!: save streamed /v1/responses spend rows under the encrypted resp_ id the client received (#49493)
* fix(responses): log the encrypted response id on streamed /v1/responses spend rows
With response id security on, the stream hook encrypts the resp_ id on every
frame, but logging had already received the SDK's own copy of the terminal
response, so the spend row, the standard logging payload and callbacks got the
managed id (or the chatcmpl id on a bridged stream) the client never saw. The
hook now records the exact id it advertises on the request's logging object,
and the logging layer stamps that id on the logged streamed response
* fix(responses): return the id the client sent on GET and cancel, and keep read rows off the create row's id
Every GET /v1/responses/{id} re-encrypted the response id, so a client that continued a bridged conversation with the
id a GET returned referenced a ciphertext no spend row was filed under. The follow-up lost turn 1 and opened a new
session. GET and cancel now answer with the exact id the client sent when it is an id this proxy issued and it decodes
to the response the provider returned. A create with previous_response_id still gets a fresh id for its own response,
and a raw provider id still gets a managed id
Because a read now carries the create's id, a GET or cancel spend row is filed as <id>_retrieve_<call id> or
<id>_cancel_<call id>. That keeps one row per read, as before, and a read row can never take the create row's
request_id, so the create row survives whichever row is written first and a follow-up replays turn 1 once. The
background settlement row keeps <id>_background_cost, and the cost poll's own status checks keep the response id
* fix(responses): log a fresh id for each read and cancel, echo the held id only to its owner
A GET or cancel of a stored response now logs a new encrypted id of its own,
so callbacks and spend rows never reuse the create's id, while the owner of
the id still gets back the exact id they sent. An admin reading someone
else's response gets an id bound to the admin, as before. The spend row id
suffix for reads and cancels is gone, so with response id security off the
rows and session history match the default branch. A streamed GET follows
the same rules.
* fix(responses): echo the sent id to every reader that passes the response access check
A GET, streamed GET or cancel of a response now answers with the id the caller sent whenever the caller passed
the access check, not only when the caller is the user or team the id was issued to. An admin, the master key
or a teammate who read someone else's streamed response got a fresh id back, and the session lookup could not
reach the first turn from it, so a bridged follow-up through that id lost the conversation history
Every read and cancel still logs a fresh id bound to the reader, never the create's id
* fix(responses): keep a read or cancel row in session history only when it brings new output
A GET, streamed GET or cancel sent with a session header joins the
create's session, and its stored response repeats the create's output.
The replay then sent the same assistant turn once per read.
Session history now keeps an aget_responses or acancel_responses row
only when its stored Responses output has an item id that no earlier
kept row carried. A polled background response keeps the first poll
that returns the finished output, since its queued create row has none,
and later polls of the same output are dropped. Every other row keeps
today's rule
* fix(responses): replay each history item once, echo the sent id on delete, keep background cost files apart
A create row now owns its output items in session history, and a read or cancel row adds only items no
create and no earlier read carried, in one linear pass, so a background response polled while it grows,
or a read that ends before its streamed create, replays the answer once and after its prompt
DELETE /v1/responses/{id} answers with the id the caller sent, on a copy, and logs a fresh encrypted id
like a read does
Azure Data Lake Storage per-event files and unbatched GCS objects are named after the event's spend row
id, so a streamed background create and its cost event, which now share the client's id, keep separate
files
The integration tests that decoded the managed id now match the exact id the client received
* fix(responses): keep main's answer for a raw-id delete and main's storage names outside background cost events
A DELETE by a raw provider id returns the provider's id as main did, and logging still records a fresh
encrypted id. A streamed GET matches the held id by upstream response id, like a plain GET
Only a background response's cost event gets the _background_cost suffix on its Azure Data Lake Storage
file and unbatched GCS object, so batch retrieves and events with no id keep main's names
An item a session read replays takes the newest read's copy at the first read's place, so an item a
poll caught half written replays finished
Tests pin the client's id over the router's continued response id, and the storage test no longer reads
the clock after the upload
* fix(spend): give background cost events their own OpenMeter and DynamoDB ids, replay id-less read items
A streamed background create and its later cost event share the client's
id, so OpenMeter dropped the second event and DynamoDB overwrote the
create's item. Both now key the cost event as <id>_background_cost, like
its spend row and storage file.
Session history again replays a read's output items that carry no id,
as main did.
* fix(integrations): keep OpenMeter and DynamoDB logging working without the proxy extras
* fix(integrations): give background cost events their own ClickHouse request id and DeepEval span id
A streamed background create and its later cost event now share the id the
client received, so ClickHouse and DeepEval saw both events under one id.
Route both through logged_event_key, the way the Azure storage, GCS,
OpenMeter and DynamoDB loggers already do, so the cost event's ClickHouse
request_id matches its spend row's `<id>_background_cost` and its DeepEval
span gets its own uuid. ClickHouse response_id stays the raw provider id.
* test(deepeval): annotate the background cost span test's locals with Final
* test(responses): run the DELETE wire test that pins the echoed client id
DELETE /v1/responses/{id} now answers with the id the caller sent, so the
skip that marked it as a known bug no longer applies. The test passes on
the /v1, /openai/v1 and bare routes.
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
4ede2c012f
commit
250af8f027
34 changed files with 1981 additions and 352 deletions
|
|
@ -1,7 +1,8 @@
|
|||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from functools import cache
|
||||
from typing import Final
|
||||
|
|
@ -15,6 +16,7 @@ from litellm.constants import (
|
|||
AZURE_STORAGE_MSFT_VERSION,
|
||||
)
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -47,6 +49,20 @@ def adls_safe_file_name(payload_id: str | None) -> str:
|
|||
return f"{name.translate(_ADLS_SAFE_NAME)}.json"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AzureStorageLogQueueItem:
|
||||
payload: StandardLoggingPayload
|
||||
file_name: str
|
||||
|
||||
|
||||
def _queued_log(kwargs: Mapping[str, object], payload: StandardLoggingPayload) -> AzureStorageLogQueueItem:
|
||||
payload_id: Final = payload.get("id")
|
||||
return AzureStorageLogQueueItem(
|
||||
payload=payload,
|
||||
file_name=adls_safe_file_name(payload_id and logged_event_key(kwargs, payload_id)),
|
||||
)
|
||||
|
||||
|
||||
@cache
|
||||
def _cached_credential_chain_token_provider() -> Callable[[], str]:
|
||||
return get_azure_ad_token_provider(
|
||||
|
|
@ -99,7 +115,7 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
self.flush_lock = asyncio.Lock()
|
||||
self.log_queue: list[StandardLoggingPayload] = []
|
||||
self.log_queue: list[AzureStorageLogQueueItem] = []
|
||||
super().__init__(**kwargs, flush_lock=self.flush_lock)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
|
|
@ -115,7 +131,7 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
def azure_storage_blob_endpoint(self) -> str:
|
||||
return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}"
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
async def async_log_success_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime):
|
||||
"""
|
||||
Async Log success events to Azure Blob Storage
|
||||
|
||||
|
|
@ -133,12 +149,12 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
if standard_logging_payload is None:
|
||||
raise ValueError("standard_logging_payload is not set")
|
||||
|
||||
self.log_queue.append(standard_logging_payload)
|
||||
self.log_queue.append(_queued_log(kwargs, standard_logging_payload))
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e)
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
async def async_log_failure_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime):
|
||||
"""
|
||||
Async Log failure events to Azure Blob Storage
|
||||
|
||||
|
|
@ -156,7 +172,7 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
if standard_logging_payload is None:
|
||||
raise ValueError("standard_logging_payload is not set")
|
||||
|
||||
self.log_queue.append(standard_logging_payload)
|
||||
self.log_queue.append(_queued_log(kwargs, standard_logging_payload))
|
||||
except Exception as e:
|
||||
verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e)
|
||||
|
||||
|
|
@ -195,13 +211,15 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
len(self.log_queue),
|
||||
)
|
||||
|
||||
for payload in self.log_queue:
|
||||
await self.async_upload_payload_to_azure_blob_storage(payload=payload)
|
||||
for item in self.log_queue:
|
||||
await self.async_upload_payload_to_azure_blob_storage(payload=item.payload, file_name=item.file_name)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("AzureBlobStorageLogger Error sending batch API - %s", e)
|
||||
|
||||
async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload):
|
||||
async def async_upload_payload_to_azure_blob_storage(
|
||||
self, payload: StandardLoggingPayload, file_name: str | None = None
|
||||
):
|
||||
"""
|
||||
Uploads the payload to Azure Blob Storage using a 3-step process:
|
||||
1. Create file resource
|
||||
|
|
@ -210,14 +228,14 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
"""
|
||||
try:
|
||||
if self.azure_storage_account_key:
|
||||
await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload)
|
||||
await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload, file_name=file_name)
|
||||
else:
|
||||
# Get a valid token instead of always requesting a new one
|
||||
await self.set_valid_azure_ad_token()
|
||||
async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry
|
||||
payload_bytes: Final = json_payload.encode("utf-8")
|
||||
filename: Final = adls_safe_file_name(payload.get("id"))
|
||||
filename: Final = file_name or adls_safe_file_name(payload.get("id"))
|
||||
base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}"
|
||||
|
||||
# Execute the 3-step upload process
|
||||
|
|
@ -377,7 +395,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
self._service_client_timeout = self._clock() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS
|
||||
return self._service_client
|
||||
|
||||
async def upload_to_azure_data_lake_with_azure_account_key(self, payload: StandardLoggingPayload):
|
||||
async def upload_to_azure_data_lake_with_azure_account_key(
|
||||
self, payload: StandardLoggingPayload, file_name: str | None = None
|
||||
):
|
||||
"""
|
||||
Uploads the payload to Azure Data Lake using the Azure SDK
|
||||
|
||||
|
|
@ -386,11 +406,11 @@ class AzureBlobStorageLogger(CustomBatchLogger):
|
|||
|
||||
try:
|
||||
today: Final = datetime.now().strftime("%Y-%m-%d")
|
||||
file_name: Final = adls_safe_file_name(payload.get("id"))
|
||||
resolved_file_name: Final = file_name or adls_safe_file_name(payload.get("id"))
|
||||
await self._upload_bytes_with_account_key(
|
||||
file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8")
|
||||
file_path=f"{today}/{resolved_file_name}", content=safe_dumps(payload).encode("utf-8")
|
||||
)
|
||||
verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name)
|
||||
verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, resolved_file_name)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error occurred: %s", e)
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger
|
||||
from litellm.integrations.clickhouse.context import is_lens_analysis
|
||||
from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
|
||||
from litellm.tracing.types import SpendLogPayload, SpendLogRecord
|
||||
|
|
@ -134,7 +135,7 @@ def spend_log_row_from_payload(
|
|||
usage: Mapping[str, Any] = metadata.get("usage_object") or hidden_params.get("usage_object") or MappingProxyType({})
|
||||
cache_read_tokens, cache_write_tokens = _cache_tokens(usage)
|
||||
trace_id, span_id = _find_traceparent(metadata, kwargs)
|
||||
request_id = str(payload.get("id") or "")
|
||||
response_id: Final = str(payload.get("id") or "")
|
||||
redact = litellm.turn_off_message_logging is True
|
||||
track_end_user: Final = not litellm.disable_end_user_cost_tracking
|
||||
completion_start_ms = _to_ms(payload.get("completionStartTime"))
|
||||
|
|
@ -157,8 +158,8 @@ def spend_log_row_from_payload(
|
|||
)
|
||||
)
|
||||
return SpendLogRecord(
|
||||
request_id=request_id,
|
||||
response_id=strip_cache_hit_suffix(request_id),
|
||||
request_id=logged_event_key(kwargs, response_id),
|
||||
response_id=strip_cache_hit_suffix(response_id),
|
||||
provider_request_id=(
|
||||
get_provider_request_id(kwargs.get("response_headers"))
|
||||
or get_provider_request_id(hidden_params.get("additional_headers"))
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.integrations.deepeval.utils import (
|
|||
to_zod_compatible_iso,
|
||||
validate_environment,
|
||||
)
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
|
||||
|
||||
# This file includes the custom callbacks for LiteLLM Proxy
|
||||
|
|
@ -109,8 +110,10 @@ class DeepEvalLogger(CustomLogger):
|
|||
)
|
||||
else:
|
||||
output = str(standard_logging_object.get("error_string", ""))
|
||||
logged_id: Final = standard_logging_object.get("id", uuid.uuid4())
|
||||
span_id: Final = logged_event_key(kwargs, logged_id) if isinstance(logged_id, str) else logged_id
|
||||
return BaseApiSpan(
|
||||
uuid=standard_logging_object.get("id", uuid.uuid4()),
|
||||
uuid=span_id,
|
||||
name=("litellm_success_callback" if is_success else "litellm_failure_callback"),
|
||||
status=(TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED),
|
||||
type=SpanApiType.LLM,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ from typing import Final, Protocol
|
|||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
|
||||
|
||||
class _DynamoTable(Protocol):
|
||||
|
|
@ -22,11 +24,13 @@ class _DynamoResource(Protocol):
|
|||
class DyanmoDBLogger:
|
||||
# Class variables or attributes
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, dynamodb: _DynamoResource | None = None):
|
||||
# Instance variables
|
||||
import boto3
|
||||
|
||||
self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"])
|
||||
self.dynamodb: Final[_DynamoResource] = dynamodb or boto3.resource(
|
||||
"dynamodb", region_name=os.environ["AWS_REGION_NAME"]
|
||||
)
|
||||
if litellm.dynamodb_table_name is None:
|
||||
raise ValueError(
|
||||
"LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=<your-table>`"
|
||||
|
|
@ -45,23 +49,23 @@ class DyanmoDBLogger:
|
|||
|
||||
_async_log_event = async_log_event
|
||||
|
||||
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
|
||||
def log_event(self, kwargs: Mapping[str, object], response_obj, start_time, end_time, print_verbose):
|
||||
try:
|
||||
print_verbose(f"DynamoDB Logging - Enters logging function for model {kwargs}")
|
||||
|
||||
# construct payload to send to DynamoDB
|
||||
# follows the same params as langfuse.py
|
||||
litellm_params: Final = kwargs.get("litellm_params", {})
|
||||
metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {}
|
||||
messages: Final = kwargs.get("messages")
|
||||
optional_params: Final = kwargs.get("optional_params", {})
|
||||
call_type: Final = kwargs.get("call_type", "litellm.completion")
|
||||
usage: Final = response_obj["usage"]
|
||||
id: Final = response_obj.get("id", str(uuid.uuid4()))
|
||||
response_id: Final = response_obj.get("id", str(uuid.uuid4()))
|
||||
|
||||
# Build the initial payload
|
||||
payload: Final[dict[str, object]] = {
|
||||
"id": id,
|
||||
"id": logged_event_key(kwargs, response_id) if isinstance(response_id, str) else response_id,
|
||||
"call_type": call_type,
|
||||
"startTime": start_time,
|
||||
"endTime": end_time,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.litellm_core_utils.cloud_storage_security import (
|
|||
sanitize_cloud_object_component,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.integrations.gcs_bucket import *
|
||||
from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload
|
||||
|
|
@ -296,30 +297,28 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
else:
|
||||
await self._send_individual_logs(items_to_process)
|
||||
|
||||
def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str:
|
||||
def _get_object_name(
|
||||
self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any
|
||||
) -> str:
|
||||
"""
|
||||
Get the object name to use for the current payload
|
||||
"""
|
||||
current_date: Final = self._get_object_date_from_datetime(datetime.now(timezone.utc))
|
||||
metadata: Final = as_str_mapping((as_str_mapping(kwargs.get("litellm_params")) or {}).get("metadata")) or {}
|
||||
gcs_log_id: Final = metadata.get("gcs_log_id")
|
||||
safe_log_id: Final = sanitize_cloud_object_component(
|
||||
gcs_log_id if isinstance(gcs_log_id, str) else None, fallback=""
|
||||
)
|
||||
if safe_log_id:
|
||||
return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}"
|
||||
if logging_payload.get("error_str", None) is not None:
|
||||
object_name = self._generate_failure_object_name(
|
||||
request_date_str=current_date,
|
||||
)
|
||||
else:
|
||||
object_name = self._generate_success_object_name(
|
||||
request_date_str=current_date,
|
||||
response_id=response_obj.get("id", ""),
|
||||
)
|
||||
|
||||
# used for testing
|
||||
_litellm_params: Final = kwargs.get("litellm_params", None) or {}
|
||||
_metadata: Final = _litellm_params.get("metadata", None) or {}
|
||||
if "gcs_log_id" in _metadata:
|
||||
safe_log_id: Final = sanitize_cloud_object_component(_metadata.get("gcs_log_id"), fallback="")
|
||||
if safe_log_id:
|
||||
object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}"
|
||||
|
||||
return object_name
|
||||
return self._generate_failure_object_name(request_date_str=current_date)
|
||||
raw_response_id: Final[object] = response_obj.get("id", "")
|
||||
response_id: Final = raw_response_id if isinstance(raw_response_id, str) else ""
|
||||
return self._generate_success_object_name(
|
||||
request_date_str=current_date,
|
||||
response_id=logged_event_key(kwargs, response_id),
|
||||
)
|
||||
|
||||
async def get_request_response_payload(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -3,12 +3,15 @@
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
HTTPHandler,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -48,7 +51,7 @@ class OpenMeterLogger(CustomLogger):
|
|||
if len(missing_keys) > 0:
|
||||
raise Exception(f"Missing keys={missing_keys} in environment.")
|
||||
|
||||
def _common_logic(self, kwargs: dict, response_obj):
|
||||
def _common_logic(self, kwargs: Mapping[str, object], response_obj):
|
||||
call_id: Final = response_obj.get("id", kwargs.get("litellm_call_id"))
|
||||
dt: Final = get_utc_datetime().isoformat()
|
||||
cost: Final = kwargs.get("response_cost", None)
|
||||
|
|
@ -74,8 +77,8 @@ class OpenMeterLogger(CustomLogger):
|
|||
# If no user provided directly, try to get it from token user_id
|
||||
if user_param is None:
|
||||
# Check if user_id is available from the API key metadata
|
||||
litellm_params: Final = kwargs.get("litellm_params", {})
|
||||
metadata: Final = litellm_params.get("metadata", {})
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {}
|
||||
user_api_key_user_id: Final = metadata.get("user_api_key_user_id", None)
|
||||
|
||||
if user_api_key_user_id is not None:
|
||||
|
|
@ -89,7 +92,7 @@ class OpenMeterLogger(CustomLogger):
|
|||
return {
|
||||
"specversion": "1.0",
|
||||
"type": os.getenv("OPENMETER_EVENT_TYPE", "litellm_tokens"),
|
||||
"id": call_id,
|
||||
"id": logged_event_key(kwargs, call_id) if isinstance(call_id, str) else call_id,
|
||||
"time": dt,
|
||||
"subject": subject,
|
||||
"source": "litellm-proxy",
|
||||
|
|
|
|||
|
|
@ -26,7 +26,12 @@ from litellm.constants import (
|
|||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
|
||||
|
|
@ -35,6 +40,8 @@ from litellm.types.utils import (
|
|||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost"
|
||||
|
||||
ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved"
|
||||
OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved"
|
||||
ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key"
|
||||
|
|
@ -151,6 +158,26 @@ def is_unbilled_non_inference_call_from_params(
|
|||
return is_unbilled_non_inference_call(call_type, metadata, response)
|
||||
|
||||
|
||||
def settles_background_response(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> bool:
|
||||
return kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES or (
|
||||
call_type in NON_INFERENCE_CALL_TYPES
|
||||
and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
)
|
||||
|
||||
|
||||
def logged_event_key(kwargs: Mapping[str, object], event_id: str) -> str:
|
||||
call_type: Final = kwargs.get("call_type")
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
metadata: Final = (
|
||||
as_str_mapping(litellm_params.get("litellm_metadata")) or as_str_mapping(litellm_params.get("metadata")) or {}
|
||||
)
|
||||
if not event_id or not settles_background_response(
|
||||
call_type if isinstance(call_type, str) else "", metadata, kwargs
|
||||
):
|
||||
return event_id
|
||||
return f"{event_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"
|
||||
|
||||
|
||||
def sanitize_user_api_key_auth(auth: object) -> object:
|
||||
"""Copy of the auth object with its budget reservation removed; the cost callback
|
||||
falls back to reading the reservation from inside the auth object."""
|
||||
|
|
|
|||
|
|
@ -872,6 +872,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# The model the proxy shows the client on streamed chunks. The logged streamed response carries it
|
||||
# once that response is priced, the same way a non-streamed response is logged
|
||||
self.client_facing_stream_model: str | None = None
|
||||
self.client_facing_response_id: str | None = None
|
||||
self.zero_cost_warned: bool = False
|
||||
self.llm_caching_handler: LLMCachingHandler | None = None
|
||||
|
||||
|
|
@ -3238,6 +3239,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
return response
|
||||
return response.model_copy(update={"id": response_id})
|
||||
|
||||
def _with_client_facing_response_id(
|
||||
self,
|
||||
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
|
||||
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
|
||||
response_id: Final = self.client_facing_response_id
|
||||
if response_id is None or getattr(response, "id", None) == response_id:
|
||||
return response
|
||||
return response.model_copy(update={"id": response_id})
|
||||
|
||||
def _client_facing_streaming_response(
|
||||
self,
|
||||
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
|
||||
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
|
||||
return self._with_client_facing_response_id(
|
||||
self._with_continued_response_id(self._with_client_facing_stream_model(response))
|
||||
)
|
||||
|
||||
def _success_handler_helper_fn(
|
||||
self,
|
||||
result: object = None,
|
||||
|
|
@ -3641,9 +3659,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
self._add_continued_partial_spend(complete_streaming_response)
|
||||
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
|
||||
logged_streaming_response: Final = self._with_continued_response_id(
|
||||
self._with_client_facing_stream_model(complete_streaming_response)
|
||||
)
|
||||
logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response)
|
||||
self.model_call_details["complete_streaming_response"] = logged_streaming_response
|
||||
self.capture_usage_before_logging(logged_streaming_response, start_time, end_time)
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
|
|
@ -4240,9 +4256,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
await self._prepare_baseline_cache_estimate(complete_streaming_response)
|
||||
|
||||
logged_streaming_response: Final = self._with_continued_response_id(
|
||||
self._with_client_facing_stream_model(complete_streaming_response)
|
||||
)
|
||||
logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response)
|
||||
self.model_call_details["async_complete_streaming_response"] = logged_streaming_response
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
|
|
|
|||
|
|
@ -7,22 +7,25 @@ instead of writing immediately on each request.
|
|||
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeVar
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
BaseLiteLLMOpenAIResponseObject,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -67,6 +70,8 @@ def _proxy_signing_key() -> str | None:
|
|||
|
||||
|
||||
_RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_ResponseObjectT = TypeVar("_ResponseObjectT", bound=BaseLiteLLMOpenAIResponseObject)
|
||||
_AddressedResultT = TypeVar("_AddressedResultT", bound=ResponsesAPIResponse | DeleteResponseResult)
|
||||
|
||||
|
||||
def _response_payload(response_obj: object) -> Mapping[str, object] | None:
|
||||
|
|
@ -77,9 +82,9 @@ def _response_payload(response_obj: object) -> Mapping[str, object] | None:
|
|||
|
||||
|
||||
def _rewrite_advertised_id(
|
||||
event: BaseLiteLLMOpenAIResponseObject,
|
||||
event: _ResponseObjectT,
|
||||
rewrite: Callable[[str], str],
|
||||
) -> BaseLiteLLMOpenAIResponseObject:
|
||||
) -> _ResponseObjectT:
|
||||
event_id: Final = getattr(event, "id", None)
|
||||
if isinstance(event_id, str) and event_id.startswith("resp_"):
|
||||
setattr(event, "id", rewrite(event_id))
|
||||
|
|
@ -116,6 +121,35 @@ def _streams_response_ids(request_route: str | None) -> bool:
|
|||
return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "")
|
||||
|
||||
|
||||
_NO_HELD_IDS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _upstream_response_id(response_id: str) -> str:
|
||||
return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id)
|
||||
|
||||
|
||||
def _names_the_same_response(response_id: str | None, original_response_id: str) -> bool:
|
||||
return response_id is not None and _upstream_response_id(response_id) == _upstream_response_id(original_response_id)
|
||||
|
||||
|
||||
def _held_id_naming(response_id: str | None, held_ids: Mapping[str, str]) -> str | None:
|
||||
return next(
|
||||
(
|
||||
held_id
|
||||
for original_response_id, held_id in held_ids.items()
|
||||
if _names_the_same_response(response_id, original_response_id)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _id_sent_back(response: ResponsesAPIResponse | DeleteResponseResult, held_ids: Mapping[str, str]) -> str | None:
|
||||
held_id: Final = _held_id_naming(response.id, held_ids)
|
||||
if held_id is not None:
|
||||
return held_id
|
||||
return response.id if isinstance(response, DeleteResponseResult) else None
|
||||
|
||||
|
||||
class ResponsesIDSecurity(CustomLogger):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -275,10 +309,18 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
|
||||
def _encrypt_response_id(
|
||||
self,
|
||||
response: BaseLiteLLMOpenAIResponseObject,
|
||||
response: _ResponseObjectT,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
request_cache: dict[str, str] | None = None,
|
||||
) -> BaseLiteLLMOpenAIResponseObject:
|
||||
) -> _ResponseObjectT:
|
||||
encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache)
|
||||
return response if encrypt is None else _rewrite_advertised_id(response, encrypt)
|
||||
|
||||
def _response_id_encrypter(
|
||||
self,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
request_cache: dict[str, str] | None,
|
||||
) -> Callable[[str], str] | None:
|
||||
signing_key: Final = self._get_signing_key()
|
||||
if signing_key is None:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -287,7 +329,7 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
"Skipping response ID encryption. "
|
||||
"See: https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key"
|
||||
)
|
||||
return response
|
||||
return None
|
||||
|
||||
def encrypt(original_id: str) -> str:
|
||||
cached: Final = request_cache.get(original_id) if request_cache is not None else None
|
||||
|
|
@ -304,47 +346,73 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
request_cache[original_id] = encrypted_id
|
||||
return encrypted_id
|
||||
|
||||
return _rewrite_advertised_id(response, encrypt)
|
||||
return encrypt
|
||||
|
||||
def _streamed_response_id_advertiser(
|
||||
self,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
request_data: Mapping[str, object],
|
||||
) -> Callable[[str], str] | None:
|
||||
if not _streams_response_ids(user_api_key_dict.request_route):
|
||||
return None
|
||||
if self._general_settings_reader().get("disable_responses_id_security", False):
|
||||
return None
|
||||
encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache={})
|
||||
if encrypt is None:
|
||||
return None
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
held_ids: Final = self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD))
|
||||
|
||||
def advertise(original_id: str) -> str:
|
||||
logged_id: Final = encrypt(original_id)
|
||||
if isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
logging_obj.client_facing_response_id = logged_id
|
||||
return _held_id_naming(original_id, held_ids) or logged_id
|
||||
|
||||
return advertise
|
||||
|
||||
def _ids_the_caller_holds(self, addressed_id: object) -> Mapping[str, str]:
|
||||
if not isinstance(addressed_id, str) or not self._is_encrypted_response_id(addressed_id):
|
||||
return _NO_HELD_IDS
|
||||
original_response_id, _, _ = self._decrypt_response_id(addressed_id)
|
||||
return MappingProxyType({original_response_id: addressed_id})
|
||||
|
||||
def _advertised_response(
|
||||
self,
|
||||
response: _AddressedResultT,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
request_data: Mapping[str, object],
|
||||
) -> _AddressedResultT:
|
||||
sent_back_id: Final = _id_sent_back(
|
||||
response, self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD))
|
||||
)
|
||||
logged_response: Final = self._encrypt_response_id(response, user_api_key_dict, request_cache=None)
|
||||
return logged_response if sent_back_id is None else logged_response.model_copy(update={"id": sent_back_id})
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
data: dict[str, object],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
response: LLMResponseTypes,
|
||||
) -> LLMResponseTypes:
|
||||
"""
|
||||
Queue response IDs for batch processing instead of writing directly to DB.
|
||||
|
||||
This method adds response IDs to an in-memory queue, which are then
|
||||
batch-processed by the DBSpendUpdateWriter during regular database update cycles.
|
||||
"""
|
||||
general_settings: Final = self._general_settings_reader()
|
||||
|
||||
if general_settings.get("disable_responses_id_security", False):
|
||||
return response
|
||||
if isinstance(response, ResponsesAPIResponse):
|
||||
response = cast(
|
||||
ResponsesAPIResponse,
|
||||
self._encrypt_response_id(response, user_api_key_dict, request_cache=None),
|
||||
)
|
||||
if isinstance(response, ResponsesAPIResponse | DeleteResponseResult):
|
||||
return self._advertised_response(response, user_api_key_dict, data)
|
||||
return response
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict
|
||||
self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object]
|
||||
) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]:
|
||||
general_settings: Final = self._general_settings_reader()
|
||||
|
||||
# Create a request-scoped cache for consistent encryption across streaming chunks.
|
||||
request_encryption_cache: Final[dict[str, str]] = {}
|
||||
|
||||
advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data)
|
||||
async for chunk in response:
|
||||
if (
|
||||
isinstance(chunk, BaseLiteLLMOpenAIResponseObject)
|
||||
and _streams_response_ids(user_api_key_dict.request_route)
|
||||
and not general_settings.get("disable_responses_id_security", False)
|
||||
):
|
||||
chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache)
|
||||
yield chunk
|
||||
yield (
|
||||
_rewrite_advertised_id(chunk, advertise)
|
||||
if advertise is not None and isinstance(chunk, BaseLiteLLMOpenAIResponseObject)
|
||||
else chunk
|
||||
)
|
||||
|
||||
|
||||
def decrypted_response_id(response_id: str) -> str:
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ from litellm.constants import (
|
|||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
CLI_SESSION_KEY_PREFIX,
|
||||
EMPTY_MAPPING,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
LITELLM_PROXY_MASTER_KEY_ALIAS,
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
|
|
@ -30,7 +29,6 @@ from litellm.constants import (
|
|||
MAX_SPEND_LOG_MODEL_NAME_LENGTH,
|
||||
MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB,
|
||||
MCP_SPEND_LOG_MODEL_PREFIX,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
REDACTED_BY_LITELM_STRING,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
UNKNOWN_MODEL_SPEND_LOG_MODEL,
|
||||
|
|
@ -53,7 +51,11 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
)
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
from litellm.litellm_core_utils.hidden_params import served_from_cache
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
|
||||
from litellm.litellm_core_utils.internal_call_metadata import (
|
||||
BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX,
|
||||
is_unbilled_non_inference_call,
|
||||
settles_background_response,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
UsageObservation,
|
||||
|
|
@ -79,13 +81,8 @@ from litellm.responses.litellm_completion_transformation.compaction_items import
|
|||
summary_only_compaction_content,
|
||||
with_compaction_summary,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag
|
||||
from litellm.types.utils import (
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
PROMPT_CARRYING_GUARDRAIL_FIELDS,
|
||||
AzureSpillover,
|
||||
CallTypes,
|
||||
|
|
@ -347,24 +344,21 @@ def _get_spend_logs_metadata(
|
|||
|
||||
|
||||
BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost"
|
||||
BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost"
|
||||
|
||||
|
||||
def _cost_row_request_id_suffix(
|
||||
call_type: str, metadata: Mapping[str, object], settles_background_response: bool
|
||||
) -> str:
|
||||
def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str:
|
||||
if call_type == CallTypes.aretrieve_batch.value:
|
||||
return BATCH_COST_REQUEST_ID_SUFFIX
|
||||
if settles_background_response or (
|
||||
call_type in NON_INFERENCE_CALL_TYPES
|
||||
and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
):
|
||||
if settles_background_response(call_type, metadata, kwargs):
|
||||
return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
|
||||
return ""
|
||||
|
||||
|
||||
def get_spend_logs_id(
|
||||
call_type: str, response_obj: dict, kwargs: Mapping[str, object], metadata: Mapping[str, object] = EMPTY_MAPPING
|
||||
call_type: str,
|
||||
response_obj: Mapping[str, object],
|
||||
kwargs: Mapping[str, object],
|
||||
metadata: Mapping[str, object] = EMPTY_MAPPING,
|
||||
) -> str | None:
|
||||
standard_logging_payload = kwargs.get("standard_logging_object")
|
||||
candidate_ids: Final = (
|
||||
|
|
@ -377,10 +371,7 @@ def get_spend_logs_id(
|
|||
)
|
||||
if resolved_id is None:
|
||||
return None
|
||||
settles_background_response: Final = (
|
||||
kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES
|
||||
)
|
||||
return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, settles_background_response)}"
|
||||
return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}"
|
||||
|
||||
|
||||
_MISSING_ATTRIBUTE: Final = object()
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER
|
||||
from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.responses.litellm_completion_transformation.client_tool_search import (
|
||||
bridge_client_tool_search,
|
||||
client_tool_search_loaded_tools,
|
||||
|
|
@ -69,6 +71,90 @@ def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool:
|
|||
return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES and not _is_a_failed_turn(spend_log)
|
||||
|
||||
|
||||
_RESPONSE_READ_CALL_TYPES: Final = frozenset(("aget_responses", "acancel_responses"))
|
||||
|
||||
|
||||
def _output_item_id(item: object) -> str | None:
|
||||
match item:
|
||||
case {"id": str() as item_id}:
|
||||
return item_id
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _is_a_response_read(spend_log: "SpendLogsPayload") -> bool:
|
||||
return "call_type" in spend_log and spend_log["call_type"] in _RESPONSE_READ_CALL_TYPES
|
||||
|
||||
|
||||
def _stored_response(spend_log: "SpendLogsPayload") -> object:
|
||||
return spend_log.get("response")
|
||||
|
||||
|
||||
def _output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str | None, object]]:
|
||||
output: Final = _stored_responses_api_output(_stored_response(spend_log)) or []
|
||||
return zip(map(_output_item_id, output), output)
|
||||
|
||||
|
||||
def _identified_output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str, object]]:
|
||||
return ((item_id, item) for item_id, item in _output_items(spend_log) if item_id is not None)
|
||||
|
||||
|
||||
def _output_item_ids(spend_log: "SpendLogsPayload") -> frozenset[str]:
|
||||
return frozenset(item_id for item_id, _ in _identified_output_items(spend_log))
|
||||
|
||||
|
||||
def _first_index_holding_each_item(indexed_turns: Iterable[tuple[int, "SpendLogsPayload"]]) -> Mapping[str, int]:
|
||||
return dict(
|
||||
chain.from_iterable(
|
||||
((item_id, index) for item_id in _output_item_ids(turn)) for index, turn in reversed(tuple(indexed_turns))
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _turn_replaying_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, int]:
|
||||
indexed_turns: Final = tuple(enumerate(turns))
|
||||
first_read_holding: Final = _first_index_holding_each_item(
|
||||
(index, turn) for index, turn in indexed_turns if _is_a_response_read(turn)
|
||||
)
|
||||
first_other_turn_holding: Final = _first_index_holding_each_item(
|
||||
(index, turn) for index, turn in indexed_turns if not _is_a_response_read(turn)
|
||||
)
|
||||
return MappingProxyType({**first_read_holding, **first_other_turn_holding})
|
||||
|
||||
|
||||
def _latest_read_copy_of_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
dict(chain.from_iterable(_identified_output_items(turn) for turn in turns if _is_a_response_read(turn)))
|
||||
)
|
||||
|
||||
|
||||
def _with_only_new_output(
|
||||
index: int, turn: "SpendLogsPayload", turn_replaying: Mapping[str, int], latest_copy: Mapping[str, object]
|
||||
) -> "SpendLogsPayload | None":
|
||||
if not _is_a_response_read(turn):
|
||||
return turn
|
||||
stored_response: Final = as_str_mapping(_stored_response(turn))
|
||||
new_output: Final = [
|
||||
item if item_id is None else latest_copy[item_id]
|
||||
for item_id, item in _output_items(turn)
|
||||
if item_id is None or turn_replaying.get(item_id) == index
|
||||
]
|
||||
if stored_response is None or not new_output:
|
||||
return None
|
||||
replayed_turn: Final[SpendLogsPayload] = {**turn, "response": {**stored_response, "output": new_output}}
|
||||
return replayed_turn
|
||||
|
||||
|
||||
def _conversation_turns(spend_logs: Sequence["SpendLogsPayload"]) -> tuple["SpendLogsPayload", ...]:
|
||||
turns: Final = tuple(filter(_is_a_conversation_turn, spend_logs))
|
||||
turn_replaying: Final = _turn_replaying_each_item(turns)
|
||||
latest_copy: Final = _latest_read_copy_of_each_item(turns)
|
||||
replayed: Final = (
|
||||
_with_only_new_output(index, turn, turn_replaying, latest_copy) for index, turn in enumerate(turns)
|
||||
)
|
||||
return tuple(turn for turn in replayed if turn is not None)
|
||||
|
||||
|
||||
_MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results"))
|
||||
|
||||
|
||||
|
|
@ -101,7 +187,7 @@ class StoredTurns:
|
|||
def _compaction_point(
|
||||
index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None
|
||||
) -> _ReplayStart | None:
|
||||
if response_holds_compaction_block(spend_log.get("response")):
|
||||
if response_holds_compaction_block(_stored_response(spend_log)):
|
||||
return _ReplayStart(index=index, at_output=True)
|
||||
if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")):
|
||||
return _ReplayStart(index=index, at_output=False)
|
||||
|
|
@ -271,12 +357,10 @@ class ResponsesSessionHandler:
|
|||
previous_response_id: str,
|
||||
client_response_ids: tuple[str, ...] = (),
|
||||
) -> StoredTurns:
|
||||
spend_logs: Final = tuple(
|
||||
spend_log
|
||||
for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(
|
||||
spend_logs: Final = _conversation_turns(
|
||||
await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(
|
||||
previous_response_id, client_response_ids
|
||||
)
|
||||
if _is_a_conversation_turn(spend_log)
|
||||
)
|
||||
return StoredTurns(
|
||||
spend_logs=spend_logs,
|
||||
|
|
@ -378,7 +462,7 @@ class ResponsesSessionHandler:
|
|||
############################################################
|
||||
# Add Output messages for this Spend Log
|
||||
############################################################
|
||||
_response_output: Final[object] = spend_log.get("response", "{}")
|
||||
_response_output: Final = _stored_response(spend_log)
|
||||
if (responses_api_output := _stored_responses_api_output(_response_output)) is not None:
|
||||
chat_completion_message_history.extend(
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
|
|
|
|||
|
|
@ -69,9 +69,6 @@ def _retrieve_response(proxy: ProxyClient, key: str, response_id: str) -> Respon
|
|||
|
||||
|
||||
class TestResponsesRetrieve:
|
||||
@pytest.mark.skip(
|
||||
reason="stage red: product gap (LIT-5446), retrieve returns a different id than the stored response (non-idempotent response-id re-encryption)"
|
||||
)
|
||||
@pytest.mark.covers("llm.responses.openai.basic.nonstream.works")
|
||||
@meta(
|
||||
Subject(
|
||||
|
|
|
|||
|
|
@ -630,19 +630,6 @@ def _spend_row(identity: str) -> Mapping[str, JsonValue]:
|
|||
return rows[0]
|
||||
|
||||
|
||||
def _spend_row_keyed_by_any(identities: frozenset[str]) -> Mapping[str, JsonValue]:
|
||||
keys: Final = tuple(sorted(identities))
|
||||
placeholders: Final = ", ".join("%s" for _ in keys)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
f'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', keys
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _spend_request_ids(identities: tuple[str, str, str]) -> frozenset[str]:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities),
|
||||
|
|
@ -815,7 +802,7 @@ def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) ->
|
|||
served: Final = _served(wire)
|
||||
assert _routes(served) == (("POST", "/v1/responses"),), served
|
||||
assert _body(served[0])["stream"] is True, served[0].body
|
||||
assert _spend_row_keyed_by_any(response_identities(last.response.id))["model"] == DEPLOYMENT_MODEL
|
||||
assert _spend_row(last.response.id)["model"] == DEPLOYMENT_MODEL
|
||||
|
||||
|
||||
def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None:
|
||||
|
|
|
|||
|
|
@ -107,9 +107,11 @@ def test_langfuse_otel_sink_outage_mid_burst_lands_every_response_id_once(gatewa
|
|||
)
|
||||
payloads: Final = tuple(json.loads(line) for line in events)
|
||||
found: Final = next(
|
||||
payload.get("response_id") or (payload.get("message") or {}).get("id") or payload.get("id")
|
||||
(payload.get("response") or {}).get("id")
|
||||
or (payload.get("message") or {}).get("id")
|
||||
or payload.get("id")
|
||||
for payload in payloads
|
||||
if payload.get("response_id") or payload.get("message") or payload.get("id")
|
||||
if payload.get("response") or payload.get("message") or payload.get("id")
|
||||
)
|
||||
return str(found)
|
||||
response: Final = rig.candidate.client.post(path, json=body, headers={"Authorization": f"Bearer {key}"})
|
||||
|
|
|
|||
|
|
@ -162,8 +162,13 @@ def test_langfuse_otel_metadata_reaches_langfuse_for_responses_sse_stream(gatewa
|
|||
for line in stream.iter_lines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
response_id: Final = next(
|
||||
event["response_id"] for event in events if event.get("type") == "response.output_text.delta"
|
||||
advertised: Final = frozenset(
|
||||
event["response"]["id"]
|
||||
for event in events
|
||||
if event.get("type") in ("response.created", "response.completed")
|
||||
)
|
||||
assert len(advertised) == 1, events
|
||||
response_id: Final = next(iter(advertised))
|
||||
attrs: Final = exported_llm_span(rig.collector, response_id)
|
||||
assert attrs["gen_ai.response.id"] == response_id, attrs
|
||||
assert_identity(attrs, chat_identity_fields(identity))
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import signal
|
||||
|
|
@ -118,17 +117,6 @@ def _responses_reply(identity: str, stream: bool) -> Reply:
|
|||
)
|
||||
|
||||
|
||||
def _decoded_responses_id(identity: str) -> str:
|
||||
try:
|
||||
return base64.b64decode(identity.removeprefix("resp_").encode()).decode()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return identity
|
||||
|
||||
|
||||
def _canonical_id(identity: str) -> str:
|
||||
return _decoded_responses_id(identity).rpartition("response_id:")[2]
|
||||
|
||||
|
||||
def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]"
|
||||
|
|
@ -170,8 +158,7 @@ class Collector:
|
|||
return tuple(
|
||||
attributes
|
||||
for attributes in self.attributes()
|
||||
if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str)
|
||||
and _canonical_id(logged) == _canonical_id(response_id)
|
||||
if attributes.get("gen_ai.response.id", {}).get("stringValue") == response_id
|
||||
)
|
||||
|
||||
def conversation_ids(self, response_id: str) -> tuple[str | None, ...]:
|
||||
|
|
@ -183,10 +170,6 @@ class Collector:
|
|||
def single_span(self, response_id: str) -> str | None:
|
||||
return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0]
|
||||
|
||||
def logged_id(self, response_id: str) -> str:
|
||||
spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30)
|
||||
return str(spans[0]["gen_ai.response.id"]["stringValue"])
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def collector() -> Iterator[Collector]:
|
||||
|
|
@ -437,12 +420,13 @@ def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig:
|
|||
completed: Final = tuple(event for event in events if event["type"] == "response.completed")
|
||||
assert len(completed) == 1, events
|
||||
assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed
|
||||
assert str(completed[0]["response"]["id"]).startswith("resp_"), completed
|
||||
client_id: Final = str(completed[0]["response"]["id"])
|
||||
assert client_id.startswith("resp_"), completed
|
||||
assert rig.upstream_bodies(marker) == (
|
||||
{"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True},
|
||||
)
|
||||
assert rig.sink.single_span(f"resp_{marker}") == session
|
||||
assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session
|
||||
assert rig.sink.single_span(client_id) == session
|
||||
assert rig.spend_session(client_id) == session
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw")
|
||||
|
|
@ -689,15 +673,11 @@ def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[
|
|||
return tuple(pool.map(one, range(count)))
|
||||
|
||||
|
||||
def _is_encrypted_responses_id(identity: str) -> bool:
|
||||
return identity.startswith("resp_") and _decoded_responses_id(identity) == identity
|
||||
|
||||
|
||||
def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]:
|
||||
spans: Final = rig.sink.attributes()
|
||||
return {
|
||||
session: tuple(
|
||||
_canonical_id(str(attributes["gen_ai.response.id"]["stringValue"]))
|
||||
str(attributes["gen_ai.response.id"]["stringValue"])
|
||||
for attributes in spans
|
||||
if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes
|
||||
)
|
||||
|
|
@ -707,21 +687,13 @@ def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]
|
|||
|
||||
def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None:
|
||||
spend: Final = eventually(
|
||||
lambda: {
|
||||
session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session))
|
||||
for session in expected.values()
|
||||
},
|
||||
lambda: {session: rig.spend_request_ids(session) for session in expected.values()},
|
||||
lambda rows: all(len(values) >= 1 for values in rows.values()),
|
||||
seconds=70,
|
||||
)
|
||||
assert landed == spend, (landed, spend)
|
||||
assert all(len(values) == 1 for values in landed.values()), landed
|
||||
caller_visible: Final = {
|
||||
session: (_canonical_id(identity),)
|
||||
for identity, session in expected.items()
|
||||
if not _is_encrypted_responses_id(identity)
|
||||
}
|
||||
assert {session: landed[session] for session in caller_visible} == caller_visible, landed
|
||||
assert landed == {session: (identity,) for identity, session in expected.items()}, landed
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once")
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import signal
|
||||
|
|
@ -9,7 +8,6 @@ from collections import deque
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -137,17 +135,6 @@ def _upstream(request: Request) -> Reply:
|
|||
return _chat_reply(f"chatcmpl-{marker}", stream)
|
||||
|
||||
|
||||
def _decoded_responses_id(identity: str) -> str:
|
||||
try:
|
||||
return base64.b64decode(identity.removeprefix("resp_").encode()).decode()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return identity
|
||||
|
||||
|
||||
def _canonical_id(identity: str) -> str:
|
||||
return _decoded_responses_id(identity).rpartition("response_id:")[2]
|
||||
|
||||
|
||||
def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
object_value(JSON.validate_json(line[6:]))
|
||||
|
|
@ -232,8 +219,7 @@ class Collector:
|
|||
return tuple(
|
||||
span
|
||||
for span in self.spans()
|
||||
if RESPONSE_ID in span.attributes
|
||||
and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(response_id)
|
||||
if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == response_id
|
||||
)
|
||||
|
||||
def single_span(self, response_id: str, *, elsewhere: "Collector | None" = None) -> Span:
|
||||
|
|
@ -249,11 +235,8 @@ class Collector:
|
|||
def landed(self, response_ids: Sequence[str]) -> dict[str, int]:
|
||||
spans: Final = self.spans()
|
||||
return {
|
||||
_canonical_id(identity): sum(
|
||||
1
|
||||
for span in spans
|
||||
if RESPONSE_ID in span.attributes
|
||||
and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(identity)
|
||||
identity: sum(
|
||||
1 for span in spans if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == identity
|
||||
)
|
||||
for identity in response_ids
|
||||
}
|
||||
|
|
@ -448,7 +431,7 @@ def _assert_operator_span(rig: Rig, response_id: str, marker: str) -> Span:
|
|||
span: Final = rig.sink.single_span(response_id)
|
||||
assert span.target == "/v1/traces", span
|
||||
assert span.ingestion_key == OPERATOR_KEY, span
|
||||
assert rig.tenant_sink.landed((response_id,)) == {_canonical_id(response_id): 0}
|
||||
assert rig.tenant_sink.landed((response_id,)) == {response_id: 0}
|
||||
bodies: Final = rig.upstream_bodies(marker)
|
||||
assert len(bodies) == 1, bodies
|
||||
assert "signoz" not in json.dumps(bodies[0]).replace(marker, ""), bodies[0]
|
||||
|
|
@ -522,7 +505,7 @@ def test_responses_stream_raw_httpx_span_lands_once_after_the_stream_is_consumed
|
|||
marker: Final = _marker()
|
||||
response: Final = rig.proxy.request("POST", "/v1/responses", {"model": rig.model, "input": marker, "stream": True})
|
||||
assert response.status_code == 200, response.text
|
||||
_assert_operator_span(rig, _responses_id(response, marker), marker)
|
||||
_assert_operator_span(rig, _responses_id(response), marker)
|
||||
|
||||
|
||||
def test_v2_flag_on_still_delivers_the_operator_span_with_the_ingestion_key(v2_rig: Rig) -> None:
|
||||
|
|
@ -821,7 +804,7 @@ def _chat_id(response: httpx.Response) -> str:
|
|||
return next(iter(identities))
|
||||
|
||||
|
||||
def _responses_id(response: httpx.Response, marker: str) -> str:
|
||||
def _responses_id(response: httpx.Response) -> str:
|
||||
if not _is_event_stream(response):
|
||||
return _body_id(response)
|
||||
completed: Final = tuple(
|
||||
|
|
@ -830,7 +813,7 @@ def _responses_id(response: httpx.Response, marker: str) -> str:
|
|||
if event.get("type") == "response.completed"
|
||||
)
|
||||
assert len(completed) == 1 and completed[0].startswith("resp_"), response.text
|
||||
return f"resp_{marker}"
|
||||
return completed[0]
|
||||
|
||||
|
||||
def _message_id(response: httpx.Response) -> str:
|
||||
|
|
@ -852,7 +835,7 @@ def _burst(rig: Rig, count: int) -> tuple[tuple[int, str, str | None], ...]:
|
|||
stream: Final = index % 2 == 0
|
||||
path, body, identity_of = (
|
||||
("/v1/chat/completions", {"model": rig.model, "messages": [{"role": "user", "content": marker}]}, _chat_id),
|
||||
("/v1/responses", {"model": rig.model, "input": marker}, partial(_responses_id, marker=marker)),
|
||||
("/v1/responses", {"model": rig.model, "input": marker}, _responses_id),
|
||||
(
|
||||
"/v1/messages",
|
||||
{"model": rig.model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]},
|
||||
|
|
@ -874,8 +857,8 @@ def _assert_exactly_once(rig: Rig, identities: Sequence[str]) -> None:
|
|||
landed: Final = eventually(
|
||||
lambda: rig.sink.landed(identities), lambda seen: all(count >= 1 for count in seen.values()), seconds=80
|
||||
)
|
||||
assert landed == {_canonical_id(identity): 1 for identity in identities}, landed
|
||||
charged: Final = frozenset(identity for identity in identities if not identity.startswith("resp_"))
|
||||
assert landed == {identity: 1 for identity in identities}, landed
|
||||
charged: Final = frozenset(identities)
|
||||
spend: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])',
|
||||
|
|
|
|||
|
|
@ -62,12 +62,12 @@ def _body(request: Request) -> dict[str, JsonValue]:
|
|||
# ciphertext the caller received, because the spend row id is read from response_obj["id"] before the
|
||||
# ResponsesIDSecurity hook rewrites it in place; the row is looked up under both ids until that ordering is fixed on
|
||||
# main
|
||||
def _spend_row(client_id: str, issued_id: str) -> dict[str, JsonValue]:
|
||||
def _spend_row(*request_ids: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE request_id = ANY(%s)",
|
||||
([client_id, issued_id],), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array
|
||||
(list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
|
|
@ -120,4 +120,4 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons
|
|||
assert issued.upstream == f"resp_upstream_{marker}", completed.response.id
|
||||
request: Final = _native_request(wire)
|
||||
assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body
|
||||
assert _spend_row(completed.response.id, issued.issued) == _success_row(model)
|
||||
assert _spend_row(completed.response.id) == _success_row(model)
|
||||
|
|
|
|||
|
|
@ -516,7 +516,7 @@ def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway:
|
|||
for attempt in attempts:
|
||||
assert attempt.body["max_output_tokens"] == 1000 and _limits(attempt.body) == (None, None), attempt.body
|
||||
_assert_logged(issued, model, "success")
|
||||
_assert_logged(_managed_response_id(streamed_issued), model, "success")
|
||||
_assert_logged(streamed_issued, model, "success")
|
||||
|
||||
|
||||
def test_region_prefixed_deployment_on_a_v1_base_sends_the_bare_model_with_max_completion_tokens(
|
||||
|
|
|
|||
|
|
@ -301,11 +301,7 @@ def _spend_row_id(rows: Sequence[Mapping[str, JsonValue]], tag: str) -> str:
|
|||
|
||||
|
||||
def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ...]) -> None:
|
||||
"""One spend row per tagged call, carrying the id the client saw.
|
||||
|
||||
FIXME: a streamed /v1/responses row carries the managed `resp_` id LiteLLM built before the proxy encrypted the
|
||||
advertised one, so that row is matched by tag alone until the spend log reads the id the client received.
|
||||
"""
|
||||
"""One spend row per tagged call, carrying the id the client saw."""
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_tags ?| %s',
|
||||
|
|
@ -316,9 +312,6 @@ def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ...
|
|||
)
|
||||
for (identity, path, stream), tag in zip(served, tags, strict=True):
|
||||
row_id: Final = _spend_row_id(rows, tag)
|
||||
if path == "/v1/responses" and stream:
|
||||
assert row_id.startswith("resp_"), (row_id, identity)
|
||||
continue
|
||||
assert row_id == identity, (path, stream, row_id, identity)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -172,7 +172,11 @@ def _upstream_id_inside(row_id: str) -> str | None:
|
|||
def _row_belongs_to(row_id: str, served: _Served) -> bool:
|
||||
if row_id == _response_id(served):
|
||||
return True
|
||||
return served.call.endpoint == "responses" and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}"
|
||||
return (
|
||||
served.call.endpoint == "responses"
|
||||
and not served.call.stream
|
||||
and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}"
|
||||
)
|
||||
|
||||
|
||||
def _assert_each_success_landed_once(rows: list[dict[str, JsonValue]], served: tuple[_Served, ...]) -> None:
|
||||
|
|
|
|||
|
|
@ -203,9 +203,6 @@ def test_retrieve_routes_to_the_deployment_that_created_the_response(
|
|||
|
||||
@pytest.mark.parametrize("prefix", _ALIASES)
|
||||
def test_retrieve_returns_the_id_the_client_holds(gateway: Gateway, prefix: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: GET /v1/responses/{id} returns a re-encrypted id that differs byte-for-byte from the id the client sent"
|
||||
)
|
||||
raw_id: Final = f"resp_retrieve_id_{uuid.uuid4().hex}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
|
|
@ -697,7 +694,6 @@ def test_delete_returns_the_client_held_id_not_the_provider_id(
|
|||
gateway: Gateway,
|
||||
prefix: str,
|
||||
) -> None:
|
||||
pytest.skip("BUG: DELETE /v1/responses/{id} returns the raw provider id instead of the proxy-issued id")
|
||||
raw_id: Final = f"resp_delete_id_{uuid.uuid4().hex}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
|
|
|
|||
465
tests/integration/spend/test_responses_stream_spend_row_id.py
Normal file
465
tests/integration/spend/test_responses_stream_spend_row_id.py
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.responses_stream import frame, response_object
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_NATIVE_MODEL: Final = "openai/gpt-5.6"
|
||||
_BRIDGED_MODEL: Final = "deepseek/deepseek-v4-pro"
|
||||
_ANSWER: Final = "streamed answer"
|
||||
_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35}
|
||||
_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
|
||||
|
||||
def _sse(payload: dict[str, JsonValue]) -> bytes:
|
||||
return f"data: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _native_completed(upstream_id: str) -> dict[str, JsonValue]:
|
||||
message: Final[dict[str, JsonValue]] = {
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": _ANSWER, "annotations": []}],
|
||||
}
|
||||
return response_object(upstream_id, "completed", output=[message], usage=_USAGE)
|
||||
|
||||
|
||||
def _native_stream(upstream_id: str) -> Reply:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
frame(
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": response_object(upstream_id, "in_progress"),
|
||||
}
|
||||
),
|
||||
frame(
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": _ANSWER,
|
||||
}
|
||||
),
|
||||
frame({"type": "response.completed", "sequence_number": 2, "response": _native_completed(upstream_id)}),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _chat_stream(upstream_id: str, model: str) -> Reply:
|
||||
def chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": upstream_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_sse(chunk({"role": "assistant", "content": _ANSWER}, None)),
|
||||
_sse(
|
||||
{
|
||||
**chunk({}, "stop"),
|
||||
"usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35},
|
||||
}
|
||||
),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
|
||||
if request.method == "GET":
|
||||
stored_id: Final = request.target.removeprefix("/v1/responses/")
|
||||
assert stored_id.startswith("resp_upstream-"), request.target
|
||||
return Reply(body=json.dumps(_native_completed(stored_id)).encode())
|
||||
if request.target.endswith("/cancel"):
|
||||
cancelled_id: Final = request.target.removeprefix("/v1/responses/").removesuffix("/cancel")
|
||||
assert cancelled_id.startswith("resp_upstream-"), request.target
|
||||
return Reply(body=json.dumps(response_object(cancelled_id, "cancelled")).encode())
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
upstream_id: Final = f"upstream-{uuid4().hex}"
|
||||
if request.target == "/v1/responses" and body.get("stream") is not True:
|
||||
return Reply(body=json.dumps(_native_completed(f"resp_{upstream_id}")).encode())
|
||||
assert body.get("stream") is True, body
|
||||
if request.target == "/v1/responses":
|
||||
return _native_stream(f"resp_{upstream_id}")
|
||||
assert request.target == "/v1/chat/completions", request.target
|
||||
return _chat_stream(f"chatcmpl-{upstream_id}", string_value(body["model"]))
|
||||
|
||||
|
||||
def _terminal_response_id(
|
||||
gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None
|
||||
) -> str:
|
||||
response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
frames: Final = tuple(
|
||||
JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
assert frames and frames[-1]["type"] == "response.completed", response.text
|
||||
advertised: Final = tuple(
|
||||
dict.fromkeys(string_value(object_value(event["response"])["id"]) for event in frames if "response" in event)
|
||||
)
|
||||
assert len(advertised) == 1, advertised
|
||||
return advertised[0]
|
||||
|
||||
|
||||
def _created_response_id(
|
||||
gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None
|
||||
) -> str:
|
||||
if body.get("stream") is True:
|
||||
return _terminal_response_id(gateway, key, body, headers)
|
||||
response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
return string_value(JSON_OBJECT.validate_json(response.content)["id"])
|
||||
|
||||
|
||||
def _read_and_cancel(
|
||||
gateway: Gateway, key: str, response_id: str, headers: Mapping[str, str] | None = None
|
||||
) -> tuple[httpx.Response, ...]:
|
||||
responses: Final = tuple(
|
||||
gateway.request(method, path, key=key, headers=headers)
|
||||
for method, path in (
|
||||
("GET", f"/v1/responses/{response_id}"),
|
||||
("GET", f"/v1/responses/{response_id}"),
|
||||
("POST", f"/v1/responses/{response_id}/cancel"),
|
||||
)
|
||||
)
|
||||
assert [response.status_code for response in responses] == [200, 200, 200], [r.text for r in responses]
|
||||
return responses
|
||||
|
||||
|
||||
def _spend_logs_lookup(gateway: Gateway, request_id: str) -> list[JsonValue]:
|
||||
response: Final = gateway.request("GET", "/spend/logs", params={"request_id": request_id})
|
||||
assert response.status_code == 200, response.text
|
||||
return [row["request_id"] for row in _JSON_LIST.validate_json(response.content)]
|
||||
|
||||
|
||||
def _request_ids_for_key(gateway: Gateway, key: str, count: int) -> list[str]:
|
||||
rows: Final = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "endTime"',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda found: len(found) >= count,
|
||||
seconds=70,
|
||||
)
|
||||
return [string_value(row["request_id"]) for row in rows]
|
||||
|
||||
|
||||
def _logs_page_search(gateway: Gateway, client_id: str) -> list[str]:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
page: Final = gateway.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"search": client_id,
|
||||
"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"),
|
||||
},
|
||||
)
|
||||
return [string_value(object_value(row)["request_id"]) for row in list_value(page["data"])]
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("litellm_model", [_NATIVE_MODEL, _BRIDGED_MODEL], ids=["native", "bridge"])
|
||||
def test_streamed_responses_spend_row_is_filed_under_the_id_the_client_saw(
|
||||
gateway: Gateway, litellm_model: str
|
||||
) -> None:
|
||||
with wire_server(_respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=litellm_model, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[model])
|
||||
|
||||
client_id: Final = _terminal_response_id(
|
||||
gateway, key, {"model": model, "input": f"marker-{uuid4().hex}", "stream": True}
|
||||
)
|
||||
|
||||
assert client_id.startswith("resp_"), client_id
|
||||
assert _request_ids_for_key(gateway, key, 1) == [client_id]
|
||||
assert _logs_page_search(gateway, client_id) == [client_id]
|
||||
assert _spend_logs_lookup(gateway, client_id) == [client_id]
|
||||
|
||||
|
||||
def _proxy_config(
|
||||
directory: Path,
|
||||
general_settings: Mapping[str, JsonValue],
|
||||
litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
) -> Path:
|
||||
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
path: Final = directory / "responses-spend.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**stock,
|
||||
"general_settings": {**object_value(stock["general_settings"]), **general_settings},
|
||||
"litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _prompt_storing_config(directory: Path) -> Path:
|
||||
return _proxy_config(directory, {"store_prompts_in_spend_logs": True})
|
||||
|
||||
|
||||
def test_a_follow_up_turn_on_a_streamed_bridged_response_replays_the_first_turn(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
first_prompt: Final = f"first turn {uuid4().hex}"
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing,
|
||||
storing.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[model])
|
||||
first_id: Final = _terminal_response_id(storing, key, {"model": model, "input": first_prompt, "stream": True})
|
||||
assert _request_ids_for_key(storing, key, 1) == [first_id]
|
||||
wire.drain()
|
||||
|
||||
second_id: Final = _terminal_response_id(
|
||||
storing,
|
||||
key,
|
||||
{"model": model, "input": "second turn", "previous_response_id": first_id, "stream": True},
|
||||
)
|
||||
|
||||
replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"])
|
||||
assert [(message["role"], message["content"]) for message in replayed] == [
|
||||
("user", first_prompt),
|
||||
("assistant", _ANSWER),
|
||||
("user", "second turn"),
|
||||
], replayed
|
||||
assert _request_ids_for_key(storing, key, 2) == [first_id, second_id]
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
@pytest.mark.parametrize("stream_first_turn", [True, False], ids=["streamed-first-turn", "non-streamed-first-turn"])
|
||||
def test_a_follow_up_through_the_id_a_retrieve_returned_replays_the_first_turn_once(
|
||||
gateway: Gateway, tmp_path: Path, stream_first_turn: bool
|
||||
) -> None:
|
||||
first_prompt: Final = f"first turn {uuid4().hex}"
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing,
|
||||
storing.scenario() as scenario,
|
||||
):
|
||||
native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[native, bridged])
|
||||
first_id: Final = _created_response_id(
|
||||
storing, key, {"model": native, "input": first_prompt, "stream": stream_first_turn}
|
||||
)
|
||||
|
||||
retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=key)
|
||||
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"])
|
||||
assert retrieved_id == first_id
|
||||
stored: Final = _request_ids_for_key(storing, key, 2)
|
||||
assert stored.count(first_id) == 1, stored
|
||||
assert len(set(stored)) == 2, stored
|
||||
assert _spend_logs_lookup(storing, first_id) == [first_id]
|
||||
wire.drain()
|
||||
|
||||
_terminal_response_id(
|
||||
storing,
|
||||
key,
|
||||
{"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True},
|
||||
)
|
||||
|
||||
replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"])
|
||||
assert [(message["role"], message["content"]) for message in replayed] == [
|
||||
("user", first_prompt),
|
||||
("assistant", _ANSWER),
|
||||
("user", "second turn"),
|
||||
], replayed
|
||||
|
||||
|
||||
def _replayed_turns(wire_body: bytes) -> list[tuple[JsonValue, JsonValue]]:
|
||||
replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire_body)["messages"])
|
||||
return [(message["role"], message["content"]) for message in replayed]
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
@pytest.mark.parametrize("reader", ["proxy-admin", "master-key"])
|
||||
def test_a_reader_other_than_the_creator_continues_through_the_id_its_retrieve_returned(
|
||||
gateway: Gateway, tmp_path: Path, reader: str
|
||||
) -> None:
|
||||
first_prompt: Final = f"first turn {uuid4().hex}"
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing,
|
||||
storing.scenario() as scenario,
|
||||
):
|
||||
native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
creator: Final = scenario.key(user_id=scenario.user(), models=[native, bridged])
|
||||
reader_key: Final = (
|
||||
scenario.key(user_id=scenario.user(user_role="proxy_admin"), models=[native, bridged])
|
||||
if reader == "proxy-admin"
|
||||
else storing.key
|
||||
)
|
||||
first_id: Final = _terminal_response_id(
|
||||
storing, creator, {"model": native, "input": first_prompt, "stream": True}
|
||||
)
|
||||
assert _request_ids_for_key(storing, creator, 1) == [first_id]
|
||||
|
||||
retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=reader_key)
|
||||
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"])
|
||||
wire.drain()
|
||||
_terminal_response_id(
|
||||
storing,
|
||||
reader_key,
|
||||
{"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True},
|
||||
)
|
||||
assert _replayed_turns(wire.drain()[-1].body) == [
|
||||
("user", first_prompt),
|
||||
("assistant", _ANSWER),
|
||||
("user", "second turn"),
|
||||
]
|
||||
assert retrieved_id == first_id
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_reads_and_cancels_sent_with_a_session_header_add_no_replayed_turns(gateway: Gateway, tmp_path: Path) -> None:
|
||||
first_prompt: Final = f"first turn {uuid4().hex}"
|
||||
session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"}
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing,
|
||||
storing.scenario() as scenario,
|
||||
):
|
||||
native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[native, bridged])
|
||||
first_id: Final = _created_response_id(storing, key, {"model": native, "input": first_prompt}, session)
|
||||
|
||||
_read_and_cancel(storing, key, first_id, session)
|
||||
|
||||
assert len(_request_ids_for_key(storing, key, 4)) == 4
|
||||
wire.drain()
|
||||
_terminal_response_id(
|
||||
storing,
|
||||
key,
|
||||
{"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True},
|
||||
session,
|
||||
)
|
||||
assert _replayed_turns(wire.drain()[-1].body) == [
|
||||
("user", first_prompt),
|
||||
("assistant", _ANSWER),
|
||||
("user", "second turn"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_each_read_and_cancel_reaches_callbacks_and_spend_logs_under_an_id_of_its_own(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
logged_call_types: Final = frozenset({"aresponses", "aget_responses", "acancel_responses"})
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
wire_server(lambda _: Reply()) as sink,
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"GENERIC_LOGGER_ENDPOINT": sink.url},
|
||||
config=_proxy_config(tmp_path, {}, {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}),
|
||||
) as logging_proxy,
|
||||
logging_proxy.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[model])
|
||||
first_id: Final = _created_response_id(logging_proxy, key, {"model": model, "input": f"marker-{uuid4().hex}"})
|
||||
|
||||
reads: Final = _read_and_cancel(logging_proxy, key, first_id)
|
||||
|
||||
assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3
|
||||
batches: Final[list[bytes]] = [] # mutable-ok: drain() consumes, later polls keep earlier batches
|
||||
|
||||
def delivered() -> tuple[tuple[str, str], ...]:
|
||||
batches.extend(batch.body for batch in sink.drain())
|
||||
return tuple(
|
||||
(string_value(event["call_type"]), string_value(event["id"]))
|
||||
for event in (object_value(item) for batch in batches for item in _JSON_LIST.validate_json(batch))
|
||||
if event["call_type"] in logged_call_types
|
||||
)
|
||||
|
||||
events: Final = eventually(delivered, lambda found: len(found) >= 4, seconds=30)
|
||||
read_ids: Final = [logged_id for call_type, logged_id in events if call_type != "aresponses"]
|
||||
assert [logged_id for call_type, logged_id in events if call_type == "aresponses"] == [first_id], events
|
||||
assert len(read_ids) == 3, events
|
||||
assert first_id not in read_ids, events
|
||||
assert len(set(read_ids)) == 3, events
|
||||
assert sorted(_request_ids_for_key(logging_proxy, key, 4)) == sorted([first_id, *read_ids])
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_with_id_security_off_reads_and_cancels_add_no_rows_and_no_replayed_turns(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
first_prompt: Final = f"first turn {uuid4().hex}"
|
||||
session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"}
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{},
|
||||
config=_proxy_config(
|
||||
tmp_path, {"store_prompts_in_spend_logs": True, "disable_responses_id_security": True}
|
||||
),
|
||||
) as insecure,
|
||||
insecure.scenario() as scenario,
|
||||
):
|
||||
native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key")
|
||||
key: Final = scenario.key(models=[native, bridged])
|
||||
first_id: Final = _created_response_id(insecure, key, {"model": native, "input": first_prompt}, session)
|
||||
|
||||
reads: Final = _read_and_cancel(insecure, key, first_id, session)
|
||||
|
||||
assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3
|
||||
later_id: Final = _created_response_id(insecure, key, {"model": native, "input": "a later request"})
|
||||
assert _request_ids_for_key(insecure, key, 2) == [first_id, later_id]
|
||||
wire.drain()
|
||||
|
||||
_terminal_response_id(
|
||||
insecure,
|
||||
key,
|
||||
{"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True},
|
||||
session,
|
||||
)
|
||||
|
||||
replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"])
|
||||
assert [(message["role"], message["content"]) for message in replayed] == [
|
||||
("user", first_prompt),
|
||||
("assistant", _ANSWER),
|
||||
("user", "second turn"),
|
||||
], replayed
|
||||
|
|
@ -5,19 +5,28 @@ import re
|
|||
import sys
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
|
||||
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.azure_storage import azure_storage as azure_storage_module
|
||||
from litellm.integrations.azure_storage.azure_storage import (
|
||||
AzureBlobStorageLogger,
|
||||
_cached_credential_chain_token_provider,
|
||||
adls_safe_file_name,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType
|
||||
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
|
||||
from litellm.types.utils import (
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
StandardAuditLogPayload,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -648,3 +657,171 @@ async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_an_account_
|
|||
assert json.loads(uploaded) == audit_log
|
||||
assert file_client.append_data.await_args.kwargs == {"data": uploaded, "offset": 0, "length": len(uploaded)}
|
||||
file_client.flush_data.assert_awaited_once_with(position=len(uploaded), offset=0)
|
||||
|
||||
|
||||
BACKGROUND_RESPONSE_ID: Final = "resp_YWJjZA=="
|
||||
BACKGROUND_COST_POLL_KWARGS: Final = {
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {
|
||||
"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN},
|
||||
},
|
||||
}
|
||||
PASSTHROUGH_SETTLEMENT_KWARGS: Final = {
|
||||
"call_type": "pass_through_endpoint",
|
||||
"litellm_params": {"metadata": {}},
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES,
|
||||
}
|
||||
|
||||
|
||||
def _background_create_cost_and_other_events(
|
||||
cost_kwargs: dict[str, object],
|
||||
) -> tuple[tuple[dict[str, object], StandardLoggingPayload], ...]:
|
||||
create: Final = StandardLoggingPayload(
|
||||
id=BACKGROUND_RESPONSE_ID,
|
||||
call_type="aresponses",
|
||||
messages=[{"role": "user", "content": "background create input"}],
|
||||
)
|
||||
cost: Final = StandardLoggingPayload(id=BACKGROUND_RESPONSE_ID, call_type="aget_responses", response_cost=0.42)
|
||||
ordinary: Final = StandardLoggingPayload(id="chatcmpl-ordinary", call_type="acompletion")
|
||||
batch_retrieve: Final = StandardLoggingPayload(id="batch_abc123", call_type="aretrieve_batch", response_cost=0.07)
|
||||
return (
|
||||
({"call_type": "aresponses", "litellm_params": {"metadata": {}}, "standard_logging_object": create}, create),
|
||||
({**cost_kwargs, "standard_logging_object": cost}, cost),
|
||||
(
|
||||
{"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": ordinary},
|
||||
ordinary,
|
||||
),
|
||||
(
|
||||
{
|
||||
"call_type": "aretrieve_batch",
|
||||
"litellm_params": {"metadata": {}},
|
||||
"standard_logging_object": batch_retrieve,
|
||||
},
|
||||
batch_retrieve,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _expected_files(
|
||||
events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...],
|
||||
) -> dict[str, StandardLoggingPayload]:
|
||||
create, cost, ordinary, batch_retrieve = (payload for _, payload in events)
|
||||
return {
|
||||
"resp_YWJjZA.json": create,
|
||||
"resp_YWJjZA_background_cost.json": cost,
|
||||
"chatcmpl-ordinary.json": ordinary,
|
||||
"batch_abc123.json": batch_retrieve,
|
||||
}
|
||||
|
||||
|
||||
async def _log_and_flush(
|
||||
logger: AzureBlobStorageLogger, events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...]
|
||||
) -> None:
|
||||
for kwargs, payload in events:
|
||||
await logger.async_log_success_event(kwargs, {"id": payload["id"]}, None, None)
|
||||
await logger.async_send_batch()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"cost_kwargs",
|
||||
(
|
||||
pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"),
|
||||
pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"),
|
||||
),
|
||||
)
|
||||
async def test_background_cost_event_keeps_its_own_file_beside_the_create_on_the_entra_path(
|
||||
mock_env_vars, monkeypatch, cost_kwargs
|
||||
):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
events: Final = _background_create_cost_and_other_events(cost_kwargs)
|
||||
with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests
|
||||
"litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
mock_http_client = AsyncMock()
|
||||
mock_http_client.put.return_value = MagicMock()
|
||||
mock_http_client.patch.return_value = MagicMock()
|
||||
mock_get_client.return_value = mock_http_client
|
||||
|
||||
logger = AzureBlobStorageLogger()
|
||||
logger.azure_auth_token = "mock-azure-ad-token"
|
||||
logger.token_expiry = None
|
||||
await _log_and_flush(logger, events)
|
||||
|
||||
base_url: Final = "https://test-account.dfs.core.windows.net/test-container/"
|
||||
appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]]
|
||||
uploaded: Final = {
|
||||
call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends
|
||||
}
|
||||
assert uploaded == _expected_files(events), (
|
||||
"the create, its background cost event and every other event must each keep their own file"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_event_is_uploaded_to_a_file_named_after_its_id(mock_env_vars, monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
failure: Final = StandardLoggingPayload(id="chatcmpl-failed", call_type="acompletion", error_str="provider failure")
|
||||
with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests
|
||||
"litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
mock_http_client = AsyncMock()
|
||||
mock_http_client.put.return_value = MagicMock()
|
||||
mock_http_client.patch.return_value = MagicMock()
|
||||
mock_get_client.return_value = mock_http_client
|
||||
|
||||
logger = AzureBlobStorageLogger()
|
||||
logger.azure_auth_token = "mock-azure-ad-token"
|
||||
logger.token_expiry = None
|
||||
await logger.async_log_failure_event(
|
||||
{"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": failure},
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
await logger.async_send_batch()
|
||||
|
||||
base_url: Final = "https://test-account.dfs.core.windows.net/test-container/"
|
||||
appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]]
|
||||
assert {call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends} == {
|
||||
"chatcmpl-failed.json": failure
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_cost_event_keeps_its_own_file_beside_the_create_with_an_account_key(
|
||||
workload_identity_env_vars, monkeypatch
|
||||
):
|
||||
monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime)
|
||||
events: Final = _background_create_cost_and_other_events(BACKGROUND_COST_POLL_KWARGS)
|
||||
|
||||
uploaded: Final[dict[str, StandardLoggingPayload]] = {}
|
||||
|
||||
def file_client_for(file_path: str) -> MagicMock:
|
||||
file_client = MagicMock()
|
||||
file_client.create_file = AsyncMock()
|
||||
file_client.flush_data = AsyncMock()
|
||||
|
||||
async def append_data(data: bytes, offset: int, length: int) -> None:
|
||||
uploaded[file_path] = json.loads(data)
|
||||
|
||||
file_client.append_data = append_data
|
||||
return file_client
|
||||
|
||||
file_system_client = MagicMock()
|
||||
file_system_client.get_file_client = MagicMock(side_effect=file_client_for)
|
||||
service_client = MagicMock()
|
||||
service_client.get_file_system_client = MagicMock(return_value=file_system_client)
|
||||
fake_aio_module = MagicMock()
|
||||
fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client)
|
||||
|
||||
with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
|
||||
logger = AzureBlobStorageLogger()
|
||||
await _log_and_flush(logger, events)
|
||||
|
||||
today: Final = _FrozenDatetime.now().strftime("%Y-%m-%d")
|
||||
assert uploaded == {f"{today}/{name}": payload for name, payload in _expected_files(events).items()}, (
|
||||
"the create, its background cost event and every other event must each keep their own file"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -20,10 +20,16 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import (
|
|||
strip_cache_hit_suffix,
|
||||
)
|
||||
from litellm.integrations.clickhouse.context import lens_analysis
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.litellm_core_utils.secret_redaction import REDACTED
|
||||
from litellm.tracing.types import SpendLogRecord
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, StandardLoggingPayload
|
||||
|
||||
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
|
||||
|
|
@ -349,6 +355,41 @@ def test_cache_hit_id_is_stripped_for_response_id():
|
|||
assert strip_cache_hit_suffix("chatcmpl-xyz") == "chatcmpl-xyz"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_kwargs",
|
||||
(
|
||||
pytest.param(
|
||||
{
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {
|
||||
"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}
|
||||
},
|
||||
},
|
||||
id="cost-poll",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"call_type": "pass_through_endpoint",
|
||||
"litellm_params": {"metadata": {}},
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES,
|
||||
},
|
||||
id="passthrough-settlement",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_background_cost_row_keeps_its_own_request_id_beside_its_create(cost_kwargs: Mapping[str, object]) -> None:
|
||||
payload: Final = cast(StandardLoggingPayload, _payload(id="resp_YWJjZA", call_type="aresponses"))
|
||||
create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}}
|
||||
|
||||
rows: Final = [spend_log_row_from_payload(payload, kwargs) for kwargs in (create_kwargs, cost_kwargs)]
|
||||
|
||||
assert [row["request_id"] for row in rows] == [
|
||||
"resp_YWJjZA",
|
||||
f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}",
|
||||
]
|
||||
assert [row["response_id"] for row in rows] == ["resp_YWJjZA", "resp_YWJjZA"]
|
||||
|
||||
|
||||
def test_parse_traceparent_valid_missing_malformed():
|
||||
assert parse_traceparent(TRACEPARENT) == (TRACE_ID, SPAN_ID)
|
||||
assert parse_traceparent(None) == ("", "")
|
||||
|
|
|
|||
|
|
@ -1,14 +1,18 @@
|
|||
import functools
|
||||
import http.server
|
||||
import importlib.abc
|
||||
import ipaddress
|
||||
import os
|
||||
import queue
|
||||
import ssl
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from importlib.machinery import ModuleSpec
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -75,6 +79,36 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]:
|
|||
return certificate_path, key_path
|
||||
|
||||
|
||||
PROXY_ONLY_PACKAGES: Final = frozenset({"fastapi", "starlette"})
|
||||
|
||||
|
||||
class _ProxyOnlyPackageBlocker(importlib.abc.MetaPathFinder):
|
||||
def find_spec(
|
||||
self,
|
||||
fullname: str,
|
||||
path: Sequence[str] | None = None,
|
||||
target: ModuleType | None = None,
|
||||
) -> ModuleSpec | None:
|
||||
if fullname.partition(".")[0] in PROXY_ONLY_PACKAGES:
|
||||
raise ModuleNotFoundError(f"No module named {fullname!r}")
|
||||
return None
|
||||
|
||||
|
||||
def _needs_proxy_extras(module_name: str) -> bool:
|
||||
return (
|
||||
module_name.partition(".")[0] in PROXY_ONLY_PACKAGES
|
||||
or module_name == "litellm.proxy"
|
||||
or module_name.startswith("litellm.proxy.")
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sdk_only_install(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for module_name in tuple(filter(_needs_proxy_extras, sys.modules)):
|
||||
monkeypatch.delitem(sys.modules, module_name)
|
||||
monkeypatch.setattr(sys, "meta_path", [_ProxyOnlyPackageBlocker(), *sys.meta_path])
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_process_environment() -> Iterator[None]:
|
||||
original: Final = dict(os.environ)
|
||||
|
|
|
|||
|
|
@ -13,9 +13,11 @@ import litellm
|
|||
import pytest
|
||||
import respx
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
|
||||
mock_response_data: Final = {
|
||||
|
|
@ -172,9 +174,7 @@ async def test_get_payload_not_found(monkeypatch):
|
|||
assert payload is None
|
||||
|
||||
|
||||
def _gcs_logger_with_upload_status(
|
||||
monkeypatch: pytest.MonkeyPatch, status_code: int
|
||||
) -> tuple[GCSBucketLogger, Mock]:
|
||||
def _gcs_logger_with_upload_status(monkeypatch: pytest.MonkeyPatch, status_code: int) -> tuple[GCSBucketLogger, Mock]:
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr(google.auth, "default", _google_default_credentials)
|
||||
monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600")
|
||||
|
|
@ -394,3 +394,98 @@ async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: p
|
|||
|
||||
(thread,) = threads
|
||||
assert thread is not threading.main_thread()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unbatched_background_cost_event_keeps_its_own_object_readable_by_its_spend_log_id(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false")
|
||||
logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200)
|
||||
object_prefix: Final = "/storage/v1/b/test-bucket/o/"
|
||||
stored: Final[dict[str, object]] = {}
|
||||
|
||||
def bucket(request: httpx.Request) -> httpx.Response:
|
||||
if request.method == "POST":
|
||||
stored[request.url.params["name"]] = json.loads(request.content)
|
||||
return httpx.Response(200, json={"name": request.url.params["name"]})
|
||||
name: Final = request.url.path.removeprefix(object_prefix)
|
||||
if name not in stored:
|
||||
return httpx.Response(404, text="No such object")
|
||||
return httpx.Response(200, content=json.dumps(stored[name]).encode("utf-8"))
|
||||
|
||||
logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket))
|
||||
response_id: Final = "resp_YmFja2dyb3VuZA"
|
||||
create: Final = {
|
||||
"id": response_id,
|
||||
"messages": [{"role": "user", "content": "background create input"}],
|
||||
"error_str": None,
|
||||
}
|
||||
cost: Final = {"id": response_id, "response_cost": 0.42, "error_str": None}
|
||||
ordinary: Final = {"id": "chatcmpl-ordinary", "error_str": None}
|
||||
batch_retrieve: Final = {"id": "batch_abc123", "response_cost": 0.07, "error_str": None}
|
||||
events: Final = (
|
||||
({"call_type": "aresponses", "litellm_params": {"metadata": {}}}, create),
|
||||
(
|
||||
{
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {
|
||||
"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}
|
||||
},
|
||||
},
|
||||
cost,
|
||||
),
|
||||
({"call_type": "acompletion", "litellm_params": {"metadata": {}}}, ordinary),
|
||||
({"call_type": "aretrieve_batch", "litellm_params": {"metadata": {}}}, batch_retrieve),
|
||||
)
|
||||
|
||||
for kwargs, payload in events:
|
||||
await logger.async_log_success_event(
|
||||
{**kwargs, "standard_logging_object": payload}, {"id": payload["id"]}, None, None
|
||||
)
|
||||
await logger.flush_queue()
|
||||
|
||||
assert {name.split("/", maxsplit=1)[1]: payload for name, payload in stored.items()} == {
|
||||
response_id: create,
|
||||
f"{response_id}_background_cost": cost,
|
||||
"chatcmpl-ordinary": ordinary,
|
||||
"batch_abc123": batch_retrieve,
|
||||
}, "the create, its background cost event and every other event must each keep their own object"
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
assert await logger.get_request_response_payload(f"{response_id}_background_cost", now, None) == cost
|
||||
assert await logger.get_request_response_payload(response_id, now, None) == create
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unbatched_gcs_log_id_still_names_a_background_cost_event_object(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false")
|
||||
logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200)
|
||||
stored: Final[dict[str, object]] = {}
|
||||
|
||||
def bucket(request: httpx.Request) -> httpx.Response:
|
||||
stored[request.url.params["name"]] = json.loads(request.content)
|
||||
return httpx.Response(200, json={"name": request.url.params["name"]})
|
||||
|
||||
logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket))
|
||||
cost: Final = {"id": "resp_YmFja2dyb3VuZA", "response_cost": 0.42, "error_str": None}
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
"gcs_log_id": "qa-run-7",
|
||||
}
|
||||
},
|
||||
"standard_logging_object": cost,
|
||||
},
|
||||
{"id": cost["id"]},
|
||||
None,
|
||||
None,
|
||||
)
|
||||
await logger.flush_queue()
|
||||
|
||||
((name, payload),) = stored.items()
|
||||
object_name: Final = name.split("/", maxsplit=1)[1]
|
||||
assert object_name.startswith("custom-") and object_name.endswith("-qa-run-7"), object_name
|
||||
assert payload == cost
|
||||
|
|
|
|||
|
|
@ -2,11 +2,15 @@ import os
|
|||
import unittest
|
||||
from litellm._uuid import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from litellm.integrations.deepeval.api import Endpoints, HttpMethods
|
||||
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
|
||||
from litellm.integrations.deepeval.types import SpanApiType, TraceSpanApiStatus
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
|
||||
class TestDeepEvalLogger(unittest.TestCase):
|
||||
|
|
@ -94,6 +98,37 @@ class TestDeepEvalLogger(unittest.TestCase):
|
|||
self.assertEqual(llm_span["inputTokenCount"], 10)
|
||||
self.assertEqual(llm_span["outputTokenCount"], 20)
|
||||
|
||||
def test_background_cost_span_keeps_its_own_uuid_beside_its_create(self) -> None:
|
||||
standard_logging_object: Final = {
|
||||
"id": "resp_YWJjZA",
|
||||
"trace_id": self.trace_id,
|
||||
"model": self.model,
|
||||
"response": {"usage": {}, "choices": [{"message": {"content": "OK"}}]},
|
||||
}
|
||||
create_kwargs: Final = {
|
||||
"input": self.input_str,
|
||||
"call_type": "aresponses",
|
||||
"litellm_params": {"metadata": {}},
|
||||
"standard_logging_object": standard_logging_object,
|
||||
}
|
||||
cost_kwargs: Final = {
|
||||
"input": self.input_str,
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {
|
||||
"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}
|
||||
},
|
||||
"standard_logging_object": standard_logging_object,
|
||||
}
|
||||
|
||||
for kwargs in (create_kwargs, cost_kwargs):
|
||||
self.logger.log_success_event(kwargs, self.mock_response_obj, self.start_time, self.end_time)
|
||||
|
||||
span_uuids: Final = [
|
||||
call.kwargs["body"]["llmSpans"][0]["uuid"]
|
||||
for call in self.mock_api_instance.send_request.call_args_list
|
||||
]
|
||||
self.assertEqual(span_uuids, ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"])
|
||||
|
||||
def test_log_failure_event(self):
|
||||
error_message = "This is an error."
|
||||
kwargs = {
|
||||
|
|
|
|||
52
tests/unit/integrations/test_dynamodb.py
Normal file
52
tests/unit/integrations/test_dynamodb.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.dynamodb import DyanmoDBLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = {
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}},
|
||||
}
|
||||
PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = {
|
||||
"call_type": "pass_through_endpoint",
|
||||
"litellm_params": {"metadata": {}},
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_kwargs",
|
||||
(
|
||||
pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"),
|
||||
pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.usefixtures("sdk_only_install")
|
||||
def test_background_cost_event_keeps_its_own_item_beside_its_create(
|
||||
monkeypatch: pytest.MonkeyPatch, cost_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "dynamodb_table_name", "litellm-logs")
|
||||
dynamodb: Final = MagicMock()
|
||||
logger: Final = DyanmoDBLogger(dynamodb=dynamodb)
|
||||
response_obj: Final = {"id": "resp_YWJjZA", "usage": {"total_tokens": 3}}
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}}
|
||||
|
||||
for kwargs in (create_kwargs, cost_kwargs):
|
||||
logger.log_event(kwargs, response_obj, now, now, print_verbose=print)
|
||||
|
||||
dynamodb.Table.assert_called_with("litellm-logs")
|
||||
written_ids: Final = [call.kwargs["Item"]["id"] for call in dynamodb.Table.return_value.put_item.call_args_list]
|
||||
assert written_ids == ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"]
|
||||
|
|
@ -1,11 +1,30 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.openmeter import OpenMeterLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG,
|
||||
BackgroundResponseSettlement,
|
||||
)
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = {
|
||||
"call_type": "aget_responses",
|
||||
"litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}},
|
||||
}
|
||||
PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = {
|
||||
"call_type": "pass_through_endpoint",
|
||||
"litellm_params": {"metadata": {}},
|
||||
SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES,
|
||||
}
|
||||
|
||||
|
||||
class TestOpenMeterIntegration:
|
||||
|
|
@ -36,6 +55,31 @@ class TestOpenMeterIntegration:
|
|||
with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"):
|
||||
OpenMeterLogger()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_kwargs",
|
||||
(
|
||||
pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"),
|
||||
pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.usefixtures("sdk_only_install")
|
||||
def test_background_cost_event_gets_its_own_event_id_beside_its_create(
|
||||
self, cost_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
logger: Final = OpenMeterLogger()
|
||||
response_obj: Final = {"id": "resp_YWJjZA"}
|
||||
create_kwargs: Final[Mapping[str, object]] = {
|
||||
"user": "user-1",
|
||||
"call_type": "aresponses",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
create_event: Final = logger._common_logic(create_kwargs, response_obj)
|
||||
cost_event: Final = logger._common_logic({**cost_kwargs, "user": "user-1", "response_cost": 0.42}, response_obj)
|
||||
|
||||
assert create_event["id"] == "resp_YWJjZA"
|
||||
assert cost_event["id"] == f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"
|
||||
|
||||
def test_common_logic_with_string_user(self):
|
||||
"""Test that _common_logic correctly handles string user parameter"""
|
||||
logger = OpenMeterLogger()
|
||||
|
|
|
|||
|
|
@ -11090,6 +11090,40 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost()
|
|||
assert logging_obj.response_cost_calculator(result=assembled) == 0.0042
|
||||
|
||||
|
||||
@pytest.mark.parametrize("continued_response_id", ["resp_continuedafteracutstream", None])
|
||||
@pytest.mark.parametrize("client_facing_response_id", ["resp_encryptedidtheclientsaw", None])
|
||||
def test_sync_success_handler_logs_a_stream_under_the_id_the_client_saw(
|
||||
client_facing_response_id, continued_response_id
|
||||
):
|
||||
class Recorder(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.logged = []
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.logged.append((response_obj.id, kwargs["standard_logging_object"]["id"]))
|
||||
|
||||
recorder: Final = Recorder()
|
||||
logging_obj: Final = _responses_stream_logging_obj()
|
||||
logging_obj.dynamic_success_callbacks = [recorder]
|
||||
logging_obj.client_facing_response_id = client_facing_response_id
|
||||
if continued_response_id is not None:
|
||||
logging_obj.record_continued_partial(
|
||||
Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), 0.0, continued_response_id
|
||||
)
|
||||
now: Final = datetime.datetime.now()
|
||||
|
||||
logging_obj.success_handler(
|
||||
result=_completed_responses_event(ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
cache_hit=False,
|
||||
)
|
||||
|
||||
expected_id: Final = client_facing_response_id or continued_response_id or "resp-1"
|
||||
assert recorder.logged == [(expected_id, expected_id)]
|
||||
|
||||
|
||||
def test_response_cost_calculator_prices_terminal_responses_event_from_its_response():
|
||||
logging_obj: Final = _responses_stream_logging_obj()
|
||||
inner_response: Final = ResponsesAPIResponse(
|
||||
|
|
|
|||
|
|
@ -5150,6 +5150,23 @@ def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_th
|
|||
assert request_ids == ("resp_shared", "resp_shared_background_cost")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"])
|
||||
def test_a_read_or_cancel_row_carries_the_id_its_callbacks_logged(call_type: str):
|
||||
row: Final = get_logging_payload(
|
||||
kwargs={
|
||||
"call_type": call_type,
|
||||
"model": "gpt-5.6",
|
||||
"litellm_call_id": "call-read",
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
|
||||
},
|
||||
response_obj={"id": "resp_logged_by_callbacks", "object": "response", "status": "completed"},
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert row["request_id"] == "resp_logged_by_callbacks"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("forged_marker", [True, BackgroundResponseSettlement.SETTLES.value])
|
||||
def test_a_request_body_cannot_move_its_spend_row_onto_a_background_settlement_id(forged_marker):
|
||||
"""A failed call logs its raw request body as kwargs, so a body field named like the settlement
|
||||
|
|
|
|||
|
|
@ -1083,6 +1083,128 @@ async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored
|
|||
]
|
||||
|
||||
|
||||
def _stored_response_turn(
|
||||
request_id: str, call_type: str, status: str, output: list[dict[str, object]], prompt: str | None = None
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"call_type": call_type,
|
||||
"session_id": "session-with-reads",
|
||||
"proxy_server_request": json.dumps({} if prompt is None else {"input": [{"role": "user", "content": prompt}]}),
|
||||
"response": {"id": request_id, "object": "response", "status": status, "output": output},
|
||||
}
|
||||
|
||||
|
||||
def _identified_output_message(item_id: str, text: str) -> dict[str, object]:
|
||||
return {**_output_message(text), "id": item_id}
|
||||
|
||||
|
||||
def _replayed_roles_and_texts(session: ChatCompletionSession) -> list[tuple[object, object]]:
|
||||
return [
|
||||
(role, _text_of(content)) for role, content in (_role_and_content(message) for message in session["messages"])
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reads_and_cancels_of_a_response_the_session_already_holds_replay_no_turns():
|
||||
answer: Final = [_identified_output_message("msg_answer", "OK")]
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Reply with OK"),
|
||||
_stored_response_turn("resp_read", "aget_responses", "completed", answer),
|
||||
_stored_response_turn("resp_streamed_read", "aget_responses", "completed", answer),
|
||||
_stored_response_turn("resp_cancel", "acancel_responses", "completed", answer),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_create")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [("user", "Reply with OK"), ("assistant", "OK")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_polled_background_response_replays_its_finished_output_once():
|
||||
answer: Final = [_identified_output_message("msg_answer", "One, two, three")]
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"),
|
||||
_stored_response_turn("resp_poll_running", "aget_responses", "in_progress", []),
|
||||
_stored_response_turn("resp_poll_done", "aget_responses", "completed", answer),
|
||||
_stored_response_turn("resp_poll_again", "aget_responses", "completed", answer),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_queued")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [("user", "Count to three"), ("assistant", "One, two, three")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_background_response_that_grows_between_polls_replays_each_output_item_once():
|
||||
first: Final = _identified_output_message("msg_first", "One")
|
||||
second: Final = _identified_output_message("msg_second", "Two")
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"),
|
||||
_stored_response_turn("resp_poll_partial", "aget_responses", "in_progress", [first]),
|
||||
_stored_response_turn("resp_poll_done", "aget_responses", "completed", [first, second]),
|
||||
_stored_response_turn("resp_poll_again", "aget_responses", "completed", [first, second]),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_queued")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_item_a_poll_caught_half_written_replays_its_finished_text_once():
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"),
|
||||
_stored_response_turn(
|
||||
"resp_poll_partial", "aget_responses", "in_progress", [_identified_output_message("msg_answer", "One, t")]
|
||||
),
|
||||
_stored_response_turn(
|
||||
"resp_poll_done",
|
||||
"aget_responses",
|
||||
"completed",
|
||||
[_identified_output_message("msg_answer", "One, two, three"), _identified_output_message("msg_more", "Done")],
|
||||
),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_queued")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [
|
||||
("user", "Count to three"),
|
||||
("assistant", "One, two, three"),
|
||||
("assistant", "Done"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_polled_output_item_without_an_id_still_replays_in_order():
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"),
|
||||
_stored_response_turn(
|
||||
"resp_poll_done",
|
||||
"aget_responses",
|
||||
"completed",
|
||||
[_output_message("One"), _identified_output_message("msg_two", "Two")],
|
||||
),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_queued")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_read_that_ends_before_its_streamed_create_replays_the_answer_once_after_the_prompt():
|
||||
answer: Final = [_identified_output_message("msg_answer", "OK")]
|
||||
spend_logs: Final = (
|
||||
_stored_response_turn("resp_read_mid_stream", "aget_responses", "completed", answer),
|
||||
_stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Say OK"),
|
||||
)
|
||||
|
||||
session: Final = await _history_from_stored_turns(spend_logs, "resp_create")
|
||||
|
||||
assert _replayed_roles_and_texts(session) == [("user", "Say OK"), ("assistant", "OK")]
|
||||
|
||||
|
||||
def _tool_call_ids(message: object) -> tuple[str, ...]:
|
||||
tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None)
|
||||
return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ())
|
||||
|
|
|
|||
|
|
@ -4,23 +4,33 @@ Tests for ResponsesIDSecurity hook.
|
|||
Tests the security hook that prevents user B from seeing response from user A.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
from collections.abc import Callable, Mapping
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.responses_id_security import (
|
||||
ResponsesIDSecurity,
|
||||
_streams_response_ids,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
GenericEvent,
|
||||
ResponseAPIUsage,
|
||||
ResponseCompletedEvent,
|
||||
ResponseCreatedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
from litellm.types.utils import SpecialEnums
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.utils import CallTypesLiteral, Choices, Message, ModelResponse, SpecialEnums, Usage
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -55,12 +65,12 @@ class TestIsEncryptedResponseId:
|
|||
import litellm.proxy.hooks.responses_id_security as responses_module
|
||||
|
||||
with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt:
|
||||
mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456"
|
||||
|
||||
result = responses_id_security._is_encrypted_response_id(
|
||||
"resp_encrypted_value"
|
||||
mock_decrypt.return_value = (
|
||||
f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456"
|
||||
)
|
||||
|
||||
result = responses_id_security._is_encrypted_response_id("resp_encrypted_value")
|
||||
|
||||
assert result is True
|
||||
mock_decrypt.assert_called_once()
|
||||
|
||||
|
|
@ -88,9 +98,7 @@ class TestDecryptResponseId:
|
|||
with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt:
|
||||
mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_original_123;user_id:user-456;team_id:team-789"
|
||||
|
||||
original_id, user_id, team_id = responses_id_security._decrypt_response_id(
|
||||
"resp_encrypted_value"
|
||||
)
|
||||
original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_encrypted_value")
|
||||
|
||||
assert original_id == "resp_original_123"
|
||||
assert user_id == "user-456"
|
||||
|
|
@ -104,9 +112,7 @@ class TestDecryptResponseId:
|
|||
with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt:
|
||||
mock_decrypt.return_value = None
|
||||
|
||||
original_id, user_id, team_id = responses_id_security._decrypt_response_id(
|
||||
"resp_plain_value"
|
||||
)
|
||||
original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_plain_value")
|
||||
|
||||
assert original_id == "resp_plain_value"
|
||||
assert user_id is None
|
||||
|
|
@ -116,9 +122,7 @@ class TestDecryptResponseId:
|
|||
class TestCheckUserAccessToResponseId:
|
||||
"""Test check_user_access_to_response_id function"""
|
||||
|
||||
def test_check_user_access_same_user(
|
||||
self, responses_id_security, mock_user_api_key_dict
|
||||
):
|
||||
def test_check_user_access_same_user(self, responses_id_security, mock_user_api_key_dict):
|
||||
"""Test that same user has access to their response ID"""
|
||||
result = responses_id_security.check_user_access_to_response_id(
|
||||
response_id_user_id="test-user-123",
|
||||
|
|
@ -128,9 +132,7 @@ class TestCheckUserAccessToResponseId:
|
|||
|
||||
assert result is True
|
||||
|
||||
def test_check_user_access_different_user_raises_exception(
|
||||
self, responses_id_security, mock_user_api_key_dict
|
||||
):
|
||||
def test_check_user_access_different_user_raises_exception(self, responses_id_security, mock_user_api_key_dict):
|
||||
"""Test that different user is denied access to response ID"""
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {}):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -143,9 +145,7 @@ class TestCheckUserAccessToResponseId:
|
|||
assert exc_info.value.status_code == 403
|
||||
assert "Forbidden" in exc_info.value.detail
|
||||
|
||||
def test_check_user_access_different_team_raises_exception(
|
||||
self, responses_id_security, mock_user_api_key_dict
|
||||
):
|
||||
def test_check_user_access_different_team_raises_exception(self, responses_id_security, mock_user_api_key_dict):
|
||||
"""Test that different team is denied access to response ID"""
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {}):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -158,9 +158,7 @@ class TestCheckUserAccessToResponseId:
|
|||
assert exc_info.value.status_code == 403
|
||||
assert "Forbidden" in exc_info.value.detail
|
||||
|
||||
def test_check_user_access_team_a_to_team_b_without_user_id(
|
||||
self, responses_id_security
|
||||
):
|
||||
def test_check_user_access_team_a_to_team_b_without_user_id(self, responses_id_security):
|
||||
"""Test that key from team A (without user_id) cannot access response from team B (without user_id)"""
|
||||
# Create a mock user from team A without user_id
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -179,9 +177,7 @@ class TestCheckUserAccessToResponseId:
|
|||
assert exc_info.value.status_code == 403
|
||||
assert "team" in exc_info.value.detail.lower()
|
||||
|
||||
def test_check_user_access_team_a_to_team_b_with_user_id(
|
||||
self, responses_id_security
|
||||
):
|
||||
def test_check_user_access_team_a_to_team_b_with_user_id(self, responses_id_security):
|
||||
"""Test that key from team A (without user_id) cannot access response from team B (with user_id)"""
|
||||
# Create a mock user from team A without user_id
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -217,9 +213,7 @@ class TestCheckUserAccessToResponseId:
|
|||
|
||||
assert result is True
|
||||
|
||||
def test_check_user_access_admin_can_access_any_response(
|
||||
self, responses_id_security
|
||||
):
|
||||
def test_check_user_access_admin_can_access_any_response(self, responses_id_security):
|
||||
"""Test that proxy admin can access any response ID"""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
|
|
@ -238,9 +232,7 @@ class TestCheckUserAccessToResponseId:
|
|||
|
||||
assert result is True
|
||||
|
||||
def test_check_user_access_security_disabled(
|
||||
self, responses_id_security, mock_user_api_key_dict
|
||||
):
|
||||
def test_check_user_access_security_disabled(self, responses_id_security, mock_user_api_key_dict):
|
||||
"""Test that when security is disabled, any user can access any response"""
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
|
|
@ -266,9 +258,7 @@ class TestAsyncPreCallHook:
|
|||
"""Test pre-call hook decrypts previous_response_id for aresponses call"""
|
||||
data = {"previous_response_id": "resp_encrypted_value"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -284,15 +274,11 @@ class TestAsyncPreCallHook:
|
|||
assert result["previous_response_id"] == "resp_original_123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_aget_responses(
|
||||
self, responses_id_security, mock_user_api_key_dict, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_aget_responses(self, responses_id_security, mock_user_api_key_dict, mock_cache):
|
||||
"""Test pre-call hook decrypts response_id for aget_responses call"""
|
||||
data = {"response_id": "resp_encrypted_456"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -308,9 +294,7 @@ class TestAsyncPreCallHook:
|
|||
assert result["response_id"] == "resp_original_456"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_team_a_accessing_team_b_response(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_team_a_accessing_team_b_response(self, responses_id_security, mock_cache):
|
||||
"""Test pre-call hook prevents team A from accessing team B response"""
|
||||
# Create a mock user from team A
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -320,9 +304,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_b"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -341,9 +323,7 @@ class TestAsyncPreCallHook:
|
|||
assert "team" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_team_a_accessing_team_b_with_user(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_team_a_accessing_team_b_with_user(self, responses_id_security, mock_cache):
|
||||
"""Test pre-call hook prevents team A (no user) from accessing team B response (with user)"""
|
||||
# Create a mock user from team A without user_id
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -353,9 +333,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_b_with_user"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -375,9 +353,7 @@ class TestAsyncPreCallHook:
|
|||
assert "forbidden" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_same_team_access(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_same_team_access(self, responses_id_security, mock_cache):
|
||||
"""Test pre-call hook allows team A to access their own team's response"""
|
||||
# Create a mock user from team A
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -387,9 +363,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_a"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -405,9 +379,7 @@ class TestAsyncPreCallHook:
|
|||
assert result["response_id"] == "resp_original_team_a"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_adelete_responses_team_security(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_adelete_responses_team_security(self, responses_id_security, mock_cache):
|
||||
"""Test pre-call hook prevents team A from deleting team B's response"""
|
||||
# Create a mock user from team A
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -417,9 +389,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_b"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -438,9 +408,7 @@ class TestAsyncPreCallHook:
|
|||
assert "team" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_acancel_responses_team_security(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_acancel_responses_team_security(self, responses_id_security, mock_cache):
|
||||
"""Test pre-call hook prevents team A from canceling team B's response"""
|
||||
# Create a mock user from team A
|
||||
mock_auth_team_a = MagicMock()
|
||||
|
|
@ -450,9 +418,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_b"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -470,16 +436,13 @@ class TestAsyncPreCallHook:
|
|||
assert exc_info.value.status_code == 403
|
||||
assert "team" in exc_info.value.detail.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_alist_input_items_decrypts_response_id(
|
||||
self, responses_id_security, mock_user_api_key_dict, mock_cache
|
||||
):
|
||||
data = {"response_id": "resp_encrypted_789"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -496,9 +459,7 @@ class TestAsyncPreCallHook:
|
|||
assert result["response_id"] == "resp_original_789"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_alist_input_items_team_security(
|
||||
self, responses_id_security, mock_cache
|
||||
):
|
||||
async def test_async_pre_call_hook_alist_input_items_team_security(self, responses_id_security, mock_cache):
|
||||
mock_auth_team_a = MagicMock()
|
||||
mock_auth_team_a.user_id = None
|
||||
mock_auth_team_a.team_id = "team-a"
|
||||
|
|
@ -506,9 +467,7 @@ class TestAsyncPreCallHook:
|
|||
|
||||
data = {"response_id": "resp_encrypted_team_b"}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_is_encrypted_response_id", return_value=True
|
||||
):
|
||||
with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True):
|
||||
with patch.object(
|
||||
responses_id_security,
|
||||
"_decrypt_response_id",
|
||||
|
|
@ -620,9 +579,7 @@ class TestAsyncPostCallStreamingIteratorHook:
|
|||
"/openai/v1/responses/resp_rawprovider123",
|
||||
],
|
||||
)
|
||||
async def test_streamed_id_encrypted_on_all_responses_routes(
|
||||
self, responses_id_security, route, monkeypatch
|
||||
):
|
||||
async def test_streamed_id_encrypted_on_all_responses_routes(self, responses_id_security, route, monkeypatch):
|
||||
streamed_id = await self._drain_streamed_id(responses_id_security, route, monkeypatch)
|
||||
|
||||
assert streamed_id != "resp_rawprovider123"
|
||||
|
|
@ -634,12 +591,8 @@ class TestAsyncPostCallStreamingIteratorHook:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_id_untouched_on_non_responses_route(
|
||||
self, responses_id_security, monkeypatch
|
||||
):
|
||||
streamed_id = await self._drain_streamed_id(
|
||||
responses_id_security, "/chat/completions", monkeypatch
|
||||
)
|
||||
async def test_streamed_id_untouched_on_non_responses_route(self, responses_id_security, monkeypatch):
|
||||
streamed_id = await self._drain_streamed_id(responses_id_security, "/chat/completions", monkeypatch)
|
||||
|
||||
assert streamed_id == "resp_rawprovider123"
|
||||
assert not responses_id_security._is_encrypted_response_id(streamed_id)
|
||||
|
|
@ -698,12 +651,8 @@ class TestStreamedGenericEventIdEncryption:
|
|||
def _advertised_ids(events):
|
||||
nested = (getattr(event, "response", None) for event in events)
|
||||
return [
|
||||
payload["id"] if isinstance(payload, dict) else payload.id
|
||||
for payload in nested
|
||||
if payload is not None
|
||||
] + [
|
||||
event.id for event in events if isinstance(getattr(event, "id", None), str)
|
||||
]
|
||||
payload["id"] if isinstance(payload, dict) else payload.id for payload in nested if payload is not None
|
||||
] + [event.id for event in events if isinstance(getattr(event, "id", None), str)]
|
||||
|
||||
async def _drain(self, responses_id_security, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij")
|
||||
|
|
@ -723,9 +672,7 @@ class TestStreamedGenericEventIdEncryption:
|
|||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_every_event_advertises_the_same_encrypted_id(
|
||||
self, responses_id_security, monkeypatch
|
||||
):
|
||||
async def test_every_event_advertises_the_same_encrypted_id(self, responses_id_security, monkeypatch):
|
||||
events = await self._drain(responses_id_security, monkeypatch)
|
||||
advertised = self._advertised_ids(events)
|
||||
|
||||
|
|
@ -742,22 +689,14 @@ class TestStreamedGenericEventIdEncryption:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_provider_id_never_reaches_the_client(
|
||||
self, responses_id_security, monkeypatch
|
||||
):
|
||||
async def test_raw_provider_id_never_reaches_the_client(self, responses_id_security, monkeypatch):
|
||||
events = await self._drain(responses_id_security, monkeypatch)
|
||||
|
||||
assert [self.RAW_ID in event.model_dump_json() for event in events] == [
|
||||
False
|
||||
] * len(events)
|
||||
assert [self.RAW_ID in event.model_dump_json() for event in events] == [False] * len(events)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sibling_fields_survive_the_rewrite(
|
||||
self, responses_id_security, monkeypatch
|
||||
):
|
||||
_, queued, keepalive, later, _ = await self._drain(
|
||||
responses_id_security, monkeypatch
|
||||
)
|
||||
async def test_sibling_fields_survive_the_rewrite(self, responses_id_security, monkeypatch):
|
||||
_, queued, keepalive, later, _ = await self._drain(responses_id_security, monkeypatch)
|
||||
|
||||
assert queued.response["status"] == "queued"
|
||||
assert later.response["status"] == "in_progress"
|
||||
|
|
@ -769,27 +708,19 @@ class TestAsyncPostCallSuccessHook:
|
|||
"""Test async_post_call_success_hook function"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_success_hook_encrypts_response(
|
||||
self, responses_id_security, mock_user_api_key_dict
|
||||
):
|
||||
async def test_async_post_call_success_hook_encrypts_response(self, responses_id_security, mock_user_api_key_dict):
|
||||
"""Test post-call hook encrypts ResponsesAPIResponse"""
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_789", created_at=1234567890, output=[], status="completed"
|
||||
)
|
||||
mock_response = ResponsesAPIResponse(id="resp_789", created_at=1234567890, output=[], status="completed")
|
||||
data = {}
|
||||
|
||||
with patch.object(
|
||||
responses_id_security, "_encrypt_response_id", return_value=mock_response
|
||||
) as mock_encrypt:
|
||||
with patch.object(responses_id_security, "_encrypt_response_id", return_value=mock_response) as mock_encrypt:
|
||||
result = await responses_id_security.async_post_call_success_hook(
|
||||
data=data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
response=mock_response,
|
||||
)
|
||||
|
||||
mock_encrypt.assert_called_once_with(
|
||||
mock_response, mock_user_api_key_dict, request_cache=None
|
||||
)
|
||||
mock_encrypt.assert_called_once_with(mock_response, mock_user_api_key_dict, request_cache=None)
|
||||
assert result == mock_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -846,9 +777,7 @@ def _auth(user_id="owner-user", team_id="owner-team", user_role=None):
|
|||
def _issue_managed_id(hook, owner, provider_response_id=_FABRICATED_PROVIDER_RESPONSE_ID):
|
||||
"""Mint an id exactly the way the proxy hands one to a client on create."""
|
||||
issued = hook._encrypt_response_id(
|
||||
ResponsesAPIResponse(
|
||||
id=provider_response_id, created_at=1234567890, output=[], status="completed"
|
||||
),
|
||||
ResponsesAPIResponse(id=provider_response_id, created_at=1234567890, output=[], status="completed"),
|
||||
owner,
|
||||
)
|
||||
return issued.id
|
||||
|
|
@ -1080,9 +1009,7 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization:
|
|||
assert exc_info.value.status_code == 403
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response(
|
||||
self, mock_cache, salt_key_env
|
||||
):
|
||||
async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response(self, mock_cache, salt_key_env):
|
||||
hook = _hook()
|
||||
stranger = _auth(user_id="stranger-user", team_id="stranger-team")
|
||||
stranger_id = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc")
|
||||
|
|
@ -1098,3 +1025,436 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization:
|
|||
|
||||
assert result["response_id"] == "resp_strangerownprovideridcccccccc"
|
||||
assert result["response_id"] != victim_provider_id
|
||||
|
||||
|
||||
def _admin():
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
return _auth(user_id="admin-user", team_id="admin-team", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
|
||||
def _provider_answer(response_id: str) -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(id=response_id, created_at=1234567890, output=[], status="completed")
|
||||
|
||||
|
||||
class TestReadsAnswerWithTheIdTheClientHolds:
|
||||
"""A follow-up looks its history up by the id the client holds, so a read must not hand out a new one."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"])
|
||||
@pytest.mark.parametrize(
|
||||
"issuer, reader",
|
||||
[
|
||||
(_auth(), _auth()),
|
||||
(_auth(team_id=None), _auth(team_id="another-team")),
|
||||
(_auth(user_id=None), _auth(user_id=None)),
|
||||
(_auth(user_id=None), _auth(user_id="teammate-user")),
|
||||
(_auth(), _admin()),
|
||||
(_auth(user_id=None), _auth(user_id=None, team_id="admin-team", user_role=_admin().user_role)),
|
||||
],
|
||||
ids=[
|
||||
"same-key",
|
||||
"same-user-new-team",
|
||||
"same-team-no-user",
|
||||
"teammate-reading-a-team-key-id",
|
||||
"admin-user",
|
||||
"admin-without-user",
|
||||
],
|
||||
)
|
||||
async def test_a_caller_with_access_gets_back_the_exact_id_they_sent(
|
||||
self, mock_cache, salt_key_env, call_type, issuer, reader
|
||||
):
|
||||
hook = _hook()
|
||||
issued_id = _issue_managed_id(hook, issuer)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type=call_type
|
||||
)
|
||||
|
||||
result = await hook.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=reader, response=_provider_answer(data["response_id"])
|
||||
)
|
||||
|
||||
assert result.id == issued_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_follow_up_gets_an_id_for_its_own_response_never_the_previous_one(self, mock_cache, salt_key_env):
|
||||
hook = _hook()
|
||||
owner = _auth()
|
||||
previous_id = _issue_managed_id(hook, owner)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=owner,
|
||||
cache=mock_cache,
|
||||
data={"model": "gpt-fake", "previous_response_id": previous_id},
|
||||
call_type="aresponses",
|
||||
)
|
||||
|
||||
result = await hook.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=owner, response=_provider_answer("resp_followupprovideridzzzzzzzzzzzzzzz")
|
||||
)
|
||||
|
||||
assert result.id != previous_id
|
||||
assert hook._decrypt_response_id(result.id) == (
|
||||
"resp_followupprovideridzzzzzzzzzzzzzzz",
|
||||
"owner-user",
|
||||
"owner-team",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_admin_read_by_a_raw_provider_id_gets_a_managed_id_back(self, mock_cache, salt_key_env):
|
||||
hook = _hook()
|
||||
admin = _admin()
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=admin,
|
||||
cache=mock_cache,
|
||||
data={"response_id": _FABRICATED_UNMANAGED_ID},
|
||||
call_type="aget_responses",
|
||||
)
|
||||
|
||||
result = await hook.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=admin, response=_provider_answer(_FABRICATED_UNMANAGED_ID)
|
||||
)
|
||||
|
||||
assert hook._is_encrypted_response_id(result.id)
|
||||
assert hook._decrypt_response_id(result.id) == (_FABRICATED_UNMANAGED_ID, "admin-user", "admin-team")
|
||||
|
||||
|
||||
_LOGGED_MANAGED_ID = "resp_managedidlitellmbuiltbeforethehookran"
|
||||
|
||||
|
||||
def _streamed_response(response_id):
|
||||
return ResponsesAPIResponse(
|
||||
id=response_id,
|
||||
created_at=0,
|
||||
model="gpt-5.1",
|
||||
object="response",
|
||||
output=[],
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
usage=ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5),
|
||||
)
|
||||
|
||||
|
||||
def _upstream_frames():
|
||||
return (
|
||||
ResponseCreatedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=_streamed_response(_LOGGED_MANAGED_ID)
|
||||
),
|
||||
ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _native_logged_result():
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID)
|
||||
)
|
||||
|
||||
|
||||
def _bridged_logged_result():
|
||||
return ModelResponse(
|
||||
id="chatcmpl-bridgedchatcompletionid",
|
||||
model="gpt-5.1",
|
||||
choices=[Choices(message=Message(role="assistant", content="hi"), finish_reason="stop", index=0)],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
|
||||
)
|
||||
|
||||
|
||||
class _SpendRowRecorder(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.spend_row_ids = []
|
||||
self.standard_logging_ids = []
|
||||
self.callback_response_ids = []
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.spend_row_ids.append(get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"])
|
||||
self.callback_response_ids.append(response_obj.id)
|
||||
if "standard_logging_object" in kwargs:
|
||||
self.standard_logging_ids.append(kwargs["standard_logging_object"]["id"])
|
||||
|
||||
|
||||
async def _stream_logged_mid_stream(
|
||||
hook: ResponsesIDSecurity,
|
||||
route: str,
|
||||
logged_result: ResponseCompletedEvent | ModelResponse,
|
||||
caller: UserAPIKeyAuth | None = None,
|
||||
request_data: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, _SpendRowRecorder]:
|
||||
streaming_caller = caller if caller is not None else _auth()
|
||||
recorder = _SpendRowRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-5.1",
|
||||
messages=[],
|
||||
stream=True,
|
||||
call_type="aresponses",
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id="streamed-responses-call",
|
||||
function_id="streamed-responses-fn",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""}
|
||||
)
|
||||
upstream = _upstream_frames()
|
||||
|
||||
async def frames():
|
||||
for frame in upstream:
|
||||
yield frame
|
||||
|
||||
client_stream = hook.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=streaming_caller.model_copy(update={"request_route": route}),
|
||||
response=frames(),
|
||||
request_data={**(request_data or {}), "litellm_logging_obj": logging_obj},
|
||||
)
|
||||
first_frame = await anext(client_stream)
|
||||
now = datetime.datetime.now()
|
||||
await logging_obj.async_success_handler(result=logged_result, start_time=now, end_time=now, cache_hit=False)
|
||||
client_frames = [first_frame, *[frame async for frame in client_stream]]
|
||||
return client_frames[-1].response.id, recorder
|
||||
|
||||
|
||||
class TestStreamedSpendRowCarriesTheIdTheClientSaw:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("logged_result", [_native_logged_result, _bridged_logged_result], ids=["native", "bridge"])
|
||||
@pytest.mark.parametrize("route", ["/v1/responses", "/responses", "/openai/v1/responses"])
|
||||
async def test_spend_row_and_callbacks_log_the_encrypted_id_of_the_terminal_frame(
|
||||
self, salt_key_env, route, logged_result
|
||||
):
|
||||
hook = _hook()
|
||||
|
||||
client_id, recorder = await _stream_logged_mid_stream(hook, route, logged_result())
|
||||
|
||||
assert hook._decrypt_response_id(client_id) == (_LOGGED_MANAGED_ID, "owner-user", "owner-team")
|
||||
assert recorder.spend_row_ids == [client_id]
|
||||
assert recorder.standard_logging_ids == [client_id]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"hook, route",
|
||||
[
|
||||
(_hook(), "/chat/completions"),
|
||||
(_hook(general_settings={"disable_responses_id_security": True}), "/v1/responses"),
|
||||
(_hook(signing_key=None), "/v1/responses"),
|
||||
],
|
||||
ids=["not-a-create-route", "security-disabled", "no-signing-key"],
|
||||
)
|
||||
async def test_stream_the_hook_leaves_alone_keeps_its_logged_id(self, salt_key_env, hook, route):
|
||||
client_id, recorder = await _stream_logged_mid_stream(hook, route, _native_logged_result())
|
||||
|
||||
assert client_id == _LOGGED_MANAGED_ID
|
||||
assert recorder.spend_row_ids == [_LOGGED_MANAGED_ID]
|
||||
assert recorder.standard_logging_ids == [_LOGGED_MANAGED_ID]
|
||||
|
||||
|
||||
class TestStreamedReadAnswersWithTheIdTheClientHolds:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route_prefix", ["/v1/responses", "/responses", "/openai/v1/responses"])
|
||||
async def test_the_owner_sees_their_id_while_callbacks_log_one_of_the_reads_own(
|
||||
self, mock_cache, salt_key_env, route_prefix
|
||||
):
|
||||
hook = _hook()
|
||||
owner = _auth()
|
||||
issued_id = _issue_managed_id(hook, owner, provider_response_id=_LOGGED_MANAGED_ID)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses"
|
||||
)
|
||||
|
||||
client_id, recorder = await _stream_logged_mid_stream(
|
||||
hook, f"{route_prefix}/{issued_id}", _native_logged_result(), caller=owner, request_data=data
|
||||
)
|
||||
|
||||
assert client_id == issued_id
|
||||
assert recorder.spend_row_ids == recorder.standard_logging_ids
|
||||
assert len(recorder.standard_logging_ids) == 1
|
||||
assert recorder.standard_logging_ids[0] != issued_id
|
||||
assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == (
|
||||
_LOGGED_MANAGED_ID,
|
||||
"owner-user",
|
||||
"owner-team",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_admin_streamed_read_of_someone_elses_response_sees_the_id_sent_and_logs_one_of_its_own(
|
||||
self, mock_cache, salt_key_env
|
||||
):
|
||||
hook = _hook()
|
||||
admin = _admin()
|
||||
issued_id = _issue_managed_id(hook, _auth(), provider_response_id=_LOGGED_MANAGED_ID)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=admin, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses"
|
||||
)
|
||||
|
||||
client_id, recorder = await _stream_logged_mid_stream(
|
||||
hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=admin, request_data=data
|
||||
)
|
||||
|
||||
assert client_id == issued_id
|
||||
assert recorder.spend_row_ids == recorder.standard_logging_ids
|
||||
assert len(recorder.standard_logging_ids) == 1
|
||||
assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == (
|
||||
_LOGGED_MANAGED_ID,
|
||||
"admin-user",
|
||||
"admin-team",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_frame_naming_the_held_response_under_another_wrapping_still_gets_the_id_sent(
|
||||
self, mock_cache, salt_key_env
|
||||
):
|
||||
hook = _hook()
|
||||
owner = _auth()
|
||||
issued_id = _issue_managed_id(hook, owner, provider_response_id=_routed_response_id(_LOGGED_MANAGED_ID))
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses"
|
||||
)
|
||||
|
||||
client_id, _ = await _stream_logged_mid_stream(
|
||||
hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=owner, request_data=data
|
||||
)
|
||||
|
||||
assert client_id == issued_id
|
||||
|
||||
|
||||
async def _read_logged_by_callbacks(
|
||||
hook: ResponsesIDSecurity,
|
||||
caller: UserAPIKeyAuth,
|
||||
issued_id: str,
|
||||
call_type: CallTypesLiteral,
|
||||
cache: DualCache,
|
||||
answer: Callable[[str], ResponsesAPIResponse | DeleteResponseResult] = _provider_answer,
|
||||
) -> tuple[str, _SpendRowRecorder]:
|
||||
recorder = _SpendRowRecorder()
|
||||
now = datetime.datetime.now()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-5.1",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type=call_type,
|
||||
start_time=now,
|
||||
litellm_call_id=f"{call_type}-call",
|
||||
function_id=f"{call_type}-fn",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""}
|
||||
)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=caller, cache=cache, data={"response_id": issued_id}, call_type=call_type
|
||||
)
|
||||
provider_answer = answer(data["response_id"])
|
||||
|
||||
client_response = await hook.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=caller, response=provider_answer
|
||||
)
|
||||
await logging_obj.async_success_handler(result=provider_answer, start_time=now, end_time=now, cache_hit=False)
|
||||
return client_response.id, recorder
|
||||
|
||||
|
||||
class TestEveryReadLogsAnIdOfItsOwn:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"])
|
||||
@pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"])
|
||||
async def test_callbacks_and_the_spend_row_log_a_new_id_per_read_never_the_one_the_client_holds(
|
||||
self, mock_cache, salt_key_env, call_type, reader
|
||||
):
|
||||
hook = _hook()
|
||||
issued_id = _issue_managed_id(hook, _auth())
|
||||
|
||||
reads = [await _read_logged_by_callbacks(hook, reader, issued_id, call_type, mock_cache) for _ in range(2)]
|
||||
|
||||
assert [client_id for client_id, _ in reads] == [issued_id, issued_id]
|
||||
logged_ids = [id_ for _, recorder in reads for id_ in recorder.standard_logging_ids]
|
||||
assert [recorder.spend_row_ids for _, recorder in reads] == [[logged_id] for logged_id in logged_ids]
|
||||
assert len(set(logged_ids)) == 2
|
||||
assert issued_id not in logged_ids
|
||||
assert {hook._decrypt_response_id(logged_id) for logged_id in logged_ids} == {
|
||||
(_FABRICATED_PROVIDER_RESPONSE_ID, reader.user_id, reader.team_id)
|
||||
}
|
||||
|
||||
|
||||
_RAW_UPSTREAM_RESPONSE_ID = "resp_06de59f3upstreamiddddddddddddddddd"
|
||||
|
||||
|
||||
def _routed_response_id(upstream_response_id: str) -> str:
|
||||
return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id(
|
||||
_provider_answer(upstream_response_id), "openai", {"model_info": {"id": "deployment-1"}}
|
||||
).id
|
||||
|
||||
|
||||
def _deletion_of(response_id: str) -> DeleteResponseResult:
|
||||
return DeleteResponseResult(
|
||||
id=ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id),
|
||||
object="response.deleted",
|
||||
deleted=True,
|
||||
)
|
||||
|
||||
|
||||
class TestDeleteAnswersWithTheIdTheClientHolds:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"])
|
||||
@pytest.mark.parametrize(
|
||||
"stored_id",
|
||||
[_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID],
|
||||
ids=["routed-deployment", "plain-id"],
|
||||
)
|
||||
async def test_a_caller_with_access_gets_back_the_exact_id_they_sent(
|
||||
self, mock_cache, salt_key_env, reader, stored_id
|
||||
):
|
||||
hook = _hook()
|
||||
issued_id = _issue_managed_id(hook, _auth(), provider_response_id=stored_id)
|
||||
data = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type="adelete_responses"
|
||||
)
|
||||
|
||||
result = await hook.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=reader, response=_deletion_of(data["response_id"])
|
||||
)
|
||||
|
||||
assert (result.id, result.object, result.deleted) == (issued_id, "response.deleted", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"])
|
||||
async def test_callbacks_and_the_spend_row_log_an_id_of_its_own_never_the_one_the_client_holds(
|
||||
self, mock_cache, salt_key_env, reader
|
||||
):
|
||||
hook = _hook()
|
||||
issued_id = _issue_managed_id(
|
||||
hook, _auth(), provider_response_id=_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID)
|
||||
)
|
||||
|
||||
client_id, recorder = await _read_logged_by_callbacks(
|
||||
hook, reader, issued_id, "adelete_responses", mock_cache, answer=_deletion_of
|
||||
)
|
||||
|
||||
assert client_id == issued_id
|
||||
assert recorder.spend_row_ids == recorder.callback_response_ids
|
||||
assert len(recorder.callback_response_ids) == 1
|
||||
assert recorder.callback_response_ids[0] != issued_id
|
||||
assert hook._decrypt_response_id(recorder.callback_response_ids[0]) == (
|
||||
_RAW_UPSTREAM_RESPONSE_ID,
|
||||
reader.user_id,
|
||||
reader.team_id,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"sent_id",
|
||||
[_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID],
|
||||
ids=["routed-deployment", "plain-id"],
|
||||
)
|
||||
async def test_an_admin_delete_by_a_raw_provider_id_gets_the_provider_id_back_and_logs_an_id_of_its_own(
|
||||
self, mock_cache, salt_key_env, sent_id
|
||||
):
|
||||
hook = _hook()
|
||||
|
||||
client_id, recorder = await _read_logged_by_callbacks(
|
||||
hook, _admin(), sent_id, "adelete_responses", mock_cache, answer=_deletion_of
|
||||
)
|
||||
|
||||
assert client_id == _RAW_UPSTREAM_RESPONSE_ID
|
||||
assert [hook._decrypt_response_id(logged_id) for logged_id in recorder.callback_response_ids] == [
|
||||
(_RAW_UPSTREAM_RESPONSE_ID, "admin-user", "admin-team")
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue