fix(cost): resolve vertex location from global config and env for cost tracking

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-13 22:34:56 +00:00
parent b95a4d303d
commit 41afb13fc6
2 changed files with 41 additions and 5 deletions

View file

@ -1397,6 +1397,13 @@ class Logging(LiteLLMLoggingBaseClass):
if margin_total_amount is not None:
self.cost_breakdown["margin_total_amount"] = margin_total_amount
def _resolve_vertex_location(self) -> str | None:
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
if not (hasattr(self, "litellm_params") and self.litellm_params):
return None
return VertexBase.safe_get_vertex_ai_location(litellm_params=self.litellm_params)
def _response_cost_calculator(
self,
result: Union[
@ -1484,11 +1491,7 @@ class Logging(LiteLLMLoggingBaseClass):
if hasattr(self, "litellm_params") and self.litellm_params
else None
),
"vertex_location": (
self.litellm_params.get("vertex_location") or self.litellm_params.get("vertex_ai_location")
if hasattr(self, "litellm_params") and self.litellm_params
else None
),
"vertex_location": self._resolve_vertex_location(),
}
except Exception as e: # error creating kwargs for cost calculation
debug_info = StandardLoggingModelCostFailureDebugInformation(

View file

@ -4575,3 +4575,36 @@ def test_response_cost_calculator_passes_vertex_location(monkeypatch, vertex_loc
cost = logging_obj._response_cost_calculator(result=response)
assert cost == pytest.approx((1000 * 3e-6 + 200 * 1.5e-5) * uplift, rel=1e-9)
def test_response_cost_calculator_uses_env_vertex_location(monkeypatch):
"""A deployment can leave vertex_location unset and let the region come from the environment, which
still routes to a regional endpoint, so cost has to pick the same fallback the request dispatch does."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
logging_obj = LitellmLogging(
model="vertex_ai/claude-sonnet-4-6",
messages=[{"role": "user", "content": "Hey"}],
stream=False,
call_type="completion",
start_time=time.time(),
litellm_call_id="vertex-location-env-123",
function_id="test-fn",
)
logging_obj.update_environment_variables(
model="vertex_ai/claude-sonnet-4-6",
user="",
optional_params={},
litellm_params={"vertex_project": "test-project"},
)
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
response = ModelResponse(
model="claude-sonnet-4-6",
usage=litellm.Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200),
)
cost = logging_obj._response_cost_calculator(result=response)
assert cost == pytest.approx((1000 * 3e-6 + 200 * 1.5e-5) * 1.1, rel=1e-9)