mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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:
parent
b95a4d303d
commit
41afb13fc6
2 changed files with 41 additions and 5 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue