cherry-pick: tag query fix + MCP metadata support (#25145)
* added support for metadata (#24261) * added support for metadata * fix: PR review - meta truthiness, BlobResourceContents mimeType, add Blob+empty meta tests Made-with: Cursor * pyproject to .25 * feat(teams): resolve access group models/MCPs/agents in team endpoints Add access_group_models, access_group_mcp_server_ids, and access_group_agent_ids to /team/info and /v2/team/list responses. These fields contain resources inherited from access groups, kept separate from direct assignments so the UI can distinguish the source. Backend: _resolve_access_group_resources() helper resolves access group resources via existing _get_*_from_access_groups() functions. UI: Teams table and detail view show direct models as blue badges and access-group-sourced models as green badges. * perf(teams): single-pass access group resolution + asyncio.gather in list endpoint - Fetch each access group object once and extract all 3 resource fields in a single pass instead of 3 separate calls (3N → N lookups) - Use asyncio.gather to resolve access groups across teams concurrently in list_team_v2 instead of sequential awaits - Add 5 unit tests for _resolve_access_group_resources * docs: add default_team_params to config reference and update examples - Add default_team_params to litellm_settings reference table in config_settings.md with all sub-fields documented - Update self_serve.md and msft_sso.md examples to include team_member_permissions, tpm_limit, and rpm_limit - Fix misleading comment that implied default_team_params only applies to SSO auto-created teams — it applies to all /team/new calls * docs: clarify that models sub-field only applies to SSO auto-created teams * fix: lazy import get_access_object to break cyclic import + short-circuit all-proxy-models display - Remove get_access_object from module-level import in team_endpoints.py and use a lazy _get_access_object wrapper to avoid cyclic dependency - Add _prisma_client is None early-exit guard in _resolve_access_group_resources - Short-circuit UI to show "All Proxy Models" when team.models is empty or contains "all-proxy-models", skipping access group model resolution * add: making organizations a select instead of read only badges * fix(ui): only send organization_id when changed and use raw initial value * fix(ui): add paginated team search to usage page filter Replace the static team dropdown on the usage page with a new TeamMultiSelect component that uses the paginated v2/team/list endpoint with debounced server-side search and infinite scroll. * fix(ui): fix imports and update placeholder for team multi select * fix(ui): wire team_id filter to key alias dropdown on Virtual Keys tab The Key Alias dropdown on the Virtual Keys page was showing aliases from all teams regardless of which team was selected. The team_id was never passed through the frontend chain to the backend /key/aliases endpoint. - Backend: add optional team_id query param to /key/aliases endpoint - networking.tsx: add team_id param to keyAliasesCall - useKeyAliases: accept and forward team_id to API call and query key - filter.tsx: pass allFilters context to custom filter components - PaginatedKeyAliasSelect: read Team ID from allFilters and pass to hook * fix(tests): correct mock targets in TestResolveAccessGroupResources Three tests were patching the non-existent `get_access_object` instead of `_get_access_object` (the lazy-import wrapper), causing AttributeError. Also added missing `prisma_client` mock so tests get past the early-exit guard and actually exercise the resolution logic. * fix: use direct attribute access with or [] fallback in _resolve_access_group_resources Replace getattr(ag, "field", []) with ag.field or [] for cleaner access and safe handling if a field is None. * fix(ui): remove model source legend from team detail view The blue/green color distinction is self-explanatory; the legend added visual clutter without providing enough value. * fix(ui): add missing access_group fields to TeamData.team_info type The TeamData interface was missing access_group_models, access_group_mcp_server_ids, and access_group_agent_ids fields, causing a TypeScript build failure. * perf(teams): batch-fetch access groups in single DB query Replace per-ID _resolve_access_group_resources loop with a single find_many call that deduplicates IDs across all teams. Removes the N+1 query pattern on cold cache for the team list endpoint. * refactor(proxy): extract helpers to fix PLR0915 violations Extract `_apply_non_admin_alias_scope` from `key_aliases`, `_resolve_team_access_group_resources` from `team_info`, and `_enforce_list_team_v2_access` from `list_team_v2` to bring each function under ruff's 50-statement limit. No behavior changes. * test(ui): update tests to match new team_id / access-group signatures - useKeyAliases, PaginatedKeyAliasSelect: add trailing `undefined` to spy matchers for the new `team_id` param on `useInfiniteKeyAliases` and `keyAliasesCall`. - EntityUsage: mock new `TeamMultiSelect` child so QueryClientProvider is not required for team-entity tests. - ModelsCell: replace the overflow-accordion test with one that verifies the new collapse-on-`all-proxy-models` behavior (no accordion, single badge). * fix(ui): send null (not '') for cleared organization_id on team update AntD <Select allowClear> returns undefined when the user clears the selection. Coalescing to "" caused the team-update payload to carry organization_id: "" instead of null, relying on the backend to coerce it. Send null directly so the intent is explicit at the source. * poetry * chore: regen poetry.lock for litellm-proxy-extras 0.4.64 bump * chore: update Next.js build artifacts (2026-04-04 17:55 UTC, node v22.16.0) --------- Co-authored-by: shivam <shivam@uni.minerva.edu> Co-authored-by: Ryan Crabbe <ryan@berri.ai> Co-authored-by: yuneng-jiang <yuneng@berri.ai> * Tag query fix (#25094) * feat(tag-spend): implement separate scheduler job for daily tag spend updates * fix(docker): add g++ to build dependencies in Dockerfile * initial test cases. TODO: check scheduler init and test cases in proxy_server related to it * resolved QPS issue when redis transaction buffer is enabled * resolving circular import error flagged by greptile * fix(mypy): use Optional[str] for api_base in PydanticAI provider to match superclass signature --------- Co-authored-by: Shivam Rawat <shivam@berri.ai> Co-authored-by: shivam <shivam@uni.minerva.edu> Co-authored-by: Ryan Crabbe <ryan@berri.ai> Co-authored-by: yuneng-jiang <yuneng@berri.ai> Co-authored-by: Harish <harishgokul01@gmail.com> Co-authored-by: Ishaan Jaffer <ishaan@berri.ai>
This commit is contained in:
ishaan-berri
GitHub
parent
693ad49719
commit
61b295238b
@ -15,6 +15,7 @@ USER root
|
||||
# Install build dependencies in one layer
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
gcc \
|
||||
g++ \
|
||||
python3-dev \
|
||||
libssl-dev \
|
||||
pkg-config \
|
||||
|
||||
@ -1344,6 +1344,7 @@ LITELLM_UI_SESSION_DURATION = os.getenv("LITELLM_UI_SESSION_DURATION", "24h")
|
||||
|
||||
########################### DB CRON JOB NAMES ###########################
|
||||
DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
|
||||
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job"
|
||||
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
|
||||
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
|
||||
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
||||
@ -1397,6 +1398,10 @@ APSCHEDULER_REPLACE_EXISTING = os.getenv(
|
||||
"1",
|
||||
] # always replace existing jobs
|
||||
|
||||
# The number of tag entries are higher than number of user, team entries. This leads to a higher QPS.
|
||||
# This will run tag spcific tasks at a later time to smooth QPS
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER = 2.3
|
||||
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL = int(
|
||||
os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)
|
||||
) # 5 minutes
|
||||
|
||||
@ -172,6 +172,34 @@ if MCP_AVAILABLE:
|
||||
mcp_info: Optional[MCPInfo] = None
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
def _normalize_resource_contents(contents: list) -> List[ReadResourceContents]:
|
||||
"""Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+)."""
|
||||
normalized: List[ReadResourceContents] = []
|
||||
for content in contents:
|
||||
meta = getattr(content, "meta", None)
|
||||
if meta is None and hasattr(content, "model_dump"):
|
||||
d = content.model_dump()
|
||||
meta = d.get("meta")
|
||||
if meta is None:
|
||||
meta = d.get("_meta")
|
||||
if isinstance(content, TextResourceContents):
|
||||
normalized.append(
|
||||
ReadResourceContents(
|
||||
content=content.text,
|
||||
mime_type=content.mimeType,
|
||||
meta=meta,
|
||||
)
|
||||
)
|
||||
elif isinstance(content, BlobResourceContents):
|
||||
normalized.append(
|
||||
ReadResourceContents(
|
||||
content=content.blob,
|
||||
mime_type=content.mimeType,
|
||||
meta=meta,
|
||||
)
|
||||
)
|
||||
return normalized
|
||||
|
||||
########################################################
|
||||
############ Initialize the MCP Server #################
|
||||
########################################################
|
||||
@ -632,26 +660,7 @@ if MCP_AVAILABLE:
|
||||
raw_headers=raw_headers,
|
||||
)
|
||||
|
||||
normalized_contents: List[ReadResourceContents] = []
|
||||
for content in read_resource_result.contents:
|
||||
if isinstance(content, TextResourceContents):
|
||||
text_content: TextResourceContents = content
|
||||
normalized_contents.append(
|
||||
ReadResourceContents(
|
||||
content=text_content.text,
|
||||
mime_type=text_content.mimeType,
|
||||
)
|
||||
)
|
||||
elif isinstance(content, BlobResourceContents):
|
||||
blob_content: BlobResourceContents = content
|
||||
normalized_contents.append(
|
||||
ReadResourceContents(
|
||||
content=blob_content.blob,
|
||||
mime_type=None,
|
||||
)
|
||||
)
|
||||
|
||||
return normalized_contents
|
||||
return _normalize_resource_contents(read_resource_result.contents)
|
||||
|
||||
########################################################
|
||||
############ End of MCP Server Routes ##################
|
||||
|
||||
@ -28,7 +28,7 @@ from typing import (
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import DualCache, RedisCache
|
||||
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
|
||||
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME,DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.proxy._types import (
|
||||
DB_CONNECTION_ERROR_TYPES,
|
||||
@ -797,7 +797,6 @@ class DBSpendUpdateWriter:
|
||||
daily_org_spend_update_queue=self.daily_org_spend_update_queue,
|
||||
daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue,
|
||||
daily_agent_spend_update_queue=self.daily_agent_spend_update_queue,
|
||||
daily_tag_spend_update_queue=self.daily_tag_spend_update_queue,
|
||||
)
|
||||
|
||||
# Only commit from redis to db if this pod is the leader
|
||||
@ -814,7 +813,6 @@ class DBSpendUpdateWriter:
|
||||
daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions,
|
||||
daily_tag_spend_update_transactions,
|
||||
) = (
|
||||
await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
)
|
||||
@ -890,13 +888,6 @@ class DBSpendUpdateWriter:
|
||||
daily_spend_transactions=daily_org_spend_update_transactions,
|
||||
)
|
||||
|
||||
if daily_tag_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
if daily_end_user_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_end_user_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@ -991,19 +982,7 @@ class DBSpendUpdateWriter:
|
||||
daily_spend_transactions=daily_org_spend_update_transactions,
|
||||
)
|
||||
|
||||
################## Daily Tag Spend Update Transactions ##################
|
||||
# Aggregate all in memory daily tag spend transactions and commit to db
|
||||
daily_tag_spend_update_transactions = cast(
|
||||
Dict[str, DailyTagSpendTransaction],
|
||||
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
|
||||
)
|
||||
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
# NOTE: Daily tag spend is committed by a separate scheduler job.
|
||||
|
||||
################## Daily End-User Spend Update Transactions ##################
|
||||
# Aggregate all in memory daily end-user spend transactions and commit to db
|
||||
@ -1032,10 +1011,75 @@ class DBSpendUpdateWriter:
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_agent_spend_update_transactions,
|
||||
)
|
||||
|
||||
|
||||
################## Tool Registry Upserts ##################
|
||||
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
|
||||
|
||||
async def _commit_daily_tag_spend_to_db(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
n_retry_times: int,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""
|
||||
Commit only tag spend updates to database.
|
||||
This is called by a separate scheduler job at a longer interval.
|
||||
"""
|
||||
daily_tag_spend_update_transactions = cast(
|
||||
Dict[str, DailyTagSpendTransaction],
|
||||
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
|
||||
)
|
||||
|
||||
if daily_tag_spend_update_transactions:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
|
||||
async def _commit_daily_tag_spend_to_db_with_redis(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
n_retry_times: int,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""
|
||||
Commit daily tag spend updates using Redis buffering.
|
||||
|
||||
This lets the dedicated daily tag scheduler drain both in-memory and
|
||||
Redis-backed tag transactions.
|
||||
"""
|
||||
await self.redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(
|
||||
daily_tag_spend_update_queue=self.daily_tag_spend_update_queue,
|
||||
)
|
||||
|
||||
if await self.pod_lock_manager.acquire_lock(
|
||||
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
|
||||
):
|
||||
verbose_proxy_logger.debug("acquired lock for daily tag spend updates")
|
||||
try:
|
||||
daily_tag_spend_update_transactions = await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
|
||||
|
||||
if daily_tag_spend_update_transactions:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - failed to commit daily tag spend updates from Redis to DB. "
|
||||
"Data already popped from Redis may be lost. Error: %s\n%s",
|
||||
str(e),
|
||||
traceback.format_exc(),
|
||||
)
|
||||
finally:
|
||||
await self.pod_lock_manager.release_lock(
|
||||
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
|
||||
)
|
||||
|
||||
async def _flush_tool_discovery_queue(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
|
||||
@ -131,7 +131,6 @@ class RedisUpdateBuffer:
|
||||
daily_org_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_agent_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_tag_spend_update_queue: DailySpendUpdateQueue,
|
||||
):
|
||||
"""
|
||||
Stores the in-memory spend updates to Redis
|
||||
@ -202,9 +201,6 @@ class RedisUpdateBuffer:
|
||||
daily_agent_spend_update_transactions = (
|
||||
await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
daily_tag_spend_update_transactions = (
|
||||
await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions
|
||||
@ -245,11 +241,6 @@ class RedisUpdateBuffer:
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
|
||||
),
|
||||
(
|
||||
daily_tag_spend_update_transactions,
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
|
||||
),
|
||||
]
|
||||
|
||||
rpush_list: List[RedisPipelineRpushOperation] = []
|
||||
@ -376,22 +367,20 @@ class RedisUpdateBuffer:
|
||||
Optional[Dict[str, DailyOrganizationSpendTransaction]],
|
||||
Optional[Dict[str, DailyEndUserSpendTransaction]],
|
||||
Optional[Dict[str, DailyAgentSpendTransaction]],
|
||||
Optional[Dict[str, DailyTagSpendTransaction]],
|
||||
]:
|
||||
"""
|
||||
Drains all 7 Redis buffer queues in a single pipeline round-trip.
|
||||
Drains the main 6 Redis buffer queues in a single pipeline round-trip.
|
||||
|
||||
Returns a 7-tuple of parsed results in this order:
|
||||
Returns a 6-tuple of parsed results in this order:
|
||||
0: DBSpendUpdateTransactions
|
||||
1: daily user spend
|
||||
2: daily team spend
|
||||
3: daily org spend
|
||||
4: daily end-user spend
|
||||
5: daily agent spend
|
||||
6: daily tag spend
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
return None, None, None, None, None, None, None
|
||||
return None, None, None, None, None, None
|
||||
|
||||
lpop_list: List[RedisPipelineLpopOperation] = [
|
||||
RedisPipelineLpopOperation(
|
||||
@ -417,16 +406,12 @@ class RedisUpdateBuffer:
|
||||
key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
),
|
||||
RedisPipelineLpopOperation(
|
||||
key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
),
|
||||
]
|
||||
|
||||
raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
# Pad with None if pipeline returned fewer results than expected
|
||||
while len(raw_results) < 7:
|
||||
while len(raw_results) < 6:
|
||||
raw_results.append(None)
|
||||
|
||||
# Slot 0: DBSpendUpdateTransactions
|
||||
@ -436,9 +421,9 @@ class RedisUpdateBuffer:
|
||||
if len(parsed) > 0:
|
||||
db_spend = self._combine_list_of_transactions(parsed)
|
||||
|
||||
# Slots 1-6: daily spend categories
|
||||
# Slots 1-5: daily spend categories
|
||||
daily_results: List[Optional[Dict[str, Any]]] = []
|
||||
for slot in range(1, 7):
|
||||
for slot in range(1, 6):
|
||||
if raw_results[slot] is None:
|
||||
daily_results.append(None)
|
||||
else:
|
||||
@ -457,7 +442,22 @@ class RedisUpdateBuffer:
|
||||
),
|
||||
cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]),
|
||||
cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]),
|
||||
cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]),
|
||||
)
|
||||
|
||||
async def store_in_memory_daily_tag_spend_updates_in_redis(
|
||||
self,
|
||||
daily_tag_spend_update_queue: DailySpendUpdateQueue,
|
||||
) -> None:
|
||||
"""
|
||||
Flush in-memory daily tag spend updates and append them to Redis.
|
||||
"""
|
||||
daily_tag_spend_update_transactions = (
|
||||
await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_tag_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
async def get_all_daily_spend_update_transactions_from_redis_buffer(
|
||||
|
||||
@ -54,6 +54,7 @@ from litellm.constants import (
|
||||
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
|
||||
LITELLM_UI_ALLOW_HEADERS,
|
||||
LITELLM_UI_SESSION_DURATION,
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
@ -6284,6 +6285,25 @@ class ProxyStartupEvent:
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ###
|
||||
## Reduces QPS as there are more tags for a single request
|
||||
tag_spend_update_interval = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER)
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
|
||||
scheduler.add_job(
|
||||
update_daily_tag_spend,
|
||||
"interval",
|
||||
seconds=tag_spend_update_interval,
|
||||
args=[prisma_client, proxy_logging_obj],
|
||||
id="update_daily_tag_spend_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Tag spend update job scheduled at {tag_spend_update_interval}s interval "
|
||||
f"({tag_spend_update_interval / batch_writing_interval:.1f}x main job interval)"
|
||||
)
|
||||
|
||||
### MONITOR SPEND LOGS QUEUE (queue-size-based job) ###
|
||||
if general_settings.get("disable_spend_logs", False) is False:
|
||||
from litellm.proxy.utils import _monitor_spend_logs_queue
|
||||
|
||||
@ -4827,6 +4827,9 @@ async def update_spend( # noqa: PLR0915
|
||||
|
||||
Triggered every minute.
|
||||
|
||||
NOTE: This job now skips tag spend updates, which are handled by a separate
|
||||
scheduler job (update_daily_tag_spend) at a longer interval to reduce contention.
|
||||
|
||||
Requires:
|
||||
user_id_list: dict,
|
||||
keys_list: list,
|
||||
@ -4859,6 +4862,46 @@ async def update_spend( # noqa: PLR0915
|
||||
)
|
||||
|
||||
|
||||
async def update_daily_tag_spend(
|
||||
prisma_client: PrismaClient,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""
|
||||
Separate scheduler job to commit daily tag spend updates.
|
||||
|
||||
Runs at a longer interval (2.3x default) than the main update_spend job
|
||||
to reduce query contention for DailyTagSpend table.
|
||||
|
||||
This is called by a dedicated scheduler job and does NOT process:
|
||||
- Regular spend updates (user, key, team, org)
|
||||
- End-user spend
|
||||
- Agent spend
|
||||
- Spend logs
|
||||
|
||||
Only processes tag spend transactions from the daily_tag_spend_update_queue.
|
||||
|
||||
Args:
|
||||
prisma_client: PrismaClient instance
|
||||
proxy_logging_obj: ProxyLogging instance for error handling
|
||||
"""
|
||||
n_retry_times = 3
|
||||
try:
|
||||
if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis():
|
||||
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis(
|
||||
prisma_client=prisma_client,
|
||||
n_retry_times=n_retry_times,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
else:
|
||||
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db(
|
||||
prisma_client=prisma_client,
|
||||
n_retry_times=n_retry_times,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error updating daily tag spend: {e}")
|
||||
|
||||
|
||||
async def update_spend_logs_job(
|
||||
prisma_client: PrismaClient,
|
||||
db_writer_client: Optional[AsyncHTTPHandler],
|
||||
|
||||
134
tests/proxy_unit_tests/test_update_daily_tag_spend.py
Normal file
134
tests/proxy_unit_tests/test_update_daily_tag_spend.py
Normal file
@ -0,0 +1,134 @@
|
||||
from typing import Dict
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
from litellm.proxy._types import DailyTagSpendTransaction
|
||||
import httpx
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_daily_tag_spend_delegates_to_tag_commit_writer():
|
||||
prisma_client = MagicMock()
|
||||
proxy_logging_obj = MagicMock()
|
||||
redis_update_buffer = MagicMock()
|
||||
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False
|
||||
proxy_logging_obj.db_spend_update_writer = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock()
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
|
||||
|
||||
await update_daily_tag_spend(
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once_with(
|
||||
prisma_client=prisma_client,
|
||||
n_retry_times=3,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_daily_tag_spend_logs_error_and_does_not_raise():
|
||||
prisma_client = MagicMock()
|
||||
proxy_logging_obj = MagicMock()
|
||||
redis_update_buffer = MagicMock()
|
||||
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False
|
||||
proxy_logging_obj.db_spend_update_writer = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock(
|
||||
side_effect=ValueError("boom")
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.utils.verbose_proxy_logger.error") as error_logger:
|
||||
await update_daily_tag_spend(
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once()
|
||||
error_logger.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_daily_tag_spend_uses_redis_writer_when_enabled():
|
||||
prisma_client = MagicMock()
|
||||
proxy_logging_obj = MagicMock()
|
||||
redis_update_buffer = MagicMock()
|
||||
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = True
|
||||
proxy_logging_obj.db_spend_update_writer = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock()
|
||||
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
|
||||
|
||||
await update_daily_tag_spend(
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_awaited_once_with(
|
||||
prisma_client=prisma_client,
|
||||
n_retry_times=3,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_tag_spend_retries_then_succeeds():
|
||||
prisma_client = MagicMock()
|
||||
proxy_logging_obj = MagicMock()
|
||||
|
||||
mock_batcher = MagicMock()
|
||||
mock_table = MagicMock()
|
||||
mock_batcher.litellm_dailytagspend = mock_table
|
||||
|
||||
# Fail entering batch context 3 times with retryable DB errors, then succeed.
|
||||
prisma_client.db.batch_.return_value.__aenter__ = AsyncMock(
|
||||
side_effect=[
|
||||
httpx.ConnectError("x"),
|
||||
httpx.ConnectError("x"),
|
||||
httpx.ConnectError("x"),
|
||||
mock_batcher,
|
||||
]
|
||||
)
|
||||
|
||||
daily_spend_transactions: Dict[str, DailyTagSpendTransaction] = {
|
||||
"k": {
|
||||
"tag": "prod-tag",
|
||||
"date": "2026-04-03",
|
||||
"api_key": "key-1",
|
||||
"model": "gpt-4o",
|
||||
"model_group": None,
|
||||
"custom_llm_provider": "openai",
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"endpoint": "",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"spend": 0.01,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
"failed_requests": 0,
|
||||
"request_id": None,
|
||||
}
|
||||
}
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock) as sleep_mock, patch(
|
||||
"random.uniform", return_value=0
|
||||
):
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=3,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_spend_transactions,
|
||||
)
|
||||
|
||||
assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4
|
||||
assert sleep_mock.await_count == 3
|
||||
mock_table.upsert.assert_called_once()
|
||||
@ -1,12 +1,11 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from mcp import ReadResourceResult, Resource
|
||||
from mcp.types import Prompt, ResourceTemplate, TextResourceContents
|
||||
from mcp.types import BlobResourceContents, Prompt, ResourceTemplate, TextResourceContents
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_MCPServerTable,
|
||||
@ -413,6 +412,111 @@ async def test_mcp_read_resource_success():
|
||||
assert result is read_result
|
||||
|
||||
|
||||
def test_normalize_resource_contents_passes_metadata():
|
||||
"""Test that _normalize_resource_contents preserves meta from ResourceContents (MCP 1.26.0+)."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_normalize_resource_contents,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
meta = {"version": "1.0", "source": "test"}
|
||||
contents = [
|
||||
TextResourceContents(
|
||||
uri="https://example.com/resource",
|
||||
text="hello world",
|
||||
mimeType="text/plain",
|
||||
meta=meta,
|
||||
)
|
||||
]
|
||||
|
||||
result = _normalize_resource_contents(contents)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].content == "hello world"
|
||||
assert result[0].mime_type == "text/plain"
|
||||
assert result[0].meta == meta
|
||||
|
||||
|
||||
def test_normalize_resource_contents_blob_with_metadata():
|
||||
"""Test that _normalize_resource_contents preserves meta for BlobResourceContents."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_normalize_resource_contents,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
meta = {"encoding": "base64"}
|
||||
contents = [
|
||||
BlobResourceContents(
|
||||
uri="https://example.com/image.png",
|
||||
blob="aGVsbG8=",
|
||||
mimeType="image/png",
|
||||
meta=meta,
|
||||
)
|
||||
]
|
||||
|
||||
result = _normalize_resource_contents(contents)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].content == "aGVsbG8="
|
||||
assert result[0].mime_type == "image/png"
|
||||
assert result[0].meta == meta
|
||||
|
||||
|
||||
def test_normalize_resource_contents_preserves_empty_metadata():
|
||||
"""Test that empty dict meta is preserved (truthiness bug fix)."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_normalize_resource_contents,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
empty_meta: dict = {}
|
||||
contents = [
|
||||
TextResourceContents(
|
||||
uri="https://example.com/resource",
|
||||
text="hi",
|
||||
mimeType="text/plain",
|
||||
meta=empty_meta,
|
||||
)
|
||||
]
|
||||
|
||||
result = _normalize_resource_contents(contents)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].meta == empty_meta
|
||||
assert result[0].meta is not None
|
||||
assert result[0].meta == {}
|
||||
|
||||
|
||||
def test_normalize_resource_contents_without_metadata():
|
||||
"""Test that _normalize_resource_contents works when meta is absent (backward compat)."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_normalize_resource_contents,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
contents = [
|
||||
TextResourceContents(
|
||||
uri="https://example.com/resource",
|
||||
text="hello",
|
||||
mimeType="text/plain",
|
||||
)
|
||||
]
|
||||
|
||||
result = _normalize_resource_contents(contents)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].content == "hello"
|
||||
assert result[0].meta is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_read_resource_multiple_servers_error():
|
||||
try:
|
||||
@ -707,8 +811,6 @@ async def test_concurrent_initialize_session_managers():
|
||||
"""Test that concurrent calls to initialize_session_managers don't cause race conditions."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_INITIALIZATION_LOCK,
|
||||
_SESSION_MANAGERS_INITIALIZED,
|
||||
initialize_session_managers,
|
||||
)
|
||||
except ImportError:
|
||||
@ -1426,7 +1528,6 @@ async def test_list_tools_with_team_tool_permissions_inheritance():
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TeamTable,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
except ImportError:
|
||||
|
||||
@ -35,7 +35,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
|
||||
"""
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2])
|
||||
|
||||
# Create mock queues - only 3 of 7 have data
|
||||
# Create mock queues - only 3 of 6 have data
|
||||
spend_update_queue = AsyncMock()
|
||||
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={"key_list_transactions": {"key1": 1.0}}
|
||||
@ -67,11 +67,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
|
||||
return_value={}
|
||||
)
|
||||
|
||||
daily_tag_queue = AsyncMock()
|
||||
daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_queue,
|
||||
@ -79,7 +74,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
|
||||
daily_org_spend_update_queue=daily_org_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_queue,
|
||||
daily_tag_spend_update_queue=daily_tag_queue,
|
||||
)
|
||||
|
||||
# Should be called exactly once (pipeline)
|
||||
@ -117,7 +111,6 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
||||
daily_org_spend_update_queue=empty_daily_queue,
|
||||
daily_end_user_spend_update_queue=empty_daily_queue,
|
||||
daily_agent_spend_update_queue=empty_daily_queue,
|
||||
daily_tag_spend_update_queue=empty_daily_queue,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline.assert_not_called()
|
||||
@ -131,7 +124,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
||||
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
|
||||
and aggregates results from async_lpop_pipeline.
|
||||
"""
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories
|
||||
db_spend_json = json.dumps(
|
||||
{
|
||||
"key_list_transactions": {"key1": 1.0, "key2": 2.0},
|
||||
@ -154,14 +147,13 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
||||
None, # slot 3: daily org (empty)
|
||||
None, # slot 4: daily end-user (empty)
|
||||
None, # slot 5: daily agent (empty)
|
||||
None, # slot 6: daily tag (empty)
|
||||
]
|
||||
)
|
||||
|
||||
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
assert len(result) == 7
|
||||
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result
|
||||
assert len(result) == 6
|
||||
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result
|
||||
|
||||
# Verify db spend was parsed correctly
|
||||
assert db_spend is not None
|
||||
@ -181,7 +173,6 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
||||
assert daily_org is None
|
||||
assert daily_end_user is None
|
||||
assert daily_agent is None
|
||||
assert daily_tag is None
|
||||
|
||||
# Verify pipeline was called once with correct keys
|
||||
mock_redis_cache.async_lpop_pipeline.assert_called_once()
|
||||
@ -192,7 +183,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
|
||||
"""When redis_cache is None, should return all Nones"""
|
||||
buffer = RedisUpdateBuffer(redis_cache=None)
|
||||
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert result == (None, None, None, None, None, None, None)
|
||||
assert result == (None, None, None, None, None, None)
|
||||
|
||||
|
||||
def test_validate_redis_transaction_buffer_raises_without_redis():
|
||||
|
||||
Loading…
Reference in New Issue
Block a user