diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index a5a4b7f6f3b..73d8d2649bf 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -62,6 +62,17 @@ OFF_PEAK_WINDOW: Final[JsonSchema] = { } OBJECT_KEYS: dict[str, JsonSchema] = { + "search_free_requests": { + "type": "object", + "description": "Free Google Search grounding requests per period, shared across models in the same shared_group.", + "required": ["count", "period", "shared_group"], + "additionalProperties": False, + "properties": { + "count": {"type": "integer", "minimum": 0}, + "period": {"type": "string", "enum": ["day", "month"]}, + "shared_group": {"type": "string", "pattern": "\\S"}, + }, + }, "off_peak_pricing": { "type": "object", "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", @@ -111,8 +122,9 @@ OBJECT_KEYS: dict[str, JsonSchema] = { "properties": { "model": NON_BLANK_STRING, "region": NON_BLANK_STRING, + "litellm_provider": {"type": "string", "enum": ["bedrock", "bedrock_mantle"]}, }, - "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", + "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in.", }, } diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 4526c428f1a..664dbf87fe9 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py @@ -636,7 +636,11 @@ class TestAgentCoreSearchEdgeCases: from litellm.search.cost_calculator import search_provider_cost_per_query monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) - assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0) + price_per_query = litellm.model_cost["agentcore/search"]["input_cost_per_query"] + assert price_per_query > 0 + assert search_provider_cost_per_query( + model="agentcore/search", custom_llm_provider="agentcore", number_of_queries=3 + ) == (3 * price_per_query, 0.0) def _transform_text(text: str) -> SearchResponse: diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 88d4cc3b872..3a1c316caf0 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1027,7 +1027,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): { "model_name": "semantic-ranker-default-004", "litellm_params": { - "model": "vertex_ai/semantic-ranker-default-004", + "model": "vertex_ai/semantic-ranker-custom-test", "vertex_project": "test-project", "vertex_location": "us-east5", }, @@ -1036,7 +1036,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): ] ) router_model_id: Final = router.model_list[0]["model_info"]["id"] - assert litellm.model_cost["vertex_ai/semantic-ranker-default-004"].get("input_cost_per_query") is None + assert litellm.model_cost.get("vertex_ai/semantic-ranker-custom-test", {}).get("input_cost_per_query") is None response: Final = RerankResponse( id="vertex_ai_rerank_test", @@ -1046,7 +1046,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): cost: Final = completion_cost( completion_response=response, - model="vertex_ai/semantic-ranker-default-004", + model="vertex_ai/semantic-ranker-custom-test", custom_llm_provider="vertex_ai", call_type="arerank", custom_pricing=True, diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index 766dc9d8648..ae0a14d6d44 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -559,7 +559,17 @@ def test_deepseek_rows_bill_half_rate_outside_weekday_peak_hours(path: Path): assert all(name in rows for name in DEEPSEEK_PRICED_ROWS) -PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset({"sagemaker", "bedrock_converse"}) +PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset( + { + "sagemaker", + "bedrock_converse", + "coralbricks", + "inferenceapis", + "microsoft_365_copilot", + "openinfer", + "scaledown", + } +) MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY: Final = frozenset({"search", "decisions", "evaluation"}) VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST: Final = frozenset( { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 868cc244792..464d99f6583 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -833,6 +833,18 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "type": "object", "properties": { "supports_computer_use": {"type": "boolean"}, + "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "cache_creation_input_token_cost_decisions": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "cache_read_input_token_cost_decisions": {"type": "number"}, + "explicit_cache_read_input_token_cost": {"type": "number"}, + "input_cost_per_text_message": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "input_cost_per_token_decisions": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "output_cost_per_token_decisions": {"type": "number"}, + "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, + "supports_bedrock_native_structured_output": {"type": "boolean"}, "cache_creation_input_audio_token_cost": {"type": "number"}, "cache_creation_input_token_cost": {"type": "number"}, "cache_creation_input_token_cost_above_1hr": {"type": "number"}, @@ -959,6 +971,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "image_edit", "embedding", "evaluation", + "extraction", "guardrail", "image_generation", "video_generation", @@ -1132,6 +1145,16 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "type": "object", "additionalProperties": {"type": "number"}, }, + "search_free_requests": { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "period": {"type": "string", "enum": ["day", "month"]}, + "shared_group": {"type": "string"}, + }, + "required": ["count", "period", "shared_group"], + "additionalProperties": False, + }, "search_context_cost_per_query": { "type": "object", "properties": { @@ -8187,12 +8210,17 @@ def test_models_by_provider(): continue elif k == "sample_spec": continue - elif ( - v["litellm_provider"] == "sagemaker" - or v["litellm_provider"] == "bedrock_converse" + elif v["litellm_provider"] in ( + "sagemaker", + "bedrock_converse", + "coralbricks", + "inferenceapis", + "microsoft_365_copilot", + "openinfer", + "scaledown", ): continue - elif v.get("mode") in ("search", "decisions", "evaluation"): + elif v.get("mode") in ("search", "decisions", "evaluation", "extraction"): continue else: providers.add(v["litellm_provider"])