mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
* bump: version 1.82.1 → 1.82.2 * fix(gemini): preserve toolConfig on native generate_content (#23493) * chore: regenerate poetry.lock to match pyproject.toml (#23514) Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> * fix claude.md * ui logo (#23556) * fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll (#23472) * fix(proxy): cap managed-object poll size + expire stale rows + kill-switch flag to prevent OOM/Prisma connection loss * fix(constants): simplify PROXY_BATCH_POLLING_ENABLED readability * docs+test: document new polling env vars, add pagination+stale-cleanup tests * fix: exclude stale_expired from batch poll queries; fix update_many assertions in tests * fix: scope stale cleanup to file_purpose, fix file_object mocks, add CheckBatchCost tests * fix: avoid duplicate cost logging in fallback path; guard integer constants against zero/negative values * fix: cache _has_batch_processed_column; guard cleanup from aborting poll; narrow fallback except * fix: add complete/completed to primary query not_in; fix vacuous test assertion - Primary find_many was missing "complete" and "completed" in its not_in filter, creating asymmetry with the fallback query. A job whose status was set to "complete" but whose batch_processed flag update failed would be silently re-fetched and re-processed every cycle, emitting duplicate cost logs. - test_fallback_completion_update_omits_batch_processed patched _is_base64_encoded_unified_file_id to return None, causing an immediate continue — so update() was never called and the assertion looped over an empty list (vacuously true). Rewrote the test to mock the full completion pipeline, verify update() is called exactly once, and assert batch_processed is absent from the update data. - Added symmetric test (primary path) proving batch_processed IS included when the column exists. Made-with: Cursor * fix(huggingface): forward extra_headers to embedding handler (#23502) The huggingface branch in litellm.embedding() did not pass the headers kwarg to huggingface_embed.embedding(), silently dropping user-provided extra_headers like X-HF-Bill-To. Fixes #23502 Made-with: Cursor --------- Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com> Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
This commit is contained in:
palanisd
GitHub
parent
239edc472a
commit
02cf15c87d
12 changed files with 803 additions and 72 deletions
|
|
@ -91,6 +91,10 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
|
|||
- Async/await patterns throughout
|
||||
- Type hints required for all public APIs
|
||||
- **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary.
|
||||
- **Use dict spread for immutable copies** — prefer `{**original, "key": new_value}` over `dict(obj)` + mutation. The spread produces the final dict in one step and makes intent clear.
|
||||
- **Guard at resolution time** — when resolving an optional value through a fallback chain (`a or b or ""`), raise immediately if the resolved result being empty is an error. Don't pass empty strings or sentinel values downstream for the callee to deal with.
|
||||
- **Extract complex comprehensions to named helpers** — a set/dict comprehension that calls into the DB or manager (e.g. "which of these server IDs are OAuth2?") belongs in a named helper function, not inline in the caller.
|
||||
- **FastAPI parameter declarations** — mark required query/form params with `= Query(...)` / `= Form(...)` explicitly when other params in the same handler are optional. Mixing `str` (required) with `Optional[str] = None` in the same signature causes silent 422s when the required param is missing.
|
||||
|
||||
### Testing Strategy
|
||||
- Unit tests in `tests/test_litellm/`
|
||||
|
|
@ -98,6 +102,8 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
|
|||
- Proxy tests in `tests/proxy_unit_tests/`
|
||||
- Load tests in `tests/load_tests/`
|
||||
- **Always add tests when adding new entity types or features** — if the existing test file covers other entity types, add corresponding tests for the new one
|
||||
- **Keep monkeypatch stubs in sync with real signatures** — when a function gains a new optional parameter, update every `fake_*` / `stub_*` in tests that patch it to also accept that kwarg (even as `**kwargs`). Stale stubs fail with `unexpected keyword argument` and mask real bugs.
|
||||
- **Test all branches of name→ID resolution** — when adding server/resource lookup that resolves names to UUIDs, test: (1) name resolves and UUID is allowed, (2) name resolves but UUID is not allowed, (3) name does not resolve at all. The silent-fallback path is where access-control bugs hide.
|
||||
|
||||
### UI / Backend Consistency
|
||||
- When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select
|
||||
|
|
|
|||
|
|
@ -935,6 +935,9 @@ router_settings:
|
|||
| PROXY_BASE_URL | Base URL for proxy service
|
||||
| PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 10
|
||||
| PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour)
|
||||
| PROXY_BATCH_POLLING_ENABLED | Set to `false` to disable the `CheckBatchCost` and `CheckResponsesCost` background polling jobs entirely. Useful for emergency mitigation on installs with large numbers of stale managed objects. Default is `true`
|
||||
| MAX_OBJECTS_PER_POLL_CYCLE | Maximum number of managed objects (batches / responses) fetched per polling cycle. Prevents OOM on installs with many stale rows. Default is `50`
|
||||
| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS | Managed objects older than this many days in a non-terminal state are marked `stale_expired` at the start of each poll cycle and skipped. Default is `7`
|
||||
| PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605
|
||||
| PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597
|
||||
| PYTHON_GC_THRESHOLD | GC thresholds ('gen0,gen1,gen2', e.g. '1000,50,50'); defaults to Python’s values.
|
||||
|
|
|
|||
138
docs/my-website/docs/proxy/ui/ui_edit_logo.md
Normal file
138
docs/my-website/docs/proxy/ui/ui_edit_logo.md
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
import Image from '@theme/IdealImage';
|
||||
|
||||
# Customize UI Logo
|
||||
|
||||
Personalize your LiteLLM dashboard by replacing the default logo with your own company branding. You can set a custom logo via the UI or the API.
|
||||
|
||||
## Via the UI
|
||||
|
||||
### 1. Navigate to Settings
|
||||
|
||||
Click the **Settings** icon in the sidebar.
|
||||
|
||||
|
||||
|
||||
### 2. Open UI Theme Settings
|
||||
|
||||
Click **UI Theme** from the settings menu.
|
||||
|
||||
|
||||
|
||||
### 3. Click the Logo URL Field
|
||||
|
||||
Click the **Logo URL** text field to start editing.
|
||||
|
||||
|
||||
|
||||
### 4. Find Your Logo Image
|
||||
|
||||
Open a new browser tab and find the logo image you want to use (e.g., search Google Images for your company logo).
|
||||
|
||||
|
||||
|
||||
### 5. Right-Click on the Logo Image
|
||||
|
||||
Right-click the image you want to use as your logo.
|
||||
|
||||
|
||||
|
||||
### 6. Copy the Image Address
|
||||
|
||||
Select **Copy Image Address** from the context menu to copy the URL.
|
||||
|
||||
|
||||
|
||||
### 7. Switch Back to LiteLLM
|
||||
|
||||
Navigate back to the LiteLLM UI tab (e.g., press **Cmd + Left** or click the tab).
|
||||
|
||||
|
||||
|
||||
### 8. Paste the Logo URL
|
||||
|
||||
Paste the copied image URL into the **Logo URL** field with **Cmd + V**.
|
||||
|
||||
|
||||
|
||||
### 9. Save Changes
|
||||
|
||||
Click **Save Changes** to apply your new logo.
|
||||
|
||||
|
||||
|
||||
Your custom logo will now appear in the LiteLLM dashboard sidebar and login page.
|
||||
|
||||
## Via the API
|
||||
|
||||
### Set a Custom Logo
|
||||
|
||||
```bash
|
||||
curl -X PATCH 'http://localhost:4000/settings/update/ui_theme_settings' \
|
||||
-H 'Authorization: Bearer <your-admin-key>' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"logo_url": "https://example.com/your-company-logo.png"
|
||||
}'
|
||||
```
|
||||
|
||||
### Set a Custom Favicon
|
||||
|
||||
You can also customize the browser tab favicon:
|
||||
|
||||
```bash
|
||||
curl -X PATCH 'http://localhost:4000/settings/update/ui_theme_settings' \
|
||||
-H 'Authorization: Bearer <your-admin-key>' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"logo_url": "https://example.com/your-company-logo.png",
|
||||
"favicon_url": "https://example.com/your-favicon.ico"
|
||||
}'
|
||||
```
|
||||
|
||||
### Get Current Theme Settings
|
||||
|
||||
```bash
|
||||
curl -X GET 'http://localhost:4000/settings/get/ui_theme_settings'
|
||||
```
|
||||
|
||||
### Reset to Default Logo
|
||||
|
||||
Send an empty `logo_url` to restore the default LiteLLM logo:
|
||||
|
||||
```bash
|
||||
curl -X PATCH 'http://localhost:4000/settings/update/ui_theme_settings' \
|
||||
-H 'Authorization: Bearer <your-admin-key>' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"logo_url": ""
|
||||
}'
|
||||
```
|
||||
|
||||
## Via `proxy_config.yaml`
|
||||
|
||||
You can also set the logo URL in your proxy configuration file:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
ui_theme_config:
|
||||
logo_url: "https://example.com/your-company-logo.png"
|
||||
favicon_url: "https://example.com/your-favicon.ico" # optional
|
||||
```
|
||||
|
||||
Or set it as an environment variable:
|
||||
|
||||
```yaml
|
||||
environment_variables:
|
||||
UI_LOGO_PATH: "https://example.com/your-company-logo.png"
|
||||
```
|
||||
|
||||
## Supported Logo Formats
|
||||
|
||||
| Format | Supported |
|
||||
|--------|-----------|
|
||||
| JPEG / JPG | Yes |
|
||||
| PNG | Yes |
|
||||
| SVG | Yes |
|
||||
| ICO (favicon only) | Yes |
|
||||
| HTTP/HTTPS URL | Yes |
|
||||
| Local file path | Yes |
|
||||
|
|
@ -332,6 +332,7 @@ const sidebars = {
|
|||
label: "Setup & SSO",
|
||||
items: [
|
||||
"proxy/admin_ui_sso",
|
||||
"proxy/ui/ui_edit_logo",
|
||||
"proxy/custom_sso",
|
||||
"proxy/custom_root_ui",
|
||||
"tutorials/scim_litellm",
|
||||
|
|
|
|||
|
|
@ -2,11 +2,15 @@
|
|||
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
|
||||
"""
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -29,6 +33,9 @@ class CheckBatchCost:
|
|||
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
|
||||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
# Cached after the first poll cycle. Once we know the column is absent we skip
|
||||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
|
||||
async def _get_user_info(self, batch_id, user_id) -> dict:
|
||||
"""
|
||||
|
|
@ -49,6 +56,47 @@ class CheckBatchCost:
|
|||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
|
||||
return {}
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
"""
|
||||
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
|
||||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
if result > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: marked {result} stale managed objects "
|
||||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
async def _fallback_find_jobs(self) -> list:
|
||||
"""Query batch jobs without the batch_processed filter (for older schemas)."""
|
||||
return await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {
|
||||
"not_in": [
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"complete",
|
||||
"completed",
|
||||
"stale_expired",
|
||||
]
|
||||
},
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
async def check_batch_cost(self):
|
||||
"""
|
||||
Check if the batch JOB has been tracked.
|
||||
|
|
@ -70,14 +118,48 @@ class CheckBatchCost:
|
|||
get_model_id_from_unified_batch_id,
|
||||
)
|
||||
|
||||
# Look for all batches that have not yet been processed by CheckBatchCost
|
||||
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed" : False,
|
||||
"status": {"not_in": ["failed", "expired", "cancelled"]}
|
||||
}
|
||||
)
|
||||
try:
|
||||
await self._cleanup_stale_managed_objects()
|
||||
except Exception as cleanup_err:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
|
||||
)
|
||||
|
||||
# Look for all batches that have not yet been processed by CheckBatchCost.
|
||||
# self._has_batch_processed_column is cached after the first probe so that
|
||||
# older schemas don't pay a guaranteed-failing primary query + warning on
|
||||
# every subsequent poll cycle.
|
||||
if self._has_batch_processed_column:
|
||||
try:
|
||||
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
"status": {
|
||||
"not_in": [
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"complete",
|
||||
"completed",
|
||||
"stale_expired",
|
||||
]
|
||||
},
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
except Exception as query_err:
|
||||
if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
|
||||
raise
|
||||
# Permanent schema gap — cache the result so future cycles skip straight to fallback
|
||||
self._has_batch_processed_column = False
|
||||
verbose_proxy_logger.warning(
|
||||
"CheckBatchCost: batch_processed column not found, querying without it"
|
||||
)
|
||||
jobs = await self._fallback_find_jobs()
|
||||
else:
|
||||
jobs = await self._fallback_find_jobs()
|
||||
for job in jobs:
|
||||
# get the model from the job
|
||||
unified_object_id = job.unified_object_id
|
||||
|
|
@ -163,14 +245,14 @@ class CheckBatchCost:
|
|||
|
||||
# Access content - handle both direct attribute and method call
|
||||
if hasattr(_file_content, 'content'):
|
||||
content_bytes = _file_content.content
|
||||
content_bytes = _file_content.content # type: ignore[union-attr]
|
||||
elif hasattr(_file_content, 'read'):
|
||||
content_bytes = await _file_content.read()
|
||||
content_bytes = await _file_content.read() # type: ignore[misc]
|
||||
else:
|
||||
content_bytes = _file_content
|
||||
content_bytes = _file_content # type: ignore[assignment]
|
||||
|
||||
file_content_as_dict = _get_file_content_as_dictionary(
|
||||
content_bytes
|
||||
content_bytes # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
deployment_info = self.llm_router.get_deployment(model_id=model_id)
|
||||
|
|
@ -195,7 +277,7 @@ class CheckBatchCost:
|
|||
file_content_dictionary=file_content_as_dict,
|
||||
custom_llm_provider=llm_provider, # type: ignore
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info,
|
||||
model_info=deployment_model_info, # type: ignore[arg-type]
|
||||
)
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
|
|
@ -236,13 +318,15 @@ class CheckBatchCost:
|
|||
|
||||
# mark the job as complete
|
||||
try:
|
||||
update_data: dict = {
|
||||
"status": "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
}
|
||||
if self._has_batch_processed_column:
|
||||
update_data["batch_processed"] = True
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update(
|
||||
where={"id": job.id},
|
||||
data={
|
||||
"batch_processed": True,
|
||||
"status": "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
},
|
||||
data=update_data,
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
|
|||
|
|
@ -3,10 +3,15 @@ Polls LiteLLM_ManagedObjectTable to check if the response is complete.
|
|||
Cost tracking is handled automatically by litellm.aget_responses().
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -27,6 +32,27 @@ class CheckResponsesCost:
|
|||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
"""
|
||||
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
|
||||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "response",
|
||||
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
if result > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckResponsesCost: marked {result} stale managed objects "
|
||||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
async def check_responses_cost(self):
|
||||
"""
|
||||
Check if background responses are complete and track their cost.
|
||||
|
|
@ -35,11 +61,20 @@ class CheckResponsesCost:
|
|||
- Cost is automatically tracked by litellm.aget_responses()
|
||||
- Mark completed/failed/cancelled responses as complete in the database
|
||||
"""
|
||||
try:
|
||||
await self._cleanup_stale_managed_objects()
|
||||
except Exception as cleanup_err:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
|
||||
)
|
||||
|
||||
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where={
|
||||
"status": {"in": ["queued", "in_progress"]},
|
||||
"file_purpose": "response",
|
||||
}
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
|
||||
|
|
|
|||
|
|
@ -1351,6 +1351,15 @@ PROXY_BUDGET_RESCHEDULER_MIN_TIME = int(
|
|||
os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
|
||||
)
|
||||
PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
|
||||
MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(
|
||||
1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
|
||||
)
|
||||
# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
|
||||
# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
|
||||
# installations with large numbers of stale managed objects).
|
||||
_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
|
||||
PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true"
|
||||
PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(
|
||||
os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5102,6 +5102,7 @@ def embedding( # noqa: PLR0915
|
|||
client=client,
|
||||
aembedding=aembedding,
|
||||
litellm_params=litellm_params_dict,
|
||||
headers=headers or {},
|
||||
)
|
||||
elif custom_llm_provider == "bedrock":
|
||||
if isinstance(input, str):
|
||||
|
|
|
|||
|
|
@ -214,6 +214,7 @@ from litellm.constants import (
|
|||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
|
||||
PROXY_BATCH_POLLING_ENABLED,
|
||||
PROXY_BATCH_POLLING_INTERVAL,
|
||||
PROXY_BATCH_WRITE_AT,
|
||||
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
|
||||
|
|
@ -260,7 +261,6 @@ from litellm.proxy.anthropic_endpoints.claude_code_endpoints import (
|
|||
claude_code_marketplace_router,
|
||||
)
|
||||
from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router
|
||||
from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
|
||||
from litellm.proxy.anthropic_endpoints.skills_endpoints import (
|
||||
router as anthropic_skills_router,
|
||||
)
|
||||
|
|
@ -471,6 +471,7 @@ from litellm.proxy.policy_engine.policy_resolve_endpoints import (
|
|||
from litellm.proxy.prompts.prompt_endpoints import router as prompts_router
|
||||
from litellm.proxy.public_endpoints import router as public_endpoints_router
|
||||
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
|
||||
from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
|
||||
from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
|
||||
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
|
|
@ -6069,7 +6070,7 @@ class ProxyStartupEvent:
|
|||
"Invalid maximum_spend_logs_retention_interval value"
|
||||
)
|
||||
### CHECK BATCH COST ###
|
||||
if llm_router is not None:
|
||||
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
|
||||
try:
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import (
|
||||
CheckBatchCost,
|
||||
|
|
@ -6100,7 +6101,7 @@ class ProxyStartupEvent:
|
|||
pass
|
||||
|
||||
### CHECK RESPONSES COST ###
|
||||
if llm_router is not None:
|
||||
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
|
||||
try:
|
||||
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
||||
CheckResponsesCost,
|
||||
|
|
|
|||
340
tests/proxy_unit_tests/test_check_batch_cost.py
Normal file
340
tests/proxy_unit_tests/test_check_batch_cost.py
Normal file
|
|
@ -0,0 +1,340 @@
|
|||
"""
|
||||
Unit tests for CheckBatchCost class.
|
||||
Covers: stale-row cleanup (file_purpose scoping), paginated find_many,
|
||||
and the batch_processed-column fallback query.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestCheckBatchCost:
|
||||
"""Test suite for CheckBatchCost class"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma_client(self):
|
||||
client = MagicMock()
|
||||
client.db = MagicMock()
|
||||
client.db.litellm_managedobjecttable = MagicMock()
|
||||
client.db.litellm_usertable = MagicMock()
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_proxy_logging_obj(self):
|
||||
return MagicMock()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm_router(self):
|
||||
return MagicMock()
|
||||
|
||||
@pytest.fixture
|
||||
def check_batch_cost_instance(
|
||||
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
|
||||
|
||||
return CheckBatchCost(
|
||||
proxy_logging_obj=mock_proxy_logging_obj,
|
||||
prisma_client=mock_prisma_client,
|
||||
llm_router=mock_llm_router,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_scoped_to_batch_file_purpose(
|
||||
self, check_batch_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only."""
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
# Return empty so the main poll loop exits immediately
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
stale_call = calls[0]
|
||||
assert stale_call[1]["data"] == {"status": "stale_expired"}
|
||||
where = stale_call[1]["where"]
|
||||
assert where["file_purpose"] == "batch"
|
||||
assert "stale_expired" in where["status"]["not_in"]
|
||||
assert "created_at" in where
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_many_uses_pagination_and_excludes_stale(
|
||||
self, check_batch_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""find_many is called with take, order, and all terminal statuses excluded."""
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
|
||||
assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
|
||||
assert find_call[1]["order"] == {"created_at": "asc"}
|
||||
not_in = find_call[1]["where"]["status"]["not_in"]
|
||||
assert "stale_expired" in not_in
|
||||
assert "complete" in not_in
|
||||
assert "completed" in not_in
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_query_used_when_batch_processed_missing(
|
||||
self, check_batch_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""Falls back to query without batch_processed when primary query raises."""
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
# First find_many (primary query) raises with a schema error; second (fallback) returns empty
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
side_effect=[Exception("column batch_processed does not exist"), []]
|
||||
)
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
|
||||
assert len(calls) == 2
|
||||
fallback_where = calls[1][1]["where"]
|
||||
assert "batch_processed" not in fallback_where
|
||||
assert "stale_expired" in fallback_where["status"]["not_in"]
|
||||
assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
|
||||
# Column absence is now cached — next call should go straight to fallback
|
||||
assert check_batch_cost_instance._has_batch_processed_column is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_column_absence_cached_across_cycles(
|
||||
self, check_batch_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""After column absence is discovered, subsequent cycles skip the primary query entirely."""
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
# Simulate column already known absent from a previous cycle
|
||||
check_batch_cost_instance._has_batch_processed_column = False
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
# Only one find_many call — the fallback directly, no primary query attempt
|
||||
assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
|
||||
fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
|
||||
assert "batch_processed" not in fallback_where
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_completion_update_omits_batch_processed(
|
||||
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""When batch_processed column is absent, completion update must not include it.
|
||||
|
||||
If it did, the update would fail silently, the job would never be marked done,
|
||||
and every subsequent poll cycle would re-log the cost (duplicate billing).
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
mock_job = MagicMock()
|
||||
mock_job.id = "job-fallback-1"
|
||||
mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
|
||||
mock_job.created_by = "user-1"
|
||||
|
||||
# Simulate column already known absent (e.g. discovered on a previous cycle)
|
||||
check_batch_cost_instance._has_batch_processed_column = False
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
|
||||
# Build a fake batch response whose status triggers the completion branch
|
||||
mock_response = MagicMock()
|
||||
mock_response.status = "completed"
|
||||
mock_response.output_file_id = "file-output-123"
|
||||
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
|
||||
|
||||
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
|
||||
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
|
||||
return_value={"api_key": "sk-test"}
|
||||
)
|
||||
|
||||
mock_deployment = MagicMock()
|
||||
mock_deployment.litellm_params.custom_llm_provider = "openai"
|
||||
mock_deployment.litellm_params.model = "gpt-4"
|
||||
mock_deployment.model_info.model_dump.return_value = {}
|
||||
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
|
||||
|
||||
mock_file_content = MagicMock()
|
||||
mock_file_content.content = b'{"id":"req-1"}'
|
||||
|
||||
decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
|
||||
side_effect=[decoded_id, None],
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
|
||||
return_value="model-123",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
|
||||
return_value="batch-456",
|
||||
),
|
||||
patch(
|
||||
"litellm.files.main.afile_content",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_file_content,
|
||||
),
|
||||
patch(
|
||||
"litellm.batches.batch_utils._get_file_content_as_dictionary",
|
||||
return_value=[{"id": "req-1"}],
|
||||
),
|
||||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
|
||||
return_value=("gpt-4", "openai", None, None),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.Logging"
|
||||
) as mock_logging_cls,
|
||||
):
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
mock_logging_cls.return_value = mock_logging_obj
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
# The update must have been called — this is the core assertion.
|
||||
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
|
||||
"Expected update() to be called exactly once for the completed job"
|
||||
)
|
||||
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
|
||||
assert "batch_processed" not in update_data, (
|
||||
"update() must NOT include batch_processed when column is absent"
|
||||
)
|
||||
assert update_data["status"] == "complete"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_primary_path_completion_update_includes_batch_processed(
|
||||
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""When batch_processed column IS present, completion update must set it to True.
|
||||
|
||||
This is the symmetric counterpart to test_fallback_completion_update_omits_batch_processed
|
||||
and proves the conditional on _has_batch_processed_column governs the update data.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
mock_job = MagicMock()
|
||||
mock_job.id = "job-primary-1"
|
||||
mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
|
||||
mock_job.created_by = "user-1"
|
||||
|
||||
assert check_batch_cost_instance._has_batch_processed_column is True
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status = "completed"
|
||||
mock_response.output_file_id = "file-output-123"
|
||||
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
|
||||
|
||||
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
|
||||
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
|
||||
return_value={"api_key": "sk-test"}
|
||||
)
|
||||
|
||||
mock_deployment = MagicMock()
|
||||
mock_deployment.litellm_params.custom_llm_provider = "openai"
|
||||
mock_deployment.litellm_params.model = "gpt-4"
|
||||
mock_deployment.model_info.model_dump.return_value = {}
|
||||
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
|
||||
|
||||
mock_file_content = MagicMock()
|
||||
mock_file_content.content = b'{"id":"req-1"}'
|
||||
|
||||
decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
|
||||
side_effect=[decoded_id, None],
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
|
||||
return_value="model-123",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
|
||||
return_value="batch-456",
|
||||
),
|
||||
patch(
|
||||
"litellm.files.main.afile_content",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_file_content,
|
||||
),
|
||||
patch(
|
||||
"litellm.batches.batch_utils._get_file_content_as_dictionary",
|
||||
return_value=[{"id": "req-1"}],
|
||||
),
|
||||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
|
||||
return_value=("gpt-4", "openai", None, None),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.Logging"
|
||||
) as mock_logging_cls,
|
||||
):
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
mock_logging_cls.return_value = mock_logging_obj
|
||||
|
||||
await check_batch_cost_instance.check_batch_cost()
|
||||
|
||||
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
|
||||
"Expected update() to be called exactly once for the completed job"
|
||||
)
|
||||
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
|
||||
assert update_data["batch_processed"] is True, (
|
||||
"update() must include batch_processed=True when column is present"
|
||||
)
|
||||
assert update_data["status"] == "complete"
|
||||
|
|
@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
|
||||
|
||||
|
|
@ -63,21 +64,46 @@ class TestCheckResponsesCost:
|
|||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""Test check_responses_cost when there are no jobs to process"""
|
||||
# Mock empty job list
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify find_many was called with pagination params
|
||||
find_many_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
|
||||
assert find_many_call[1]["where"] == {
|
||||
"status": {"in": ["queued", "in_progress"]},
|
||||
"file_purpose": "response",
|
||||
}
|
||||
assert find_many_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
|
||||
assert find_many_call[1]["order"] == {"created_at": "asc"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_stale_managed_objects(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""Stale rows (older than cutoff) are bulk-updated to stale_expired before polling."""
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=5
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
# Should not raise any errors
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify find_many was called with correct parameters
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
||||
where={
|
||||
"status": {"in": ["queued", "in_progress"]},
|
||||
"file_purpose": "response",
|
||||
}
|
||||
)
|
||||
# The first update_many call should be the stale-row cleanup scoped to "response"
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
stale_call = calls[0]
|
||||
assert stale_call[1]["data"] == {"status": "stale_expired"}
|
||||
where = stale_call[1]["where"]
|
||||
assert where["file_purpose"] == "response"
|
||||
assert "stale_expired" in where["status"]["not_in"]
|
||||
assert "created_at" in where
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_completed_response(
|
||||
|
|
@ -89,6 +115,7 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_123"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-123"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_123"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
|
|
@ -108,8 +135,9 @@ class TestCheckResponsesCost:
|
|||
),
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check with mocked litellm.aget_responses
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -117,11 +145,12 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify the job was marked as completed
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
|
||||
call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
|
||||
assert call_args[1]["data"]["status"] == "completed"
|
||||
assert call_args[1]["where"]["id"]["in"] == ["job-123"]
|
||||
# calls[0] = stale cleanup, calls[1] = job completion
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 2
|
||||
completion_call = calls[1]
|
||||
assert completion_call[1]["data"]["status"] == "completed"
|
||||
assert completion_call[1]["where"]["id"]["in"] == ["job-123"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_failed_response(
|
||||
|
|
@ -133,6 +162,7 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_456"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-456"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_456"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
|
|
@ -148,8 +178,9 @@ class TestCheckResponsesCost:
|
|||
usage=None,
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -157,10 +188,10 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify the job was marked as completed (even though response failed)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
|
||||
call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
|
||||
assert call_args[1]["data"]["status"] == "completed"
|
||||
# calls[0] = stale cleanup, calls[1] = job completion
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 2
|
||||
assert calls[1][1]["data"]["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_cancelled_response(
|
||||
|
|
@ -172,6 +203,7 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_789"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-789"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_789"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
|
|
@ -187,8 +219,9 @@ class TestCheckResponsesCost:
|
|||
usage=None,
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -196,8 +229,10 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify the job was marked as completed
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
|
||||
# calls[0] = stale cleanup, calls[1] = job completion
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 2
|
||||
assert calls[1][1]["data"]["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_in_progress_response(
|
||||
|
|
@ -209,6 +244,7 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_in_progress"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-in-progress"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_in_progress"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
|
|
@ -224,8 +260,9 @@ class TestCheckResponsesCost:
|
|||
usage=None,
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -233,8 +270,10 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify no updates were made (response still in progress)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
|
||||
# Only the stale-cleanup call should have fired — no completion update
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"] == {"status": "stale_expired"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_queued_response(
|
||||
|
|
@ -246,6 +285,7 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_queued"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-queued"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_queued"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
|
|
@ -261,8 +301,9 @@ class TestCheckResponsesCost:
|
|||
usage=None,
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -270,8 +311,10 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify no updates were made (response still queued)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
|
||||
# Only the stale-cleanup call should have fired — no completion update
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"] == {"status": "stale_expired"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_exception(
|
||||
|
|
@ -283,13 +326,15 @@ class TestCheckResponsesCost:
|
|||
mock_job.unified_object_id = "resp_test_error"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-error"
|
||||
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_error"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check with mocked exception
|
||||
with patch(
|
||||
|
|
@ -300,8 +345,10 @@ class TestCheckResponsesCost:
|
|||
# Should not raise, just skip the job
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify no updates were made (job was skipped due to error)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
|
||||
# Only the stale-cleanup call should have fired — no completion update
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"] == {"status": "stale_expired"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_multiple_jobs(
|
||||
|
|
@ -313,16 +360,19 @@ class TestCheckResponsesCost:
|
|||
mock_job1.unified_object_id = "resp_test_1"
|
||||
mock_job1.created_by = "user1"
|
||||
mock_job1.id = "job-1"
|
||||
mock_job1.file_object = {"model": "gpt-4o", "id": "resp_test_1"}
|
||||
|
||||
mock_job2 = MagicMock()
|
||||
mock_job2.unified_object_id = "resp_test_2"
|
||||
mock_job2.created_by = "user2"
|
||||
mock_job2.id = "job-2"
|
||||
mock_job2.file_object = {"model": "gpt-4o", "id": "resp_test_2"}
|
||||
|
||||
mock_job3 = MagicMock()
|
||||
mock_job3.unified_object_id = "resp_test_3"
|
||||
mock_job3.created_by = "user3"
|
||||
mock_job3.id = "job-3"
|
||||
mock_job3.file_object = {"model": "gpt-4o", "id": "resp_test_3"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job1, mock_job2, mock_job3]
|
||||
|
|
@ -364,8 +414,9 @@ class TestCheckResponsesCost:
|
|||
),
|
||||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
# Run the check
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
|
|
@ -373,10 +424,41 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# Verify only the 2 completed jobs were marked as complete
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
|
||||
call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
|
||||
assert len(call_args[1]["where"]["id"]["in"]) == 2
|
||||
assert "job-1" in call_args[1]["where"]["id"]["in"]
|
||||
assert "job-3" in call_args[1]["where"]["id"]["in"]
|
||||
assert "job-2" not in call_args[1]["where"]["id"]["in"]
|
||||
# calls[0] = stale cleanup, calls[1] = completion of 2 finished jobs
|
||||
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
assert len(calls) == 2
|
||||
completion_call = calls[1]
|
||||
assert len(completion_call[1]["where"]["id"]["in"]) == 2
|
||||
assert "job-1" in completion_call[1]["where"]["id"]["in"]
|
||||
assert "job-3" in completion_call[1]["where"]["id"]["in"]
|
||||
assert "job-2" not in completion_call[1]["where"]["id"]["in"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_no_model_in_file_object(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""When file_object has no 'model' key, model_name is None and metadata skips model fields."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_no_model"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-no-model"
|
||||
mock_job.file_object = {} # no "model" key → model_name=None branch
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status = "completed"
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# aget_responses should be called without model metadata
|
||||
call_kwargs = mock_aget.call_args[1]
|
||||
assert "model" not in call_kwargs.get("litellm_metadata", {})
|
||||
assert "model_group" not in call_kwargs.get("litellm_metadata", {})
|
||||
|
|
|
|||
|
|
@ -124,4 +124,35 @@ class TestHuggingFaceEmbedding:
|
|||
assert "source_sentence" in request_data["inputs"]
|
||||
assert "sentences" in request_data["inputs"]
|
||||
assert request_data["inputs"]["source_sentence"] == input_text[0]
|
||||
assert request_data["inputs"]["sentences"] == input_text[1:]
|
||||
assert request_data["inputs"]["sentences"] == input_text[1:]
|
||||
|
||||
def test_extra_headers_forwarded_to_handler(self):
|
||||
"""extra_headers passed to litellm.embedding() should reach the HTTP post call."""
|
||||
input_text = ["hello world"]
|
||||
|
||||
litellm.embedding(
|
||||
model=self.model,
|
||||
input=input_text,
|
||||
input_type="embed",
|
||||
extra_headers={"X-HF-Bill-To": "my-org"},
|
||||
)
|
||||
|
||||
self.mock_http.assert_called_once()
|
||||
call_kwargs = self.mock_http.call_args
|
||||
sent_headers = call_kwargs[1].get("headers", {})
|
||||
assert sent_headers.get("X-HF-Bill-To") == "my-org"
|
||||
|
||||
def test_no_extra_headers_uses_defaults(self):
|
||||
"""When no extra_headers are provided, default headers should not include custom ones."""
|
||||
input_text = ["hello world"]
|
||||
|
||||
litellm.embedding(
|
||||
model=self.model,
|
||||
input=input_text,
|
||||
input_type="embed",
|
||||
)
|
||||
|
||||
self.mock_http.assert_called_once()
|
||||
call_kwargs = self.mock_http.call_args
|
||||
sent_headers = call_kwargs[1].get("headers", {})
|
||||
assert "X-HF-Bill-To" not in sent_headers
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue