mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test: move tests/test_litellm root and small trees into tests/unit (#43186)
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng-jiang
GitHub
parent
694783ebbe
commit
f6882246d4
248 changed files with 5722 additions and 5685 deletions
|
|
@ -14,12 +14,12 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
*.md | *.mdx) : ;;
|
||||
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py)
|
||||
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py)
|
||||
has_mcp_dependencies=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
tests/e2e/*/*.py) : ;;
|
||||
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
|
||||
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
|
||||
has_provider_harness=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ legacy_flags=(
|
|||
enterprise-package
|
||||
enterprise-routing
|
||||
mcp-integration
|
||||
misc
|
||||
proxy-db-auth-checks
|
||||
proxy-db-budgets
|
||||
proxy-db-custom-logging
|
||||
|
|
@ -22,6 +23,7 @@ legacy_flags=(
|
|||
proxy-db-proxy-utils
|
||||
proxy-extras
|
||||
proxy-infra
|
||||
responses-caching-types
|
||||
)
|
||||
|
||||
legacy_paths() {
|
||||
|
|
@ -36,6 +38,7 @@ legacy_paths() {
|
|||
echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py
|
||||
echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;;
|
||||
enterprise-routing)
|
||||
echo tests/unit/google_genai
|
||||
echo tests/unit/enterprise/enterprise_callbacks/send_emails
|
||||
echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
|
||||
echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py
|
||||
|
|
@ -48,9 +51,28 @@ legacy_paths() {
|
|||
echo tests/unit/enterprise/proxy/test_managed_files_access_check.py
|
||||
echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;;
|
||||
mcp-integration)
|
||||
echo tests/unit/experimental_mcp_client
|
||||
echo tests/unit/proxy/_experimental/mcp_server
|
||||
echo tests/unit/responses/mcp
|
||||
echo tests/mcp_tests/test_proxy_mcp_e2e.py ;;
|
||||
misc)
|
||||
find tests/unit -maxdepth 1 -name 'test_*.py'
|
||||
echo tests/unit/test_router
|
||||
echo tests/unit/a2a_protocol
|
||||
echo tests/unit/batches
|
||||
echo tests/unit/chat_completions
|
||||
echo tests/unit/completion_extras
|
||||
echo tests/unit/containers
|
||||
echo tests/unit/embeddings
|
||||
echo tests/unit/endpoints
|
||||
echo tests/unit/files
|
||||
echo tests/unit/images
|
||||
echo tests/unit/interactions
|
||||
echo tests/unit/messages
|
||||
echo tests/unit/rag
|
||||
echo tests/unit/rerank_api
|
||||
echo tests/unit/vector_stores
|
||||
echo tests/unit/videos ;;
|
||||
proxy-db-auth-checks)
|
||||
echo tests/unit/proxy/auth/test_auth_checks.py
|
||||
echo tests/unit/proxy/auth/test_user_api_key_auth.py
|
||||
|
|
@ -113,6 +135,7 @@ legacy_paths() {
|
|||
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
|
||||
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
|
||||
proxy-infra) echo tests/unit/gateway ;;
|
||||
responses-caching-types) echo tests/unit/types ;;
|
||||
*) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;;
|
||||
esac
|
||||
}
|
||||
|
|
|
|||
|
|
@ -341,6 +341,7 @@ workflows:
|
|||
flag:
|
||||
- enterprise-package
|
||||
- proxy-infra
|
||||
- responses-caching-types
|
||||
- proxy-db-auth-checks
|
||||
- proxy-db-jwt-and-keys
|
||||
- proxy-db-proxy-server-core
|
||||
|
|
@ -353,6 +354,13 @@ workflows:
|
|||
- proxy-db-endpoints-and-responses
|
||||
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
|
||||
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
|
||||
- unit:
|
||||
name: unit-misc
|
||||
flag: misc
|
||||
shards: 2
|
||||
reruns: 2
|
||||
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
|
||||
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
|
||||
- unit:
|
||||
name: unit-proxy-db-proxy-utils
|
||||
flag: proxy-db-proxy-utils
|
||||
|
|
|
|||
4
.github/merge-smoke-tests.json
vendored
4
.github/merge-smoke-tests.json
vendored
|
|
@ -5,8 +5,8 @@
|
|||
"CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport",
|
||||
"MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key",
|
||||
"MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]",
|
||||
"COST-EXPLICIT": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
|
||||
"COST-ZERO": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
|
||||
"COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
|
||||
"COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
|
||||
"LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
|
||||
"LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
|
||||
"CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
|
||||
|
|
|
|||
4
.github/workflows/test-redis-compat.yml
vendored
4
.github/workflows/test-redis-compat.yml
vendored
|
|
@ -10,7 +10,7 @@ on:
|
|||
- "litellm/_redis_credential_provider.py"
|
||||
- "litellm/caching/redis_cache.py"
|
||||
- "litellm/caching/evicted_client_closer.py"
|
||||
- "tests/test_litellm/test_redis.py"
|
||||
- "tests/unit/test_redis.py"
|
||||
- "tests/local_testing/test_caching.py"
|
||||
- "tests/test_litellm/caching/test_redis_connection_pool.py"
|
||||
- "tests/test_litellm/caching/test_redis_cluster_cache.py"
|
||||
|
|
@ -84,7 +84,7 @@ jobs:
|
|||
run: |
|
||||
redis-server --version
|
||||
uv run --no-sync pytest \
|
||||
tests/test_litellm/test_redis.py \
|
||||
tests/unit/test_redis.py \
|
||||
tests/test_litellm/caching/test_redis_connection_pool.py \
|
||||
tests/test_litellm/caching/test_redis_cluster_cache.py \
|
||||
tests/test_litellm/caching/test_evicted_client_closer.py \
|
||||
|
|
|
|||
20
.github/workflows/test-unit.yml
vendored
20
.github/workflows/test-unit.yml
vendored
|
|
@ -52,7 +52,7 @@ jobs:
|
|||
include:
|
||||
- shard: mcp-integration
|
||||
artifact-name: mcp-integration
|
||||
test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client"
|
||||
test-path: "tests/mcp_tests"
|
||||
unit-flag: mcp-integration
|
||||
workers: 2
|
||||
reruns: 0
|
||||
|
|
@ -70,7 +70,6 @@ jobs:
|
|||
- shard: enterprise-routing
|
||||
artifact-name: enterprise-routing
|
||||
test-path: >-
|
||||
tests/test_litellm/google_genai
|
||||
tests/test_litellm/router_utils
|
||||
tests/test_litellm/router_strategy
|
||||
unit-flag: enterprise-routing
|
||||
|
|
@ -106,26 +105,13 @@ jobs:
|
|||
- shard: misc
|
||||
artifact-name: misc
|
||||
test-path: >-
|
||||
tests/test_litellm/batches
|
||||
tests/test_litellm/secret_managers
|
||||
tests/test_litellm/a2a_protocol
|
||||
tests/test_litellm/chat_completions
|
||||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/endpoints
|
||||
tests/test_litellm/files
|
||||
tests/test_litellm/images
|
||||
tests/test_litellm/interactions
|
||||
tests/test_litellm/messages
|
||||
tests/test_litellm/embeddings
|
||||
tests/test_litellm/ocr
|
||||
tests/test_litellm/passthrough
|
||||
tests/test_litellm/rag
|
||||
tests/test_litellm/rerank_api
|
||||
tests/test_litellm/rust_bridge
|
||||
tests/test_litellm/vector_stores
|
||||
tests/test_litellm/videos
|
||||
tests/test_litellm/test_*.py
|
||||
unit-flag: misc
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
|
|
@ -243,7 +229,7 @@ jobs:
|
|||
test-path: >-
|
||||
tests/test_litellm/responses
|
||||
tests/test_litellm/caching
|
||||
tests/test_litellm/types
|
||||
unit-flag: responses-caching-types
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
|
|
|
|||
4
Makefile
4
Makefile
|
|
@ -332,10 +332,10 @@ test-unit-core-utils: install-test-deps
|
|||
$(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20
|
||||
|
||||
test-unit-other: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-root: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
# Proxy unit tests (tests/unit/proxy split alphabetically)
|
||||
test-proxy-unit-a: install-test-deps
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ from tests._vcr_redis_persister import (
|
|||
# network call entirely, so skip tests record nothing (NOOP) and passing tests
|
||||
# stop carrying a volatile github episode. This matches the established idiom in
|
||||
# the unit-test suite, which sets the same flag (see e.g.
|
||||
# tests/test_litellm/test_cost_calculator.py). ``setdefault`` so an explicit
|
||||
# tests/unit/test_cost_calculator.py). ``setdefault`` so an explicit
|
||||
# override still wins.
|
||||
os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
|
||||
|
|
|
|||
|
|
@ -13,15 +13,16 @@ def check_for_litellm_module_deletion(base_dir):
|
|||
del sys.modules[module]
|
||||
"""
|
||||
problematic_files = []
|
||||
test_dir = os.path.join(base_dir, "test_litellm")
|
||||
candidate_dirs = [os.path.join(base_dir, name) for name in ("test_litellm", "unit")]
|
||||
test_dirs = [test_dir for test_dir in candidate_dirs if os.path.exists(test_dir)]
|
||||
|
||||
if not os.path.exists(test_dir):
|
||||
print(f"Warning: Directory {test_dir} does not exist.")
|
||||
if not test_dirs:
|
||||
print(f"Warning: None of {candidate_dirs} exist.")
|
||||
return []
|
||||
|
||||
print(f"Checking directory: {test_dir}")
|
||||
print(f"Checking directories: {test_dirs}")
|
||||
|
||||
for root, _, files in os.walk(test_dir):
|
||||
for root, _, files in (entry for test_dir in test_dirs for entry in os.walk(test_dir)):
|
||||
for file in files:
|
||||
if file.endswith(".py"):
|
||||
file_path = os.path.join(root, file)
|
||||
|
|
@ -173,7 +174,7 @@ def main():
|
|||
f"This can cause import issues and test failures. Files: {problematic_files}"
|
||||
)
|
||||
else:
|
||||
print("✓ No litellm module deletion patterns found in test_litellm directory.")
|
||||
print("✓ No litellm module deletion patterns found in tests/test_litellm or tests/unit.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ def get_all_functions_called_in_tests(base_dir):
|
|||
specifically in files containing the word 'router'.
|
||||
"""
|
||||
called_functions = set()
|
||||
test_dirs = ["local_testing", "router_unit_tests", "test_litellm"]
|
||||
test_dirs = ["local_testing", "router_unit_tests", "test_litellm", "unit"]
|
||||
|
||||
for test_dir in test_dirs:
|
||||
dir_path = os.path.join(base_dir, test_dir)
|
||||
|
|
|
|||
|
|
@ -277,4 +277,4 @@ class BaseSkillsAPITest(ABC):
|
|||
#
|
||||
# Transformation logic (URL construction, headers, request/response parsing) is
|
||||
# covered by unit tests in:
|
||||
# tests/test_litellm/test_anthropic_skills_transformation.py
|
||||
# tests/unit/test_anthropic_skills_transformation.py
|
||||
|
|
|
|||
|
|
@ -1,387 +0,0 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
import litellm.batches.batch_utils as bu
|
||||
from litellm.types.llms.openai import Batch
|
||||
|
||||
GROUNDED_USAGE_METADATA = {
|
||||
"promptTokenCount": 19,
|
||||
"candidatesTokenCount": 59,
|
||||
"thoughtsTokenCount": 406,
|
||||
"toolUsePromptTokenCount": 73,
|
||||
"totalTokenCount": 557,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}],
|
||||
"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}],
|
||||
"toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}],
|
||||
"trafficType": "ON_DEMAND",
|
||||
}
|
||||
PASSTHROUGH_OUTPUT_URI = (
|
||||
"gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/"
|
||||
"predictions.jsonl"
|
||||
)
|
||||
UNGROUNDED_USAGE_METADATA = {
|
||||
"promptTokenCount": 20,
|
||||
"candidatesTokenCount": 48,
|
||||
"thoughtsTokenCount": 195,
|
||||
"toolUsePromptTokenCount": 73,
|
||||
"totalTokenCount": 336,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}],
|
||||
"trafficType": "ON_DEMAND",
|
||||
}
|
||||
|
||||
|
||||
def _batch(output_file_id: str) -> Batch:
|
||||
return Batch(
|
||||
id="b",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="f",
|
||||
object="batch",
|
||||
status="completed",
|
||||
output_file_id=output_file_id,
|
||||
)
|
||||
|
||||
|
||||
def _vertex_jsonl(rows: list[dict]) -> bytes:
|
||||
return "\n".join(json.dumps(row) for row in rows).encode()
|
||||
|
||||
|
||||
def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict:
|
||||
return {
|
||||
"id": f"batch_req_{custom_id}",
|
||||
"custom_id": custom_id,
|
||||
"response": {
|
||||
"status_code": 200,
|
||||
"request_id": custom_id,
|
||||
"body": {
|
||||
"id": f"chatcmpl-{custom_id}",
|
||||
"object": "chat.completion",
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
},
|
||||
},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"):
|
||||
candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}
|
||||
grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {}
|
||||
response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata}
|
||||
return {
|
||||
"request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]},
|
||||
"status": "",
|
||||
"response": {**response, **({"modelVersion": model_version} if model_version else {})},
|
||||
"processed_time": "2026-09-23T19:02:00.000+00:00",
|
||||
}
|
||||
|
||||
|
||||
def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list:
|
||||
import litellm.cost_calculator as cc
|
||||
|
||||
calls: list = []
|
||||
|
||||
def _calc(**kw):
|
||||
calls.append(kw)
|
||||
return (prompt_cost, completion_cost)
|
||||
|
||||
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
||||
return calls
|
||||
|
||||
|
||||
def test_vertex_native_cost_bills_embedding_rows(monkeypatch):
|
||||
monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7})
|
||||
rows = [
|
||||
{
|
||||
"key": "id_1",
|
||||
"status": "",
|
||||
"request": {"content": {"parts": [{"text": "hello world"}]}},
|
||||
"response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}},
|
||||
},
|
||||
{
|
||||
"key": "id_2",
|
||||
"status": "",
|
||||
"request": {"content": {"parts": [{"text": "hello"}]}},
|
||||
"response": {"embedding": {"values": [0.3]}, "tokenCount": "3"},
|
||||
},
|
||||
{"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}},
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 1)
|
||||
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5)
|
||||
assert result.cost == pytest.approx(5 * 1e-7)
|
||||
assert result.models == ["gemini-embedding-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False),
|
||||
]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
||||
)
|
||||
_capture_cost_calls(monkeypatch)
|
||||
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert result.successful_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
||||
)
|
||||
_capture_cost_calls(monkeypatch)
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
assert result.successful_requests == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)]
|
||||
|
||||
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
||||
return _vertex_jsonl(raw_rows)
|
||||
|
||||
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
|
||||
result = await bu._handle_completed_batch(
|
||||
_batch(PASSTHROUGH_OUTPUT_URI),
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert result.cost == pytest.approx(1.0)
|
||||
assert result.usage.total_tokens == 557
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)
|
||||
ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False)
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash")
|
||||
|
||||
grounded_usage, ungrounded_usage = (call["usage"] for call in calls)
|
||||
assert grounded_usage.prompt_tokens == 19
|
||||
assert grounded_usage.completion_tokens == 59 + 406
|
||||
assert grounded_usage.completion_tokens_details.reasoning_tokens == 406
|
||||
assert ungrounded_usage.prompt_tokens == 20 + 73
|
||||
assert ungrounded_usage.completion_tokens == 48 + 195
|
||||
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
||||
19 + 93,
|
||||
465 + 243,
|
||||
557 + 336,
|
||||
)
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
||||
assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert result.successful_requests == 3
|
||||
assert result.usage.total_tokens == 557 + 336 + 336
|
||||
|
||||
|
||||
def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch):
|
||||
_capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
{"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"},
|
||||
{"request": {"contents": []}, "response": {"candidates": []}},
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (1, 2)
|
||||
assert result.usage.total_tokens == 557
|
||||
|
||||
|
||||
def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0)
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
|
||||
bu.calculate_vertex_ai_batch_cost_and_usage(
|
||||
[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
"gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6}
|
||||
|
||||
await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash"]
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
|
||||
|
||||
def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
{"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}},
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (1, 1)
|
||||
assert result.usage.total_tokens == 557
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"])
|
||||
def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model)
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model]
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.usage.total_tokens == 557 + 336
|
||||
|
||||
|
||||
def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices():
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash")
|
||||
without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None)
|
||||
|
||||
twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info)
|
||||
both = bu.calculate_vertex_ai_batch_cost_and_usage(
|
||||
[with_version, without_version], "vertex_ai/*", model_info=deployment_model_info
|
||||
)
|
||||
|
||||
assert twin.cost > 0
|
||||
assert both.cost == pytest.approx(2 * twin.cost)
|
||||
assert (both.successful_requests, both.failed_requests) == (2, 0)
|
||||
|
||||
|
||||
def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch):
|
||||
import litellm.cost_calculator as cc
|
||||
|
||||
def _calc(**kw):
|
||||
if kw["model"] == "gemini-unpriced":
|
||||
raise ValueError("no pricing")
|
||||
return (0.5, 0.25)
|
||||
|
||||
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
||||
|
||||
assert result.cost == pytest.approx(0.75)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.usage.total_tokens == 557 + 336
|
||||
assert result.models == ["gemini-unpriced", "gemini-2.5-flash"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert calls == []
|
||||
assert (result.successful_requests, result.failed_requests) == (0, 1)
|
||||
|
|
@ -1,117 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.chat_completions import dispatch
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import (
|
||||
LiteLLMChatCompletionsRequest,
|
||||
NativeAcompletion,
|
||||
NativeCompletion,
|
||||
)
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_completion_calls_keep_the_python_result() -> None:
|
||||
sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
|
||||
assert isinstance(sync_response, ModelResponse)
|
||||
assert isinstance(async_response, ModelResponse)
|
||||
assert sync_response.choices[0].message.content == "ok"
|
||||
assert async_response.choices[0].message.content == "ok"
|
||||
|
||||
|
||||
def test_sync_completion_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
assert request.model == "test-model"
|
||||
assert request.messages == MESSAGES
|
||||
assert request.custom_llm_provider == "openai"
|
||||
assert request.stream is True
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{"custom_llm_provider": "openai", "stream": True},
|
||||
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_completion_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = ModelResponse()
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_acompletion_marker_bypasses_native() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
pytest.fail("acompletion's inner completion call must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{"custom_llm_provider": "openai", "acompletion": True},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
@ -14,6 +14,7 @@ from pathlib import Path
|
|||
from types import SimpleNamespace
|
||||
import httpx
|
||||
import pytest
|
||||
from pytest_socket import _remove_restrictions
|
||||
|
||||
import asyncio
|
||||
|
||||
|
|
@ -509,6 +510,14 @@ def setup_and_teardown():
|
|||
print(f"[conftest] Module teardown complete (worker: {worker_id or 'master'})")
|
||||
|
||||
|
||||
def pytest_collectstart():
|
||||
_remove_restrictions()
|
||||
|
||||
|
||||
def pytest_runtest_setup():
|
||||
_remove_restrictions()
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
"""
|
||||
Customize test collection order.
|
||||
|
|
|
|||
|
|
@ -7,10 +7,6 @@ the litellm_responses bridge provider, which calls litellm.responses() internall
|
|||
|
||||
import os
|
||||
|
||||
from litellm.interactions.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesInteractionsConfig,
|
||||
)
|
||||
from litellm.types.interactions import Turn
|
||||
from tests.test_litellm.interactions.base_interactions_test import (
|
||||
BaseInteractionsTest,
|
||||
)
|
||||
|
|
@ -30,71 +26,3 @@ class TestLiteLLMResponsesBridge(BaseInteractionsTest):
|
|||
def get_api_key(self) -> str:
|
||||
"""Return the OpenAI API key from environment."""
|
||||
return os.getenv("OPENAI_API_KEY", "")
|
||||
|
||||
|
||||
class TestBridgeInputTransformation:
|
||||
"""Regression tests for translating Interactions input into Responses API input.
|
||||
|
||||
The bridge used to pass Google content parts through raw ({"type": "text"}),
|
||||
which the Responses API rejects with a 400, and it dropped the role encoded
|
||||
in step types and in the legacy "model" turn role.
|
||||
"""
|
||||
|
||||
def test_step_input_maps_roles_and_content_types(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[
|
||||
{"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]},
|
||||
{"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]},
|
||||
{"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]},
|
||||
]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]},
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]},
|
||||
]
|
||||
|
||||
def test_legacy_turn_input_maps_model_role_to_assistant(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[
|
||||
{"role": "user", "content": [{"type": "text", "text": "I like apples."}]},
|
||||
{"role": "model", "content": [{"type": "text", "text": "I like oranges."}]},
|
||||
]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]},
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]},
|
||||
]
|
||||
|
||||
def test_turn_pydantic_model_with_string_content(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[Turn(role="model", content="I like oranges.")]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}
|
||||
]
|
||||
|
||||
def test_string_input_passes_through(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello")
|
||||
assert transformed == "Hello"
|
||||
|
||||
def test_content_list_input_becomes_single_user_message(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[{"type": "text", "text": "Hello"}, "world"]
|
||||
)
|
||||
assert transformed == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Hello"},
|
||||
{"type": "input_text", "text": "world"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
def test_non_text_content_passes_through_unchanged(self):
|
||||
image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"}
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[{"type": "user_input", "content": [image_part]}]
|
||||
)
|
||||
assert transformed == [{"role": "user", "content": [image_part]}]
|
||||
|
|
|
|||
|
|
@ -1,155 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.messages import dispatch
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.messages.entrypoints import (
|
||||
LiteLLMMessagesRequest,
|
||||
NativeAmessages,
|
||||
NativeMessages,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_anthropic_messages_keeps_the_python_result() -> None:
|
||||
response: Final = await litellm.anthropic_messages(
|
||||
model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok"
|
||||
)
|
||||
|
||||
assert isinstance(response, dict)
|
||||
content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", []))
|
||||
assert content[0]["text"] == "ok"
|
||||
|
||||
|
||||
def test_sync_messages_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
assert request.model == "claude-test"
|
||||
assert request.messages == MESSAGES
|
||||
assert request.max_tokens == 10
|
||||
assert request.custom_llm_provider == "anthropic"
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{
|
||||
"model": "claude-test",
|
||||
"messages": MESSAGES,
|
||||
"max_tokens": 10,
|
||||
"custom_llm_provider": "anthropic",
|
||||
},
|
||||
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_messages_binding_error_delegates_unchanged_to_python() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
pytest.fail("a call without max_tokens cannot project a request and must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_messages_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{"model": "claude-test", "messages": MESSAGES, "max_tokens": 10},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_is_async_marker_bypasses_native() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
pytest.fail("anthropic_messages' inner handler call must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{
|
||||
"model": "claude-test",
|
||||
"messages": MESSAGES,
|
||||
"max_tokens": 10,
|
||||
"custom_llm_provider": "anthropic",
|
||||
"is_async": True,
|
||||
},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
@ -2823,7 +2823,7 @@ def test_update_router_config_schema_includes_tag_routing_prefix():
|
|||
# UpdateRouterConfig before calling update_settings; a field missing here
|
||||
# causes model_dump(exclude_none=True) to silently drop it before
|
||||
# update_settings is ever called -- the same bug shape LIT-3152 fixed for
|
||||
# retry_policy (see tests/test_litellm/test_router_retry_policy_update.py).
|
||||
# retry_policy (see tests/unit/test_router_retry_policy_update.py).
|
||||
from litellm.types.router import UpdateRouterConfig
|
||||
|
||||
config = UpdateRouterConfig(tag_routing_prefix="route:")
|
||||
|
|
|
|||
|
|
@ -3,20 +3,13 @@ Unit tests for litellm.compress().
|
|||
"""
|
||||
|
||||
import os
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.compression.scoring.bm25 import bm25_score_messages
|
||||
from litellm.compression.scoring.embedding_scorer import embedding_score_messages
|
||||
from litellm.compression.content_detection import detect_content_type
|
||||
from litellm.compression.message_stubbing import extract_key, stub_message
|
||||
from litellm.compression.retrieval_tool import build_retrieval_tool
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
CALL_TYPE = CallTypes.completion
|
||||
ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -24,420 +17,26 @@ ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_bm25_relevance_ranking():
|
||||
query = "Fix the authentication bug in the login handler"
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "def login_handler(): authentication check bug fix",
|
||||
},
|
||||
{"role": "user", "content": "def render_template(name): css styling layout"},
|
||||
{"role": "user", "content": "def verify(): authentication token bug handler"},
|
||||
]
|
||||
scores = bm25_score_messages(query, messages)
|
||||
# Messages sharing query terms should score higher than unrelated ones
|
||||
assert scores[0] > scores[1]
|
||||
assert scores[2] > scores[1]
|
||||
|
||||
|
||||
def test_bm25_empty_query():
|
||||
scores = bm25_score_messages("", [{"role": "user", "content": "hello"}])
|
||||
assert scores == [0.0]
|
||||
|
||||
|
||||
def test_bm25_empty_messages():
|
||||
scores = bm25_score_messages("query", [])
|
||||
assert scores == []
|
||||
|
||||
|
||||
def test_bm25_empty_content():
|
||||
scores = bm25_score_messages("query", [{"role": "user", "content": ""}])
|
||||
assert scores == [0.0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Content detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_detect_code():
|
||||
code = """
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
def main():
|
||||
class Foo:
|
||||
pass
|
||||
return Foo()
|
||||
"""
|
||||
assert detect_content_type(code) == "code"
|
||||
|
||||
|
||||
def test_detect_json():
|
||||
assert detect_content_type('{"key": "value", "num": 42}') == "json"
|
||||
assert detect_content_type("[1, 2, 3]") == "json"
|
||||
|
||||
|
||||
def test_detect_text():
|
||||
assert detect_content_type("This is a plain text paragraph about dogs.") == "text"
|
||||
|
||||
|
||||
def test_detect_empty():
|
||||
assert detect_content_type("") == "text"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Message stubbing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_extract_key_with_filename():
|
||||
msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"}
|
||||
used: set = set()
|
||||
key = extract_key(msg, fallback_index=0, used_keys=used)
|
||||
assert key == "auth.py"
|
||||
|
||||
|
||||
def test_extract_key_fallback():
|
||||
msg = {"role": "user", "content": "Some random content without a filename"}
|
||||
used: set = set()
|
||||
key = extract_key(msg, fallback_index=5, used_keys=used)
|
||||
assert key == "message_5"
|
||||
|
||||
|
||||
def test_extract_key_duplicates():
|
||||
used: set = set()
|
||||
msg = {"role": "user", "content": "# auth.py\ncode here"}
|
||||
k1 = extract_key(msg, fallback_index=0, used_keys=used)
|
||||
k2 = extract_key(msg, fallback_index=1, used_keys=used)
|
||||
assert k1 == "auth.py"
|
||||
assert k2 == "auth.py_2"
|
||||
|
||||
|
||||
def test_stub_message():
|
||||
msg = {"role": "user", "content": "line1\nline2\nline3"}
|
||||
stubbed = stub_message(msg, "test_key")
|
||||
assert stubbed["role"] == "user"
|
||||
assert "test_key" in stubbed["content"]
|
||||
assert "litellm_content_retrieve" in stubbed["content"]
|
||||
assert "3 lines" in stubbed["content"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Retrieval tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_retrieval_tool_schema():
|
||||
tool = build_retrieval_tool(["auth.py", "utils.py"])
|
||||
assert tool["type"] == "function"
|
||||
assert tool["function"]["name"] == "litellm_content_retrieve"
|
||||
assert "key" in tool["function"]["parameters"]["properties"]
|
||||
assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [
|
||||
"auth.py",
|
||||
"utils.py",
|
||||
]
|
||||
assert tool["function"]["parameters"]["required"] == ["key"]
|
||||
|
||||
|
||||
def test_retrieval_tool_description_lists_keys():
|
||||
tool = build_retrieval_tool(["foo.py", "bar.js"])
|
||||
desc = tool["function"]["description"]
|
||||
assert "foo.py" in desc
|
||||
assert "bar.js" in desc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# compress() — end-to-end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_compress_below_trigger_passthrough():
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE)
|
||||
assert result["messages"] == messages
|
||||
assert result["cache"] == {}
|
||||
assert result["tools"] == []
|
||||
assert result["compression_ratio"] == 0.0
|
||||
assert result["original_tokens"] == result["compressed_tokens"]
|
||||
|
||||
|
||||
def test_compress_above_trigger():
|
||||
big_messages = [
|
||||
{"role": "system", "content": "You are a coding assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "# utils.py\n" + "def helper():\n pass\n" * 2000,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "# readme.md\n" + "This is documentation. " * 2000,
|
||||
},
|
||||
{"role": "user", "content": "Fix the bug in auth.py"},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
big_messages,
|
||||
model="gpt-4o",
|
||||
call_type=CALL_TYPE,
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
||||
assert result["compressed_tokens"] < result["original_tokens"]
|
||||
assert result["compression_ratio"] > 0
|
||||
assert len(result["cache"]) > 0
|
||||
assert len(result["tools"]) == 1
|
||||
assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve"
|
||||
|
||||
|
||||
def test_compress_anthropic_list_content_is_boundary_stable():
|
||||
messages = [
|
||||
{"role": "system", "content": [{"type": "text", "text": "System prompt"}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "# a.py\n" + "alpha " * 2000},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/a.png"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "# b.py\n" + "beta " * 2000},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/b.png"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Fix alpha bug in a.py"}],
|
||||
},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-20250514",
|
||||
call_type=ANTHROPIC_CALL_TYPE,
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
||||
assert result["compressed_tokens"] < result["original_tokens"]
|
||||
assert len(result["messages"]) == len(messages)
|
||||
assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages]
|
||||
assert len(result["cache"]) > 0
|
||||
assert len(result["tools"]) == 1
|
||||
assert result["tools"][0]["type"] == "custom"
|
||||
assert result["tools"][0]["name"] == "litellm_content_retrieve"
|
||||
assert "input_schema" in result["tools"][0]
|
||||
|
||||
|
||||
def test_compress_preserves_system_message():
|
||||
messages = [
|
||||
{"role": "system", "content": "System prompt. " * 500},
|
||||
{"role": "user", "content": "Large file content. " * 5000},
|
||||
{"role": "user", "content": "Fix the bug"},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
|
||||
)
|
||||
assert result["messages"][0]["role"] == "system"
|
||||
assert "System prompt" in result["messages"][0]["content"]
|
||||
|
||||
|
||||
def test_compress_preserves_last_user_message():
|
||||
messages = [
|
||||
{"role": "user", "content": "Big context " * 5000},
|
||||
{"role": "user", "content": "Fix the bug in auth.py"},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
|
||||
)
|
||||
last_user = [m for m in result["messages"] if m["role"] == "user"][-1]
|
||||
assert "Fix the bug in auth.py" in last_user["content"]
|
||||
|
||||
|
||||
def test_compress_preserves_last_assistant_message():
|
||||
messages = [
|
||||
{"role": "user", "content": "Big context " * 5000},
|
||||
{"role": "assistant", "content": "I'll help with that. " * 2000},
|
||||
{"role": "user", "content": "Now fix the bug"},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
|
||||
)
|
||||
assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"]
|
||||
assert len(assistant_msgs) >= 1
|
||||
# The last assistant message should be preserved (not stubbed)
|
||||
last_assistant = assistant_msgs[-1]
|
||||
assert "I'll help with that" in last_assistant["content"]
|
||||
|
||||
|
||||
def test_cache_keys_match_stubs():
|
||||
messages = [
|
||||
{"role": "user", "content": "# auth.py\n" + "code " * 5000},
|
||||
{"role": "user", "content": "Fix it"},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
|
||||
)
|
||||
if result["tools"]:
|
||||
tool_desc = result["tools"][0]["function"]["description"]
|
||||
for key in result["cache"]:
|
||||
assert key in tool_desc
|
||||
|
||||
|
||||
def test_compress_default_target():
|
||||
"""compression_target defaults to compression_trigger // 2."""
|
||||
messages = [
|
||||
{"role": "user", "content": "content " * 5000},
|
||||
{"role": "user", "content": "query"},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000
|
||||
)
|
||||
# Should have compressed — target = 1000
|
||||
assert result["compressed_tokens"] <= result["original_tokens"]
|
||||
|
||||
|
||||
def test_compress_nested_tool_result_extracts_text_only():
|
||||
messages = [
|
||||
{"role": "system", "content": [{"type": "text", "text": "System rules"}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "prefix"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": [
|
||||
{"type": "text", "text": "nested text fragment"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://example.com/secret-tool.png",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/top.png"},
|
||||
},
|
||||
{"type": "text", "text": " " + ("irrelevant " * 3000)},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "final query that must remain"}],
|
||||
},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-20250514",
|
||||
call_type=ANTHROPIC_CALL_TYPE,
|
||||
compression_trigger=500,
|
||||
compression_target=100,
|
||||
)
|
||||
|
||||
cached_text = " ".join(result["cache"].values())
|
||||
assert "nested text fragment" in cached_text
|
||||
assert "https://example.com/secret-tool.png" not in cached_text
|
||||
assert "https://example.com/top.png" not in cached_text
|
||||
|
||||
|
||||
def test_compress_default_call_type_is_completion():
|
||||
result = litellm.compress(
|
||||
messages=[
|
||||
{"role": "user", "content": "Large context " * 4000},
|
||||
{"role": "user", "content": "query"},
|
||||
],
|
||||
model="gpt-4o",
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
||||
assert result["compressed_tokens"] <= result["original_tokens"]
|
||||
assert isinstance(result["tools"], list)
|
||||
|
||||
|
||||
def test_compress_forwards_embedding_model_params(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_embedding_score_messages(
|
||||
query, messages, model, cache=None, embedding_model_params=None
|
||||
):
|
||||
captured["query"] = query
|
||||
captured["model"] = model
|
||||
captured["embedding_model_params"] = embedding_model_params
|
||||
return [0.0] * len(messages)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.compression.scoring.embedding_scorer.embedding_score_messages",
|
||||
fake_embedding_score_messages,
|
||||
)
|
||||
|
||||
result = litellm.compress(
|
||||
messages=[
|
||||
{"role": "user", "content": "Authentication code " * 2000},
|
||||
{"role": "user", "content": "Fix auth"},
|
||||
],
|
||||
model="gpt-4o",
|
||||
call_type=CALL_TYPE,
|
||||
compression_trigger=1000,
|
||||
embedding_model="text-embedding-3-small",
|
||||
embedding_model_params={"api_base": "https://example-embeddings.test"},
|
||||
)
|
||||
|
||||
assert result["compressed_tokens"] <= result["original_tokens"]
|
||||
assert captured["model"] == "text-embedding-3-small"
|
||||
assert captured["embedding_model_params"] == {
|
||||
"api_base": "https://example-embeddings.test"
|
||||
}
|
||||
|
||||
|
||||
def test_embedding_scorer_forwards_embedding_model_params(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
class _MockResponse:
|
||||
data = [
|
||||
{"embedding": [1.0, 0.0]},
|
||||
{"embedding": [1.0, 0.0]},
|
||||
{"embedding": [0.0, 1.0]},
|
||||
]
|
||||
|
||||
def fake_embedding(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return _MockResponse()
|
||||
|
||||
monkeypatch.setattr(litellm, "embedding", fake_embedding)
|
||||
|
||||
scores = embedding_score_messages(
|
||||
query="auth",
|
||||
messages=[
|
||||
{"role": "user", "content": "auth code"},
|
||||
{"role": "user", "content": "cooking recipe"},
|
||||
],
|
||||
model="text-embedding-3-small",
|
||||
embedding_model_params={"api_base": "https://example-embeddings.test"},
|
||||
)
|
||||
|
||||
assert len(scores) == 2
|
||||
assert captured["model"] == "text-embedding-3-small"
|
||||
assert captured["api_base"] == "https://example-embeddings.test"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Embedding scorer — integration test (skipped without API key)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -458,210 +57,3 @@ def test_embedding_scorer():
|
|||
)
|
||||
assert result["compression_ratio"] > 0
|
||||
assert len(result["cache"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"final_user_message, expected_content",
|
||||
[
|
||||
("How to cook?", "Unrelated cooking recipes "),
|
||||
("Fix auth", "Authentication code "),
|
||||
],
|
||||
)
|
||||
def test_simple_compression(final_user_message, expected_content):
|
||||
messages = [
|
||||
{"role": "user", "content": "Authentication code " * 2000},
|
||||
{"role": "user", "content": "Unrelated cooking recipes " * 2000},
|
||||
{"role": "user", "content": final_user_message},
|
||||
]
|
||||
result = litellm.compress(
|
||||
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
|
||||
)
|
||||
if expected_content == "Unrelated cooking recipes ":
|
||||
assert "Unrelated cooking recipes " in result["messages"][1]["content"]
|
||||
assert "Authentication code " not in result["messages"][0]["content"]
|
||||
elif expected_content == "Authentication code ":
|
||||
assert "Authentication code " in result["messages"][0]["content"]
|
||||
assert "Unrelated cooking recipes " not in result["messages"][1]["content"]
|
||||
else:
|
||||
raise ValueError(f"Unexpected expected_content: {expected_content}")
|
||||
|
||||
|
||||
def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch):
|
||||
compress_module = importlib.import_module("litellm.compression.compress")
|
||||
|
||||
def fake_bm25_score_messages(query, messages):
|
||||
assert "final query" in query
|
||||
assert len(messages) == 5
|
||||
# Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2)
|
||||
return [0.95, 0.01, 0.02, 0.8, 1.0]
|
||||
|
||||
def fake_token_counter(model, messages=None, text=None):
|
||||
if messages is not None:
|
||||
return 1000
|
||||
if text is None:
|
||||
return 0
|
||||
if "final query" in text:
|
||||
return 50
|
||||
if "assistant_tail" in text:
|
||||
return 20
|
||||
if "other_blob" in text:
|
||||
return 220
|
||||
if "tool_payload_relevant" in text:
|
||||
return 200
|
||||
if text == "":
|
||||
return 1
|
||||
return 10
|
||||
|
||||
monkeypatch.setattr(
|
||||
compress_module, "bm25_score_messages", fake_bm25_score_messages
|
||||
)
|
||||
monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "other_blob " * 300},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_drop",
|
||||
"name": "litellm_content_retrieve",
|
||||
"input": {"key": "message_1"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_drop",
|
||||
"content": [{"type": "text", "text": "tool_payload_relevant"}],
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "assistant_tail"},
|
||||
{"role": "user", "content": "final query"},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-20250514",
|
||||
call_type=ANTHROPIC_CALL_TYPE,
|
||||
compression_trigger=100,
|
||||
compression_target=280,
|
||||
)
|
||||
|
||||
# idx=1,2 should be dropped atomically (no orphan tool blocks left behind)
|
||||
assert len(result["messages"]) == 3
|
||||
assert result["messages"][0]["role"] == "user"
|
||||
assert "other_blob" in result["messages"][0]["content"]
|
||||
assert result["messages"][1]["content"] == "assistant_tail"
|
||||
assert result["messages"][2]["content"] == "final query"
|
||||
assert result["cache"] == {}
|
||||
|
||||
|
||||
def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch):
|
||||
compress_module = importlib.import_module("litellm.compression.compress")
|
||||
|
||||
def fake_bm25_score_messages(query, messages):
|
||||
assert "final query" in query
|
||||
assert len(messages) == 5
|
||||
# Prefer the tool exchange span over idx=0
|
||||
return [0.05, 0.01, 0.92, 0.8, 1.0]
|
||||
|
||||
def fake_token_counter(model, messages=None, text=None):
|
||||
if messages is not None:
|
||||
return 1000
|
||||
if text is None:
|
||||
return 0
|
||||
if "final query" in text:
|
||||
return 50
|
||||
if "assistant_tail" in text:
|
||||
return 20
|
||||
if "other_blob" in text:
|
||||
return 220
|
||||
if "tool_payload_relevant" in text:
|
||||
return 200
|
||||
if text == "":
|
||||
return 1
|
||||
return 10
|
||||
|
||||
monkeypatch.setattr(
|
||||
compress_module, "bm25_score_messages", fake_bm25_score_messages
|
||||
)
|
||||
monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "other_blob " * 300},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_keep",
|
||||
"name": "litellm_content_retrieve",
|
||||
"input": {"key": "message_1"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_keep",
|
||||
"content": [{"type": "text", "text": "tool_payload_relevant"}],
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "assistant_tail"},
|
||||
{"role": "user", "content": "final query"},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-20250514",
|
||||
call_type=ANTHROPIC_CALL_TYPE,
|
||||
compression_trigger=100,
|
||||
compression_target=280,
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == 5
|
||||
assert result["messages"][1]["role"] == "assistant"
|
||||
assert result["messages"][2]["role"] == "user"
|
||||
# idx=0 should be compressed instead
|
||||
assert "litellm_content_retrieve" in result["messages"][0]["content"]
|
||||
assert len(result["cache"]) == 1
|
||||
|
||||
|
||||
def test_compress_anthropic_malformed_tool_sequence_passes_through():
|
||||
messages = [
|
||||
{"role": "user", "content": "other_blob " * 300},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_broken",
|
||||
"name": "litellm_content_retrieve",
|
||||
"input": {"key": "message_1"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]},
|
||||
{"role": "user", "content": "final query"},
|
||||
]
|
||||
|
||||
result = litellm.compress(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-20250514",
|
||||
call_type=ANTHROPIC_CALL_TYPE,
|
||||
compression_trigger=100,
|
||||
compression_target=280,
|
||||
)
|
||||
|
||||
assert result["messages"] == messages
|
||||
assert result["cache"] == {}
|
||||
assert result["tools"] == []
|
||||
assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence"
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -2072,3 +2072,348 @@ def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch):
|
|||
)
|
||||
assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2)
|
||||
assert result.usage.total_tokens == 15
|
||||
|
||||
|
||||
GROUNDED_USAGE_METADATA = {
|
||||
"promptTokenCount": 19,
|
||||
"candidatesTokenCount": 59,
|
||||
"thoughtsTokenCount": 406,
|
||||
"toolUsePromptTokenCount": 73,
|
||||
"totalTokenCount": 557,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}],
|
||||
"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}],
|
||||
"toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}],
|
||||
"trafficType": "ON_DEMAND",
|
||||
}
|
||||
|
||||
|
||||
PASSTHROUGH_OUTPUT_URI = (
|
||||
"gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/"
|
||||
"predictions.jsonl"
|
||||
)
|
||||
|
||||
|
||||
UNGROUNDED_USAGE_METADATA = {
|
||||
"promptTokenCount": 20,
|
||||
"candidatesTokenCount": 48,
|
||||
"thoughtsTokenCount": 195,
|
||||
"toolUsePromptTokenCount": 73,
|
||||
"totalTokenCount": 336,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}],
|
||||
"trafficType": "ON_DEMAND",
|
||||
}
|
||||
|
||||
|
||||
def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"):
|
||||
candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}
|
||||
grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {}
|
||||
response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata}
|
||||
return {
|
||||
"request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]},
|
||||
"status": "",
|
||||
"response": {**response, **({"modelVersion": model_version} if model_version else {})},
|
||||
"processed_time": "2026-09-23T19:02:00.000+00:00",
|
||||
}
|
||||
|
||||
|
||||
def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list:
|
||||
import litellm.cost_calculator as cc
|
||||
|
||||
calls: list = []
|
||||
|
||||
def _calc(**kw):
|
||||
calls.append(kw)
|
||||
return (prompt_cost, completion_cost)
|
||||
|
||||
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
||||
return calls
|
||||
|
||||
|
||||
def test_vertex_native_cost_bills_embedding_rows(monkeypatch):
|
||||
monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7})
|
||||
rows = [
|
||||
{
|
||||
"key": "id_1",
|
||||
"status": "",
|
||||
"request": {"content": {"parts": [{"text": "hello world"}]}},
|
||||
"response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}},
|
||||
},
|
||||
{
|
||||
"key": "id_2",
|
||||
"status": "",
|
||||
"request": {"content": {"parts": [{"text": "hello"}]}},
|
||||
"response": {"embedding": {"values": [0.3]}, "tokenCount": "3"},
|
||||
},
|
||||
{"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}},
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 1)
|
||||
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5)
|
||||
assert result.cost == pytest.approx(5 * 1e-7)
|
||||
assert result.models == ["gemini-embedding-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False),
|
||||
]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
||||
)
|
||||
_capture_cost_calls(monkeypatch)
|
||||
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert result.successful_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
||||
)
|
||||
_capture_cost_calls(monkeypatch)
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
assert result.successful_requests == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)]
|
||||
|
||||
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
||||
return _vertex_jsonl(raw_rows)
|
||||
|
||||
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
|
||||
result = await bu._handle_completed_batch(
|
||||
_batch(PASSTHROUGH_OUTPUT_URI),
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert result.cost == pytest.approx(1.0)
|
||||
assert result.usage.total_tokens == 557
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)
|
||||
ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False)
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash")
|
||||
|
||||
grounded_usage, ungrounded_usage = (call["usage"] for call in calls)
|
||||
assert grounded_usage.prompt_tokens == 19
|
||||
assert grounded_usage.completion_tokens == 59 + 406
|
||||
assert grounded_usage.completion_tokens_details.reasoning_tokens == 406
|
||||
assert ungrounded_usage.prompt_tokens == 20 + 73
|
||||
assert ungrounded_usage.completion_tokens == 48 + 195
|
||||
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
||||
19 + 93,
|
||||
465 + 243,
|
||||
557 + 336,
|
||||
)
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
||||
assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert result.successful_requests == 3
|
||||
assert result.usage.total_tokens == 557 + 336 + 336
|
||||
|
||||
|
||||
def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch):
|
||||
_capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
{"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"},
|
||||
{"request": {"contents": []}, "response": {"candidates": []}},
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (1, 2)
|
||||
assert result.usage.total_tokens == 557
|
||||
|
||||
|
||||
def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0)
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
|
||||
bu.calculate_vertex_ai_batch_cost_and_usage(
|
||||
[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
"gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6}
|
||||
|
||||
await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.5-flash",
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
|
||||
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
||||
|
||||
|
||||
def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash"]
|
||||
assert result.models == ["gemini-2.5-flash"]
|
||||
|
||||
|
||||
def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
{"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}},
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
||||
|
||||
assert (result.successful_requests, result.failed_requests) == (1, 1)
|
||||
assert result.usage.total_tokens == 557
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"])
|
||||
def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model):
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model)
|
||||
|
||||
assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model]
|
||||
assert result.cost == pytest.approx(1.5)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.usage.total_tokens == 557 + 336
|
||||
|
||||
|
||||
def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices():
|
||||
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
||||
with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash")
|
||||
without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None)
|
||||
|
||||
twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info)
|
||||
both = bu.calculate_vertex_ai_batch_cost_and_usage(
|
||||
[with_version, without_version], "vertex_ai/*", model_info=deployment_model_info
|
||||
)
|
||||
|
||||
assert twin.cost > 0
|
||||
assert both.cost == pytest.approx(2 * twin.cost)
|
||||
assert (both.successful_requests, both.failed_requests) == (2, 0)
|
||||
|
||||
|
||||
def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch):
|
||||
import litellm.cost_calculator as cc
|
||||
|
||||
def _calc(**kw):
|
||||
if kw["model"] == "gemini-unpriced":
|
||||
raise ValueError("no pricing")
|
||||
return (0.5, 0.25)
|
||||
|
||||
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
||||
rows = [
|
||||
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"),
|
||||
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"),
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
||||
|
||||
assert result.cost == pytest.approx(0.75)
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
||||
assert result.usage.total_tokens == 557 + 336
|
||||
assert result.models == ["gemini-unpriced", "gemini-2.5-flash"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
||||
)
|
||||
calls = _capture_cost_calls(monkeypatch)
|
||||
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert calls == []
|
||||
assert (result.successful_requests, result.failed_requests) == (0, 1)
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from litellm.rust_bridge.chat_completions.entrypoints import (
|
|||
)
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.chat_completions import dispatch
|
||||
from litellm.rust_bridge.catalog import Rules
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
PYTHON_RULES: Final = ()
|
||||
|
|
@ -256,3 +258,100 @@ async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.Mo
|
|||
NATIVE_ACOMPLETION.reset()
|
||||
assert result is expected
|
||||
assert [request.model for request in captured] == ["gpt-4o"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_completion_calls_keep_the_python_result() -> None:
|
||||
sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
|
||||
assert isinstance(sync_response, ModelResponse)
|
||||
assert isinstance(async_response, ModelResponse)
|
||||
assert sync_response.choices[0].message.content == "ok"
|
||||
assert async_response.choices[0].message.content == "ok"
|
||||
|
||||
|
||||
def test_sync_completion_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
assert request.model == "test-model"
|
||||
assert request.messages == MESSAGES
|
||||
assert request.custom_llm_provider == "openai"
|
||||
assert request.stream is True
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{"custom_llm_provider": "openai", "stream": True},
|
||||
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_completion_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = ModelResponse()
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_acompletion_marker_bypasses_native() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> ModelResponse:
|
||||
pytest.fail("acompletion's inner completion call must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{"custom_llm_provider": "openai", "acompletion": True},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Coroutine, Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
from pytest_socket import enable_socket, socket_allow_hosts
|
||||
|
||||
|
|
@ -10,6 +14,14 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
|||
import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import
|
||||
import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency
|
||||
import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency
|
||||
from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency
|
||||
from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency
|
||||
image_handling as image_handling_module,
|
||||
)
|
||||
from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency
|
||||
close_litellm_async_clients,
|
||||
)
|
||||
from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency
|
||||
|
||||
LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"]
|
||||
AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = (
|
||||
|
|
@ -20,6 +32,63 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = (
|
|||
"AZURE_USERNAME",
|
||||
"AZURE_PASSWORD",
|
||||
)
|
||||
AMBIENT_AWS_ENV_VARS: Final = (
|
||||
"AWS_PROFILE",
|
||||
"AWS_DEFAULT_PROFILE",
|
||||
"AWS_CONTAINER_CREDENTIALS_FULL_URI",
|
||||
"AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"AWS_ROLE_ARN",
|
||||
"AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"AWS_REGION_NAME",
|
||||
"AWS_DEFAULT_REGION",
|
||||
)
|
||||
MODULES_WITH_AWS_AUTH_HANDLERS: Final = (
|
||||
"litellm.main",
|
||||
"litellm.files.main",
|
||||
"litellm.rerank_api.main",
|
||||
"litellm.realtime_api.main",
|
||||
)
|
||||
CALLBACK_LISTS: Final = (
|
||||
"callbacks",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"input_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
"_async_input_callback",
|
||||
)
|
||||
RESET_TO_NONE_GLOBALS: Final = ("model_fallbacks", "cache")
|
||||
RESTORED_GLOBALS: Final = (
|
||||
"disable_aiohttp_transport",
|
||||
"force_ipv4",
|
||||
"drop_params",
|
||||
"secret_manager_client",
|
||||
"_key_management_system",
|
||||
"_key_management_settings",
|
||||
"api_base",
|
||||
"num_retries",
|
||||
"modify_params",
|
||||
"ssl_verify",
|
||||
"credential_list",
|
||||
"model_group_settings",
|
||||
"default_internal_user_params",
|
||||
"default_team_params",
|
||||
"prometheus_emit_stream_label",
|
||||
"vector_store_registry",
|
||||
"model_cost",
|
||||
"cost_margin_config",
|
||||
"cost_discount_config",
|
||||
"disable_hf_tokenizer_download",
|
||||
"disable_copilot_system_to_assistant",
|
||||
"cohere_models",
|
||||
"anthropic_models",
|
||||
"token_counter",
|
||||
"initialized_langfuse_clients",
|
||||
)
|
||||
MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient")
|
||||
SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client")
|
||||
|
||||
|
||||
def _allow_loopback_only() -> None:
|
||||
|
|
@ -29,11 +98,116 @@ def _allow_loopback_only() -> None:
|
|||
_allow_loopback_only()
|
||||
|
||||
|
||||
def pytest_collectstart() -> None:
|
||||
_allow_loopback_only()
|
||||
|
||||
|
||||
@pytest.hookimpl(trylast=True)
|
||||
def pytest_runtest_setup() -> None:
|
||||
_allow_loopback_only()
|
||||
|
||||
|
||||
def _run_coroutine_if_needed(result: object) -> None:
|
||||
if not asyncio.iscoroutine(result):
|
||||
return
|
||||
coroutine: Final[Coroutine[object, object, object]] = result
|
||||
try:
|
||||
asyncio.run(coroutine)
|
||||
except RuntimeError:
|
||||
try:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
coroutine.close()
|
||||
return
|
||||
loop.create_task(coroutine)
|
||||
|
||||
|
||||
def _close_handler_if_needed(handler: object) -> None:
|
||||
close: Final = getattr(handler, "close", None)
|
||||
if not callable(close):
|
||||
return
|
||||
_run_coroutine_if_needed(close())
|
||||
|
||||
|
||||
def _reset_aws_auth_caches() -> None:
|
||||
modules: Final = tuple(importlib.import_module(name) for name in MODULES_WITH_AWS_AUTH_HANDLERS)
|
||||
flushes: Final = (
|
||||
getattr(getattr(getattr(module, attr_name), "iam_cache", None), "flush_cache", None)
|
||||
for module in modules
|
||||
for attr_name in dir(module)
|
||||
)
|
||||
for flush in filter(callable, flushes):
|
||||
flush()
|
||||
boto3.DEFAULT_SESSION = None
|
||||
|
||||
|
||||
def _flush_client_caches() -> None:
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
image_handling_module.in_memory_cache.flush_cache()
|
||||
_reset_aws_auth_caches()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
|
||||
aws_dir: Final = tmp_path_factory.mktemp("aws-config")
|
||||
credentials: Final = aws_dir / "credentials"
|
||||
config: Final = aws_dir / "config"
|
||||
credentials.write_text("", encoding="utf-8")
|
||||
config.write_text("", encoding="utf-8")
|
||||
return credentials, config
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]:
|
||||
credentials, config = isolated_aws_config_files
|
||||
with pytest.MonkeyPatch.context() as environment:
|
||||
environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials))
|
||||
environment.setenv("AWS_CONFIG_FILE", str(config))
|
||||
environment.setenv("AWS_EC2_METADATA_DISABLED", "true")
|
||||
for name in AMBIENT_AWS_ENV_VARS:
|
||||
environment.delenv(name, raising=False)
|
||||
environment.delenv("PROXY_BASE_URL", raising=False)
|
||||
environment.setenv("LITELLM_CLI_DISABLE_KEYRING", "1")
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_litellm_globals() -> Iterator[None]:
|
||||
original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS}
|
||||
original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
|
||||
original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
|
||||
original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__}
|
||||
original_loggers: Final = {
|
||||
logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters))
|
||||
for logger in ALL_LOGGERS
|
||||
}
|
||||
original_tool_policy_registry: Final = tool_registry_writer_module._tool_policy_registry
|
||||
_flush_client_caches()
|
||||
for name in CALLBACK_LISTS:
|
||||
setattr(litellm, name, [])
|
||||
for name in RESET_TO_NONE_GLOBALS:
|
||||
setattr(litellm, name, None)
|
||||
for name in MODULE_LEVEL_CLIENTS:
|
||||
litellm.__dict__.pop(name, None)
|
||||
tool_registry_writer_module._tool_policy_registry = None
|
||||
yield
|
||||
_flush_client_caches()
|
||||
leaked_clients: Final = tuple(litellm.__dict__.pop(name, None) for name in MODULE_LEVEL_CLIENTS)
|
||||
for name, client in zip(MODULE_LEVEL_CLIENTS, leaked_clients):
|
||||
if client is not original_clients.get(name):
|
||||
_close_handler_if_needed(client)
|
||||
litellm.__dict__.update(original_clients)
|
||||
for name, value in (original_callbacks | original_reset | original_restored).items():
|
||||
setattr(litellm, name, value)
|
||||
for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items():
|
||||
logger.setLevel(level)
|
||||
logger.disabled = disabled
|
||||
logger.propagate = propagate
|
||||
logger.handlers = handlers
|
||||
logger.filters = filters
|
||||
tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_router_model_cost_state() -> Iterator[None]:
|
||||
original_live_routers: Final = frozenset(litellm_router_module._live_routers)
|
||||
|
|
@ -41,6 +215,7 @@ def isolate_router_model_cost_state() -> Iterator[None]:
|
|||
model_key: dict(model_value)
|
||||
for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
|
||||
}
|
||||
litellm_utils_module._invalidate_model_cost_lowercase_map()
|
||||
yield
|
||||
for router in tuple(litellm_router_module._live_routers):
|
||||
litellm_router_module._live_routers.discard(router)
|
||||
|
|
@ -68,4 +243,9 @@ def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
|
||||
|
||||
def pytest_sessionfinish() -> None:
|
||||
for name in MODULE_LEVEL_CLIENTS:
|
||||
_close_handler_if_needed(litellm.__dict__.pop(name, None))
|
||||
for name in SESSION_CLIENTS:
|
||||
_close_handler_if_needed(getattr(litellm, name, None))
|
||||
_run_coroutine_if_needed(close_litellm_async_clients())
|
||||
enable_socket()
|
||||
|
|
|
|||
|
|
@ -2,99 +2,13 @@
|
|||
"""
|
||||
Test to verify the Google GenAI generate_content handler functionality
|
||||
"""
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler
|
||||
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def test_non_stream_response_when_stream_requested_sync():
|
||||
"""
|
||||
Test that when a non-stream response is returned but streaming was requested,
|
||||
the sync handler correctly transforms it to generate_content format.
|
||||
"""
|
||||
from litellm.types.utils import Choices
|
||||
|
||||
# Mock a non-stream response (ModelResponse with valid choices)
|
||||
mock_response = ModelResponse(
|
||||
id="test-123",
|
||||
choices=[
|
||||
Choices(
|
||||
index=0,
|
||||
message={"role": "assistant", "content": "Hello, world!"},
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
created=1234567890,
|
||||
model="gpt-3.5-turbo",
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
# Create an instance of the adapter
|
||||
adapter = GoogleGenAIAdapter()
|
||||
|
||||
# Test the adapter's translate_completion_to_generate_content method directly
|
||||
result = adapter.translate_completion_to_generate_content(mock_response)
|
||||
|
||||
# Verify the result is a valid Google GenAI format response
|
||||
assert "candidates" in result
|
||||
assert isinstance(result["candidates"], list)
|
||||
assert len(result["candidates"]) > 0
|
||||
candidate = result["candidates"][0]
|
||||
assert "content" in candidate
|
||||
assert "parts" in candidate["content"]
|
||||
assert isinstance(candidate["content"]["parts"], list)
|
||||
assert len(candidate["content"]["parts"]) > 0
|
||||
assert "text" in candidate["content"]["parts"][0]
|
||||
assert candidate["content"]["parts"][0]["text"] == "Hello, world!"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_stream_response_when_stream_requested_async():
|
||||
"""
|
||||
Test that when a non-stream response is returned but streaming was requested,
|
||||
the async handler correctly transforms it to generate_content format.
|
||||
"""
|
||||
from litellm.types.utils import Choices
|
||||
|
||||
# Mock a non-stream response (ModelResponse with valid choices)
|
||||
mock_response = ModelResponse(
|
||||
id="test-123",
|
||||
choices=[
|
||||
Choices(
|
||||
index=0,
|
||||
message={"role": "assistant", "content": "Hello, world!"},
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
created=1234567890,
|
||||
model="gpt-3.5-turbo",
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
# Create an instance of the adapter
|
||||
adapter = GoogleGenAIAdapter()
|
||||
|
||||
# Test the adapter's translate_completion_to_generate_content method directly
|
||||
result = adapter.translate_completion_to_generate_content(mock_response)
|
||||
|
||||
# Verify the result is a valid Google GenAI format response
|
||||
assert "candidates" in result
|
||||
assert isinstance(result["candidates"], list)
|
||||
assert len(result["candidates"]) > 0
|
||||
candidate = result["candidates"][0]
|
||||
assert "content" in candidate
|
||||
assert "parts" in candidate["content"]
|
||||
assert isinstance(candidate["content"]["parts"], list)
|
||||
assert len(candidate["content"]["parts"]) > 0
|
||||
assert "text" in candidate["content"]["parts"][0]
|
||||
assert candidate["content"]["parts"][0]["text"] == "Hello, world!"
|
||||
|
||||
|
||||
def test_stream_response_when_stream_requested_sync():
|
||||
80
tests/unit/interactions/test_litellm_responses_bridge.py
Normal file
80
tests/unit/interactions/test_litellm_responses_bridge.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
"""
|
||||
Tests for LiteLLM Responses bridge provider.
|
||||
|
||||
Inherits from BaseInteractionsTest to run the same test suite against
|
||||
the litellm_responses bridge provider, which calls litellm.responses() internally.
|
||||
"""
|
||||
|
||||
|
||||
from litellm.interactions.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesInteractionsConfig,
|
||||
)
|
||||
from litellm.types.interactions import Turn
|
||||
|
||||
|
||||
class TestBridgeInputTransformation:
|
||||
"""Regression tests for translating Interactions input into Responses API input.
|
||||
|
||||
The bridge used to pass Google content parts through raw ({"type": "text"}),
|
||||
which the Responses API rejects with a 400, and it dropped the role encoded
|
||||
in step types and in the legacy "model" turn role.
|
||||
"""
|
||||
|
||||
def test_step_input_maps_roles_and_content_types(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[
|
||||
{"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]},
|
||||
{"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]},
|
||||
{"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]},
|
||||
]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]},
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]},
|
||||
]
|
||||
|
||||
def test_legacy_turn_input_maps_model_role_to_assistant(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[
|
||||
{"role": "user", "content": [{"type": "text", "text": "I like apples."}]},
|
||||
{"role": "model", "content": [{"type": "text", "text": "I like oranges."}]},
|
||||
]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]},
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]},
|
||||
]
|
||||
|
||||
def test_turn_pydantic_model_with_string_content(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[Turn(role="model", content="I like oranges.")]
|
||||
)
|
||||
assert transformed == [
|
||||
{"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}
|
||||
]
|
||||
|
||||
def test_string_input_passes_through(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello")
|
||||
assert transformed == "Hello"
|
||||
|
||||
def test_content_list_input_becomes_single_user_message(self):
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[{"type": "text", "text": "Hello"}, "world"]
|
||||
)
|
||||
assert transformed == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Hello"},
|
||||
{"type": "input_text", "text": "world"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
def test_non_text_content_passes_through_unchanged(self):
|
||||
image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"}
|
||||
transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(
|
||||
[{"type": "user_input", "content": [image_part]}]
|
||||
)
|
||||
assert transformed == [{"role": "user", "content": [image_part]}]
|
||||
|
|
@ -4,7 +4,7 @@ OpenAPI compliance tests for Google Interactions API.
|
|||
Validates that our SDK requests/responses match the OpenAPI spec at:
|
||||
https://ai.google.dev/static/api/interactions.openapi.json
|
||||
|
||||
Run with: pytest tests/test_litellm/interactions/test_openapi_compliance.py -v
|
||||
Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v
|
||||
"""
|
||||
|
||||
import json
|
||||
|
|
@ -22,6 +22,8 @@ from litellm.rust_bridge.messages.entrypoints import (
|
|||
NativeMessages,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
from pydantic import TypeAdapter
|
||||
from litellm.messages import dispatch
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
PYTHON_RULES: Final[Rules] = ()
|
||||
|
|
@ -284,3 +286,137 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon
|
|||
NATIVE_AMESSAGES.reset()
|
||||
assert result is expected
|
||||
assert [request.model for request in captured] == ["claude-sonnet-4-5"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_anthropic_messages_keeps_the_python_result() -> None:
|
||||
response: Final = await litellm.anthropic_messages(
|
||||
model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok"
|
||||
)
|
||||
|
||||
assert isinstance(response, dict)
|
||||
content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", []))
|
||||
assert content[0]["text"] == "ok"
|
||||
|
||||
|
||||
def test_sync_messages_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
assert request.model == "claude-test"
|
||||
assert request.messages == MESSAGES
|
||||
assert request.max_tokens == 10
|
||||
assert request.custom_llm_provider == "anthropic"
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{
|
||||
"model": "claude-test",
|
||||
"messages": MESSAGES,
|
||||
"max_tokens": 10,
|
||||
"custom_llm_provider": "anthropic",
|
||||
},
|
||||
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_messages_binding_error_delegates_unchanged_to_python() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
pytest.fail("a call without max_tokens cannot project a request and must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_messages_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{"model": "claude-test", "messages": MESSAGES, "max_tokens": 10},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_is_async_marker_bypasses_native() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
|
||||
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
def native(
|
||||
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
|
||||
) -> AnthropicMessagesResponse:
|
||||
pytest.fail("anthropic_messages' inner handler call must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{
|
||||
"model": "claude-test",
|
||||
"messages": MESSAGES,
|
||||
"max_tokens": 10,
|
||||
"custom_llm_provider": "anthropic",
|
||||
"is_async": True,
|
||||
},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue