mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
0751886680
* feat(bedrock): support retrieve for model-invocation-job batch ARNs `bedrock.retrieve_batch` previously only handled `:async-invoke/` ARNs (Twelve Labs Marengo embeddings). The `:model-invocation-job/` ARNs returned by `CreateModelInvocationJob` (the bulk batch inference API behind `bedrock.create_batch`) fell through and returned a misleading data-plane error, leaving created jobs unretrievable through the LiteLLM batches API. The two ARN families live on different AWS service endpoints (`bedrock-runtime` data plane vs `bedrock` control plane), so they need distinct handlers. This adds: * `BedrockBatchesHandler._handle_model_invocation_job_status` — calls the control plane via boto3 (`bedrock:GetModelInvocationJob`), reusing `BaseAWSLLM.get_credentials` for credential resolution so model_list / env / role-assumption configs continue to apply. The response is reshaped into a `LiteLLMBatch` with the same status mapping `transform_create_batch_response` already uses. * Output-file-URI prediction. Bedrock surfaces the user-supplied `s3OutputDataConfig.s3Uri` *prefix* in `GetModelInvocationJob`, but results actually land at `<prefix>/<job-id>/<basename(input)>.out`. We compute that single-file URI client-side and surface it as `output_file_id`, so OpenAI-style `client.files.content(...)` works without an extra `ListObjectsV2` round-trip. The bare prefix stays in metadata for callers that want the manifest. * Dispatch in `litellm/batches/main.py` for the new ARN family, alongside the existing async-invoke branch. * Unit tests covering ARN parsing, output-URI prediction (incl. edge cases), the full status mapping, region resolution precedence, and failure-message propagation. Note: `request_counts` is intentionally `(0, 0, 0)` — `GetModelInvocationJob` does not report per-record counts; getting accurate numbers requires parsing `manifest.json.out` from the output S3 prefix, which is left to callers. Made-with: Cursor * fix(bedrock): address PR feedback on model-invocation-job retrieve Addresses Greptile P2 findings on #26834: 1. Use the bare job id (not the full ARN) when constructing the `api_base` URL for `pre_call` logging. Passing the full ARN double- counts the `model-invocation-job/` segment and embeds colons in the path, producing misleading log lines. 2. Drop the `or output_prefix` fallback when `_predict_output_file_uri` returns None. A bare prefix is not a downloadable object and surfacing it as `output_file_id` re-creates the very NoSuchKey bug this handler exists to fix. The bare prefix is still preserved in `metadata["output_s3_uri"]` for callers that want to do their own S3 listing or read `manifest.json.out`. `metadata["output_file_uri"]` uses "" rather than None to satisfy the OpenAI Batch metadata schema (`dict[str, str]`); callers should branch on the typed `output_file_id` field instead. Also expands test coverage on the new code path: - new "stay None" regression test for the prediction-fail case - pre_call/post_call logging hook assertions (incl. the bare-id URL) - explicit cancelled_at / expired_at coverage - _to_epoch type-handling matrix and the boto3 ImportError branch - defensive _extract_region_from_bedrock_arn exception path - empty-basename case for _predict_output_file_uri Patch coverage on the changed lines is now 100% (the only remaining uncovered lines in the file belong to the pre-existing `_handle_async_invoke_status` method, which this PR does not touch). Made-with: Cursor * test(bedrock): cover retrieve_batch dispatch for both ARN families Codecov flagged 8 uncovered lines on `litellm/batches/main.py` after this PR refactored the Bedrock dispatch into a single guard with two sub-branches (`async-invoke` + `model-invocation-job`). Existing tests exercised the handlers directly but not the dispatch in `main.py`. Adds `tests/test_litellm/batches/test_retrieve_batch_bedrock_dispatch.py` with 6 mocked tests that exercise `litellm.retrieve_batch` end-to-end for the dispatch logic: - async-invoke ARN routes to `_handle_async_invoke_status` - async-invoke ARN with no region falls back to "us-east-1" (preserves prior behavior on this branch) - model-invocation-job ARN routes to the new `_handle_model_invocation_job_status` handler - model-invocation-job ARN with no region forwards None (so the new handler can sniff region from the ARN itself, rather than getting silently routed to us-east-1) - unrelated bedrock ARN family falls through to the generic provider-config retrieve path (neither special handler invoked) - non-bedrock batch ids skip the bedrock dispatch entirely Both handlers are mocked at the import site so the tests don't hit AWS — the focus here is purely the new dispatch logic in main.py. Co-authored-by: Cursor <cursoragent@cursor.com> * test(bedrock): move retrieve_batch dispatch test to tests/test_litellm/ The dispatch test landed under `tests/test_litellm/batches/`, a new directory that no upstream `test-unit-*.yml` workflow's `test-path` allow-list includes. As a result, the test was never executed in CI and codecov reported `litellm/batches/main.py` patch coverage at 11.11% (8 lines uncovered) — the lines belonging to this PR's dispatch refactor itself. Move the file up one level so it matches the `tests/test_litellm/test_*.py` glob that `test-unit-misc.yml` already runs, and adjust `sys.path.insert` for the new depth. The companion handler tests under `tests/test_litellm/llms/bedrock/batches/test_handler.py` are unaffected — they're picked up by the `llms` directory in `test-unit-llm-providers.yml`. Made-with: Cursor --------- Co-authored-by: Cursor <cursoragent@cursor.com> |
||
|---|---|---|
| .. | ||
| agent_tests | ||
| audio_tests | ||
| basic_proxy_startup_tests | ||
| batches_tests | ||
| benchmarks | ||
| code_coverage_tests | ||
| documentation_tests | ||
| enterprise | ||
| guardrails_tests | ||
| image_gen_tests | ||
| litellm | ||
| litellm-proxy-extras | ||
| litellm_core_utils | ||
| litellm_utils_tests | ||
| llm_responses_api_testing | ||
| llm_translation | ||
| load_tests | ||
| local_testing | ||
| logging_callback_tests | ||
| mcp_tests | ||
| multi_instance_e2e_tests | ||
| ocr_tests | ||
| old_proxy_tests/tests | ||
| openai_endpoints_tests | ||
| otel_tests | ||
| pass_through_tests | ||
| pass_through_unit_tests | ||
| proxy_admin_ui_tests | ||
| proxy_e2e_anthropic_messages_tests | ||
| proxy_security_tests | ||
| proxy_unit_tests | ||
| router_unit_tests | ||
| scim_tests | ||
| search_tests | ||
| spend_tracking_tests | ||
| store_model_in_db_tests | ||
| test_litellm | ||
| unified_google_tests | ||
| vector_store_tests | ||
| windows_tests | ||
| __init__.py | ||
| _flush_vcr_cache.py | ||
| _vcr_conftest_common.py | ||
| _vcr_redis_persister.py | ||
| eval_swe_bench.py | ||
| gettysburg.wav | ||
| large_text.py | ||
| openai_batch_completions.jsonl | ||
| README.MD | ||
| test_budget_management.py | ||
| test_callbacks_on_proxy.py | ||
| test_config.py | ||
| test_debug_warning.py | ||
| test_default_encoding_non_root.py | ||
| test_end_users.py | ||
| test_entrypoint.py | ||
| test_fallbacks.py | ||
| test_gpt5_azure_temperature_support.py | ||
| test_health.py | ||
| test_keys.py | ||
| test_litellm_proxy_responses_config.py | ||
| test_logging.conf | ||
| test_models.py | ||
| test_new_vector_store_endpoints.py | ||
| test_openai_endpoints.py | ||
| test_organizations.py | ||
| test_otel_thread_leak.py | ||
| test_passthrough_endpoints.py | ||
| test_presidio_latency.py | ||
| test_proxy_server_non_root.py | ||
| test_ratelimit.py | ||
| test_resource_cleanup.py | ||
| test_service_logger_otel.py | ||
| test_spend_logs.py | ||
| test_team.py | ||
| test_team_logging.py | ||
| test_team_members.py | ||
| test_users.py | ||
In total litellm runs 1000+ tests
[02/20/2025] Update:
To make it easier to contribute and map what behavior is tested,
we've started mapping the litellm directory in tests/test_litellm
This folder can only run mock tests.