diff --git a/litellm/constants.py b/litellm/constants.py index 9efeefacacd..276bc903680 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2001,7 +2001,7 @@ BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = get_env_float( "BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", 60.0 ) BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = get_env_float( - "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", 3600.0 + "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", 7800.0 ) _background_interaction_cost_polling_env: Final = os.getenv( "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 03feb8dc1e6..927f37791fb 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -126,6 +126,7 @@ from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage from litellm.types.extraction import ExtractionResponse, ExtractionUsage +from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel from litellm.types.llms.openai import ( HttpxBinaryResponseContent, @@ -1205,6 +1206,8 @@ def _get_response_model(completion_response: object) -> str | None: if completion_response is None: return None + if isinstance(completion_response, InteractionsAPIResponse): + return completion_response.model or completion_response.agent if isinstance(completion_response, BaseModel): return getattr(completion_response, "model", None) elif isinstance(completion_response, dict): diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index cd144646f79..76c65f7041d 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -300,6 +300,26 @@ def create( ) else: custom_llm_provider = custom_llm_provider or "gemini" + if agent and model and model != agent: + raise litellm.BadRequestError( + message=( + f"agent '{agent}' can only run on a deployment of that agent, but the model is '{model}'. " + "Configure a deployment whose model is the agent and send its model group." + ), + model=model, + llm_provider=custom_llm_provider, + ) + if extra_body and ( + ("agent" in extra_body and extra_body["agent"] != agent) or (agent and "model" in extra_body) + ): + raise litellm.BadRequestError( + message=( + "extra_body cannot set a different agent or add a model to an agent interaction. " + "Send the agent and model as top-level fields." + ), + model=model, + llm_provider=custom_llm_provider, + ) interactions_api_config: Final = get_provider_interactions_api_config( provider=custom_llm_provider, @@ -340,7 +360,7 @@ def create( litellm_logging_obj.update_from_kwargs( kwargs=kwargs, - model=model, + model=agent or model, optional_params=dict(optional_params), litellm_params={"litellm_call_id": litellm_call_id}, custom_llm_provider=custom_llm_provider, diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index a02afa3a13f..b9a2cc6493c 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -198,11 +198,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ request_body: Final[dict[str, object]] = {} - # Model or Agent (one required) - if model: - request_body["model"] = GeminiModelInfo.get_base_model(model) or model - elif agent: + if agent: request_body["agent"] = agent + elif model: + request_body["model"] = GeminiModelInfo.get_base_model(model) or model else: raise ValueError("Either 'model' or 'agent' must be provided") @@ -225,7 +224,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): dropped_keys: Final = ( frozenset({"agent_config"}) if drops_unsupported_params(litellm_params.get("drop_params")) - and (model or not _has_agent_config_type(optional_params.get("agent_config"))) + and (not agent or not _has_agent_config_type(optional_params.get("agent_config"))) else frozenset() ) for key in optional_keys: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e57abe99418..793c4ac8b4e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54503,6 +54503,20 @@ "supports_web_search": true, "web_search_billing_unit": "per_query" }, + "vertex_ai/deep-research-preview-04-2026": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "source": "https://cloud.google.com/model-garden/pricing", + "web_search_billing_unit": "per_query" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_batches": 1e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e57abe99418..793c4ac8b4e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54503,6 +54503,20 @@ "supports_web_search": true, "web_search_billing_unit": "per_query" }, + "vertex_ai/deep-research-preview-04-2026": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "source": "https://cloud.google.com/model-garden/pricing", + "web_search_billing_unit": "per_query" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_batches": 1e-07, diff --git a/tests/unit/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py index d55994dfb19..ac80f1dc484 100644 --- a/tests/unit/interactions/test_background_cost_polling.py +++ b/tests/unit/interactions/test_background_cost_polling.py @@ -3,7 +3,7 @@ import json import time from datetime import datetime, timedelta, timezone from itertools import islice -from typing import Optional +from typing import Final, Optional import httpx import pytest @@ -21,11 +21,13 @@ from litellm.interactions.background_cost_polling import ( poll_and_log_background_interaction_cost, PollSchedule, resume_unsettled_background_interactions, + SettlementOutcome, ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.interactions import InteractionsAPIResponse +from tests._support.clock_loops import AutojumpClockLoop USAGE_BLOCK = { "total_tokens": 175, @@ -1422,3 +1424,70 @@ async def test_create_on_a_replica_without_the_settlement_table_still_settles_in task.cancel() with pytest.raises(asyncio.CancelledError): await task + + +class _JobFinishingAt: + def __init__(self, finishes_at: float) -> None: + self.finishes_at: Final = finishes_at + self.last_polled_at: float | None = None + + async def fetch(self, context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + now: Final = asyncio.get_running_loop().time() + self.last_polled_at = now + finished: Final = now >= self.finishes_at + return _response("completed" if finished else "in_progress", with_usage=finished) + + +async def _finish_owned_loop_tasks() -> None: + current: Final = asyncio.current_task() + leftover: Final = tuple(task for task in asyncio.all_tasks() if task is not current) + for task in leftover: + task.cancel() + await asyncio.gather(*leftover, return_exceptions=True) + + +def _poll_on_a_virtual_clock( + context: BackgroundInteractionPollContext, job: _JobFinishingAt +) -> SettlementOutcome | None: + loop: Final = AutojumpClockLoop() + try: + return loop.run_until_complete(poll_and_log_background_interaction_cost(context, fetch_interaction=job.fetch)) + finally: + loop.run_until_complete(_finish_owned_loop_tasks()) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +_TWO_HOURS_SECONDS: Final = 7200.0 + + +def test_a_background_job_that_runs_past_the_old_one_hour_window_settles_under_the_default_poll_window() -> None: + logging_obj: Final = _logging_obj() + job: Final = _JobFinishingAt(_TWO_HOURS_SECONDS) + context: Final = BackgroundInteractionPollContext( + interaction_id="interactions/bg-abc", custom_llm_provider="gemini", logging_obj=logging_obj + ) + + outcome: Final = _poll_on_a_virtual_clock(context, job) + + assert outcome == "billed" + assert job.last_polled_at is not None and job.last_polled_at >= _TWO_HOURS_SECONDS + assert logging_obj.model_call_details["response_cost"] > 0 + + +def test_an_explicit_shorter_poll_window_still_gives_up_before_a_job_that_outlasts_it() -> None: + logging_obj: Final = _logging_obj() + job: Final = _JobFinishingAt(_TWO_HOURS_SECONDS) + context: Final = BackgroundInteractionPollContext( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + logging_obj=logging_obj, + timeout_seconds=3600.0, + ) + + outcome: Final = _poll_on_a_virtual_clock(context, job) + + assert outcome == "unsettled" + assert job.last_polled_at is not None and job.last_polled_at <= 3600.0 + assert logging_obj.model_call_details.get("response_cost") is None diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index cd32eb11485..9b448bbb1c6 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -1,3 +1,4 @@ +import asyncio import base64 import json from collections.abc import Mapping @@ -6,7 +7,8 @@ from typing import Final import httpx import pytest -from pydantic import JsonValue, TypeAdapter +import responses +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter from respx import MockRouter import litellm @@ -96,6 +98,45 @@ def api_key(): return "test-api-key" +_GEMINI_INTERACTIONS_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions" +_AGENT: Final = "antigravity-preview-09-2026" +_OTHER_AGENT: Final = "deep-research-preview-04-2026" + +_EXTRA_BODY_SELECTOR_OVERRIDES: Final = ( + pytest.param( + {"model": f"gemini/{_AGENT}", "agent": _AGENT}, {"agent": _OTHER_AGENT}, id="matching-agent-nested-other-agent" + ), + pytest.param({"agent": _AGENT}, {"agent": None}, id="agent-nested-null-agent"), + pytest.param({"model": "gemini/gemini-4-argon"}, {"agent": _AGENT}, id="model-nested-agent"), + pytest.param({"agent": _AGENT}, {"model": "gemini-4-argon"}, id="agent-nested-model"), + pytest.param({"agent": _AGENT}, {"model": _AGENT}, id="agent-nested-model-named-like-the-agent"), +) + +_EXTRA_BODY_PASS_THROUGH: Final = ( + pytest.param( + {"agent": _AGENT}, + {"agent": _AGENT, "store": False}, + {"agent": _AGENT, "store": False}, + id="same-agent-repeated", + ), + pytest.param( + {"agent": _AGENT}, {"custom_flag": 1}, {"agent": _AGENT, "custom_flag": 1}, id="agent-unrelated-extra" + ), + pytest.param( + {"model": "gemini/gemini-4-argon"}, + {"custom_flag": 1}, + {"model": "gemini-4-argon", "custom_flag": 1}, + id="model-unrelated-extra", + ), + pytest.param( + {"model": "gemini/gemini-4-argon"}, + {"model": "gemini-4-carbon"}, + {"model": "gemini-4-carbon"}, + id="model-only-nested-model", + ), +) + + class TestGoogleInteractionsCreate: @pytest.mark.usefixtures("fake_provider_credentials") def test_missing_model_and_agent(self, api_key): @@ -305,6 +346,15 @@ class TestGoogleInteractionsCreate: }, id="agent-keeps-an-object", ), + pytest.param( + {"model": "gemini/antigravity-preview-09-2026", "agent": "antigravity-preview-09-2026"}, + {"type": "antigravity", "model": "gemini-3.8-flash"}, + { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + }, + id="matching-model-and-agent-sends-the-agent-with-its-config", + ), pytest.param( {"agent": "antigravity-preview-09-2026"}, "antigravity", @@ -347,6 +397,45 @@ class TestGoogleInteractionsCreate: assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."} + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_refuses_an_agent_with_another_model_before_any_http_call(self, respx_mock: MockRouter) -> None: + with pytest.raises(litellm.BadRequestError, match="antigravity-preview-09-2026"): + interactions.create( + model="gemini/gemini-2.5-flash", + agent="antigravity-preview-09-2026", + input="Say hi in one word.", + api_key="gemini-offline", + ) + + assert not respx_mock.calls + + @pytest.mark.parametrize(("target", "extra_body"), _EXTRA_BODY_SELECTOR_OVERRIDES) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_refuses_extra_body_that_changes_the_agent_or_model_before_any_http_call( + self, respx_mock: MockRouter, target: Mapping[str, str], extra_body: Mapping[str, object] + ) -> None: + with pytest.raises(litellm.BadRequestError, match="extra_body"): + interactions.create(**target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline") + + assert not respx_mock.calls + + @pytest.mark.parametrize(("target", "extra_body", "expected_target"), _EXTRA_BODY_PASS_THROUGH) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_sends_extra_body_that_keeps_the_agent_and_model( + self, + respx_mock: MockRouter, + target: Mapping[str, str], + extra_body: Mapping[str, object], + expected_target: Mapping[str, object], + ) -> None: + route: Final = respx_mock.post(_GEMINI_INTERACTIONS_URL).mock( + return_value=httpx.Response(200, json={"id": "interaction-extra", "status": "completed"}) + ) + + interactions.create(**target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline") + + assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."} + @pytest.mark.usefixtures("fake_provider_credentials") def test_create_rejects_agent_config_a_bridged_provider_cannot_take(self, respx_mock: MockRouter) -> None: route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( @@ -494,6 +583,39 @@ class TestInteractionsAcreateOffline: def _httpx_only_transport(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + @pytest.mark.parametrize(("target", "extra_body"), _EXTRA_BODY_SELECTOR_OVERRIDES) + @pytest.mark.usefixtures("fake_provider_credentials") + @pytest.mark.asyncio + async def test_acreate_refuses_extra_body_that_changes_the_agent_or_model_before_any_http_call( + self, respx_mock: MockRouter, target: Mapping[str, str], extra_body: Mapping[str, object] + ) -> None: + with pytest.raises(litellm.BadRequestError, match="extra_body"): + await interactions.acreate( + **target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline" + ) + + assert not respx_mock.calls + + @pytest.mark.parametrize(("target", "extra_body", "expected_target"), _EXTRA_BODY_PASS_THROUGH) + @pytest.mark.usefixtures("fake_provider_credentials") + @pytest.mark.asyncio + async def test_acreate_sends_extra_body_that_keeps_the_agent_and_model( + self, + respx_mock: MockRouter, + target: Mapping[str, str], + extra_body: Mapping[str, object], + expected_target: Mapping[str, object], + ) -> None: + route: Final = respx_mock.post(_GEMINI_INTERACTIONS_URL).mock( + return_value=httpx.Response(200, json={"id": "interaction-extra", "status": "completed"}) + ) + + await interactions.acreate( + **target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline" + ) + + assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."} + @pytest.mark.usefixtures("fake_provider_credentials") @pytest.mark.asyncio async def test_acreate_simple_gemini(self, respx_mock: MockRouter) -> None: @@ -922,3 +1044,134 @@ def test_create_sends_a_caller_continuation_token_to_resume_a_long_decode(respx_ }, sent assert response.status == "completed" assert response.continuation_token is None + + +_DEEP_RESEARCH_AGENT: Final = "deep-research-preview-04-2026" +_DEEP_RESEARCH_USAGE: Final = { + "total_input_tokens": 1000, + "total_cached_tokens": 200, + "total_output_tokens": 300, + "total_thought_tokens": 100, + "total_tokens": 1400, + "input_tokens_by_modality": [{"modality": "text", "tokens": 1000}], + "output_tokens_by_modality": [{"modality": "text", "tokens": 300}], + "grounding_tool_count": [{"type": "google_search", "count": 3}], +} + + +def _deep_research_cost_from_the_vertex_row() -> float: + row: Final = litellm.model_cost[f"vertex_ai/{_DEEP_RESEARCH_AGENT}"] + return ( + (1000 - 200) * row["input_cost_per_token"] + + 200 * row["cache_read_input_token_cost"] + + (300 + 100) * row["output_cost_per_token"] + + 3 * row["search_context_cost_per_query"]["search_context_size_medium"] + ) + + +class _LoggedModelMapValue(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + key: str + + +class _LoggedModelMap(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + model_map_value: _LoggedModelMapValue + + +class _LoggedModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + model_map_information: _LoggedModelMap + + +class _LoggedCost(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + response_cost: float + standard_logging_object: _LoggedModel + + +class _CostRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.cost: _LoggedCost | None = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.cost = _LoggedCost.model_validate(kwargs) + self.logged.set() + + +_DEEP_RESEARCH_INTERACTION: Final = { + "id": "interaction-dr", + "object": "interaction", + "agent": _DEEP_RESEARCH_AGENT, + "status": "completed", + "usage": _DEEP_RESEARCH_USAGE, +} +_DEEP_RESEARCH_URL: Final = ( + "https://aiplatform.googleapis.com/v1beta1/projects/dr-project/locations/global/interactions" +) + + +def _deep_research_reply(stream: bool) -> httpx.Response: + if not stream: + return httpx.Response(200, json=_DEEP_RESEARCH_INTERACTION) + completed: Final = json.dumps({"event_type": "interaction.completed", "interaction": _DEEP_RESEARCH_INTERACTION}) + return httpx.Response(200, text=f"data: {completed}\n\n", headers={"content-type": "text/event-stream"}) + + +@pytest.mark.parametrize("stream", [False, True], ids=["background-or-foreground", "streamed"]) +@pytest.mark.asyncio +async def test_an_agent_only_vertex_deep_research_interaction_is_billed_under_the_agent_row( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + recorder: Final = _CostRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(f"{_DEEP_RESEARCH_URL}?alt=sse" if stream else _DEEP_RESEARCH_URL).mock( + return_value=_deep_research_reply(stream) + ) + credentials: Final = { + "type": "authorized_user", + "client_id": "dr-client", + "client_secret": "dr-secret", + "refresh_token": "dr-refresh-token", + } + + with responses.RequestsMock(assert_all_requests_are_fired=False) as google_oauth: + google_oauth.post( + "https://oauth2.googleapis.com/token", + json={"access_token": "dr-token", "expires_in": 3600, "token_type": "Bearer"}, + ) + response: Final = await interactions.acreate( + agent=_DEEP_RESEARCH_AGENT, + input="Compare grid-scale storage options.", + custom_llm_provider="vertex_ai", + vertex_project="dr-project", + vertex_location="global", + vertex_credentials=json.dumps(credentials), + stream=stream, + ) + if stream: + async for _ in response: + pass + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert json.loads(route.calls.last.request.content) == { + "agent": _DEEP_RESEARCH_AGENT, + "input": "Compare grid-scale storage options.", + "stream": stream, + } + assert recorder.cost is not None + assert recorder.cost.standard_logging_object.model == _DEEP_RESEARCH_AGENT + assert recorder.cost.standard_logging_object.model_map_information.model_map_value.key == ( + f"vertex_ai/{_DEEP_RESEARCH_AGENT}" + ) + assert recorder.cost.response_cost == pytest.approx(_deep_research_cost_from_the_vertex_row()) diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index e6df7bf1d4f..ef793391788 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -425,6 +425,208 @@ def test_interaction_by_id_naming_a_model_the_key_cannot_call_is_refused_before_ assert not respx_mock.calls +_DEEP_RESEARCH_USAGE: Final = { + "total_input_tokens": 1000, + "total_cached_tokens": 200, + "total_output_tokens": 300, + "total_thought_tokens": 100, + "total_tokens": 1400, + "input_tokens_by_modality": [{"modality": "text", "tokens": 1000}], + "output_tokens_by_modality": [{"modality": "text", "tokens": 300}], + "grounding_tool_count": [{"type": "google_search", "count": 3}], +} + + +def _deep_research_cost_from_the_vertex_row() -> float: + row: Final = litellm.model_cost[f"vertex_ai/{_DEEP_RESEARCH_AGENT}"] + return ( + (1000 - 200) * row["input_cost_per_token"] + + 200 * row["cache_read_input_token_cost"] + + (300 + 100) * row["output_cost_per_token"] + + 3 * row["search_context_cost_per_query"]["search_context_size_medium"] + ) + + +def _deep_research_create(model_group: str) -> Mapping[str, str]: + return {"model": model_group, "agent": _DEEP_RESEARCH_AGENT, "input": "Compare grid-scale storage options."} + + +def _completed_deep_research(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(_VERTEX_COLLECTION_URL).mock( + return_value=httpx.Response( + 200, + json={ + "id": _INTERACTION_ID, + "object": "interaction", + "agent": _DEEP_RESEARCH_AGENT, + "status": "completed", + "usage": _DEEP_RESEARCH_USAGE, + }, + ) + ) + + +def test_routed_create_with_an_agent_sends_the_agent_with_the_deployment_credentials_and_bills_the_agent( + interactions_proxy: TestClient, respx_mock: respx.MockRouter +) -> None: + upstream: Final = _completed_deep_research(respx_mock) + + response: Final = interactions_proxy.post( + "/v1beta/interactions", json=_deep_research_create(_VERTEX_DEEP_RESEARCH_GROUP) + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == { + "agent": _DEEP_RESEARCH_AGENT, + "input": "Compare grid-scale storage options.", + } + assert upstream.calls.last.request.headers["authorization"] == "Bearer deployment-token" + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx( + _deep_research_cost_from_the_vertex_row() + ) + + +def test_routed_create_with_an_agent_on_another_models_deployment_is_refused_before_any_upstream_call( + interactions_proxy: TestClient, respx_mock: respx.MockRouter, google_oauth: responses.RequestsMock +) -> None: + response: Final = interactions_proxy.post("/v1beta/interactions", json=_deep_research_create(_VERTEX_GROUP)) + + assert response.status_code == 400, response.text + assert _DEEP_RESEARCH_AGENT in response.text + assert not respx_mock.calls + assert not google_oauth.calls + + +@pytest.mark.parametrize( + "create", + ( + pytest.param( + {**_deep_research_create(_VERTEX_DEEP_RESEARCH_GROUP), "extra_body": {"agent": "other-agent"}}, + id="nested-other-agent", + ), + pytest.param( + {**_deep_research_create(_VERTEX_DEEP_RESEARCH_GROUP), "extra_body": {"model": "gemini-4-argon"}}, + id="nested-model-beside-the-agent", + ), + pytest.param( + { + "model": _VERTEX_GROUP, + "input": "Compare grid-scale storage options.", + "extra_body": {"agent": _DEEP_RESEARCH_AGENT}, + }, + id="nested-agent-on-a-model-deployment", + ), + ), +) +def test_routed_create_with_extra_body_that_changes_the_agent_or_model_is_refused_before_any_upstream_call( + interactions_proxy: TestClient, + respx_mock: respx.MockRouter, + google_oauth: responses.RequestsMock, + create: Mapping[str, object], +) -> None: + response: Final = interactions_proxy.post("/v1beta/interactions", json=create) + + assert response.status_code == 400, response.text + assert "extra_body" in response.text + assert not respx_mock.calls + assert not google_oauth.calls + + +def test_routed_create_with_extra_body_repeating_the_agent_sends_it_once_with_the_other_extras( + interactions_proxy: TestClient, respx_mock: respx.MockRouter +) -> None: + upstream: Final = _completed_deep_research(respx_mock) + + response: Final = interactions_proxy.post( + "/v1beta/interactions", + json={ + **_deep_research_create(_VERTEX_DEEP_RESEARCH_GROUP), + "extra_body": {"agent": _DEEP_RESEARCH_AGENT, "custom_flag": 1}, + }, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == { + "agent": _DEEP_RESEARCH_AGENT, + "input": "Compare grid-scale storage options.", + "custom_flag": 1, + } + + +def test_a_key_allowed_the_agent_deployment_runs_the_agent( + interactions_proxy: TestClient, respx_mock: respx.MockRouter +) -> None: + upstream: Final = _completed_deep_research(respx_mock) + + with _cached_key(UserAPIKeyAuth(models=[_VERTEX_DEEP_RESEARCH_GROUP])) as authorization: + response: Final = interactions_proxy.post( + "/v1beta/interactions", + json=_deep_research_create(_VERTEX_DEEP_RESEARCH_GROUP), + headers={"authorization": authorization}, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content)["agent"] == _DEEP_RESEARCH_AGENT + + +@pytest.mark.parametrize( + ("key", "teams", "model_group", "status_code"), + ( + pytest.param( + UserAPIKeyAuth(models=[_VERTEX_GROUP]), (), _VERTEX_GROUP, 400, id="own-group-cannot-run-another-agent" + ), + pytest.param( + UserAPIKeyAuth(models=[_VERTEX_GROUP]), (), _VERTEX_DEEP_RESEARCH_GROUP, 403, id="agent-group-not-on-key" + ), + pytest.param( + UserAPIKeyAuth(models=[_VERTEX_DEEP_RESEARCH_GROUP]), + (), + _VERTEX_GROUP, + 403, + id="other-group-not-on-key", + ), + pytest.param( + UserAPIKeyAuth(team_id="team-without-deep-research"), + (LiteLLM_TeamTableCachedObj(team_id="team-without-deep-research", models=[_VERTEX_GROUP]),), + _VERTEX_DEEP_RESEARCH_GROUP, + 403, + id="agent-group-not-on-team", + ), + ), +) +def test_a_restricted_key_cannot_run_an_agent_it_was_not_given_the_deployment_for( + interactions_proxy: TestClient, + respx_mock: respx.MockRouter, + google_oauth: responses.RequestsMock, + key: UserAPIKeyAuth, + teams: tuple[LiteLLM_TeamTableCachedObj, ...], + model_group: str, + status_code: int, +) -> None: + with _cached_key(key, *teams) as authorization: + response: Final = interactions_proxy.post( + "/v1beta/interactions", json=_deep_research_create(model_group), headers={"authorization": authorization} + ) + + assert response.status_code == status_code, response.text + assert not respx_mock.calls + assert not google_oauth.calls + + +def test_routed_create_without_an_agent_still_sends_the_deployment_model( + interactions_proxy: TestClient, respx_mock: respx.MockRouter +) -> None: + upstream: Final = respx_mock.post(_VERTEX_COLLECTION_URL).mock( + return_value=httpx.Response(200, json={"id": _INTERACTION_ID, "object": "interaction", "status": "completed"}) + ) + + response: Final = interactions_proxy.post("/v1beta/interactions", json={"model": _VERTEX_GROUP, "input": "Hi"}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == {"model": "gemini-4-argon", "input": "Hi"} + assert upstream.calls.last.request.headers["authorization"] == "Bearer deployment-token" + + def _project_url(project: str) -> str: return f"https://aiplatform.googleapis.com/v1beta1/projects/{project}/locations/global/interactions" diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 23748790869..10ab0cd25f5 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -5250,6 +5250,37 @@ def test_completion_cost_bills_interactions_google_search_per_query(): assert cost > 3 * per_query_cost +def test_completion_cost_bills_an_agent_interaction_at_its_agent_row(): + from litellm.types.interactions import InteractionsAPIResponse + + agent: Final = "deep-research-preview-04-2026" + model_info: Final = litellm.get_model_info(model=agent, custom_llm_provider="vertex_ai") + response: Final = InteractionsAPIResponse( + id="interactions/agent123", + agent=agent, + status="completed", + usage={ + "total_input_tokens": 1000, + "input_tokens_by_modality": [{"modality": "text", "tokens": 1000}], + "total_cached_tokens": 200, + "total_output_tokens": 300, + "output_tokens_by_modality": [{"modality": "text", "tokens": 300}], + "total_thought_tokens": 100, + "grounding_tool_count": [{"type": "google_search", "count": 3}], + }, + ) + + cost: Final = completion_cost(completion_response=response, custom_llm_provider="vertex_ai") + + assert model_info["key"] == f"vertex_ai/{agent}" + assert cost == pytest.approx( + 800 * model_info["input_cost_per_token"] + + 200 * model_info["cache_read_input_token_cost"] + + 400 * model_info["output_cost_per_token"] + + 3 * model_info["search_context_cost_per_query"]["search_context_size_medium"] + ) + + def test_completion_cost_bills_interactions_video_output_at_video_rate(): from litellm.types.interactions import InteractionsAPIResponse