mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(interactions): restore provider-scoped model attribution on foreign-pod settlement
Persist the router deployment string in the settlement context and carry it into the rebuilt logging metadata so reconstruct_model_name yields the same spend log model string as the create-path billing.
This commit is contained in:
parent
b3e25df019
commit
47c57cf138
4 changed files with 14 additions and 2 deletions
|
|
@ -86,6 +86,7 @@ class BackgroundSettlementContext(BaseModel):
|
|||
custom_llm_provider: str
|
||||
model: str
|
||||
model_group: Optional[str] = None
|
||||
deployment: Optional[str] = None
|
||||
litellm_call_id: str
|
||||
litellm_trace_id: Optional[str] = None
|
||||
call_type: str
|
||||
|
|
@ -209,11 +210,13 @@ def build_settlement_context(
|
|||
metadata = get_litellm_metadata_from_kwargs(kwargs=context.logging_obj.model_call_details)
|
||||
reservation_dict = metadata.get("user_api_key_budget_reservation")
|
||||
model_group = metadata.get("deployment_model_name")
|
||||
deployment = metadata.get("deployment")
|
||||
return BackgroundSettlementContext(
|
||||
interaction_id=context.interaction_id,
|
||||
custom_llm_provider=context.custom_llm_provider,
|
||||
model=str(context.logging_obj.model or context.logging_obj.model_call_details.get("model")),
|
||||
model_group=model_group if isinstance(model_group, str) else None,
|
||||
deployment=deployment if isinstance(deployment, str) else None,
|
||||
litellm_call_id=str(context.logging_obj.litellm_call_id),
|
||||
litellm_trace_id=context.logging_obj.model_call_details.get("litellm_trace_id"),
|
||||
call_type=str(context.logging_obj.call_type),
|
||||
|
|
|
|||
|
|
@ -128,7 +128,11 @@ def rebuild_logging_for_settlement(context: BackgroundSettlementContext) -> "Lit
|
|||
)
|
||||
attribution = {key: value for key, value in context.attribution.model_dump().items() if value is not None}
|
||||
reservation = context.budget_reservation.model_dump() if context.budget_reservation is not None else None
|
||||
metadata = attribution if reservation is None else {**attribution, "user_api_key_budget_reservation": reservation}
|
||||
optional_entries = (
|
||||
("user_api_key_budget_reservation", reservation),
|
||||
("deployment", context.deployment),
|
||||
)
|
||||
metadata = {**attribution, **{key: value for key, value in optional_entries if value is not None}}
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={
|
||||
"litellm_call_id": context.litellm_call_id,
|
||||
|
|
|
|||
|
|
@ -390,6 +390,7 @@ def _attributed_litellm_params(reservation: Optional[dict] = None) -> dict:
|
|||
"user_api_key_user_id": "user-123",
|
||||
"user_api_key_team_id": "team-456",
|
||||
"deployment_model_name": "gemini-3-flash-preview",
|
||||
"deployment": "gemini/gemini-3-flash-preview",
|
||||
}
|
||||
if reservation is not None:
|
||||
metadata["user_api_key_budget_reservation"] = reservation
|
||||
|
|
@ -415,6 +416,7 @@ async def test_poller_persists_settlement_context_with_attribution_and_reservati
|
|||
assert context.attribution.user_api_key_team_id == "team-456"
|
||||
assert context.model == "gemini-2.5-flash"
|
||||
assert context.model_group == "gemini-3-flash-preview"
|
||||
assert context.deployment == "gemini/gemini-3-flash-preview"
|
||||
assert context.custom_llm_provider == "gemini"
|
||||
assert context.call_type == "acreate_interaction"
|
||||
assert context.litellm_call_id == "bg-interactions-call-id"
|
||||
|
|
|
|||
|
|
@ -39,12 +39,14 @@ INTERACTION_ID = "interactions/bg-foreign"
|
|||
def _settlement_context(
|
||||
reservation: Optional[SettlementReservation] = None,
|
||||
model_group: Optional[str] = None,
|
||||
deployment: Optional[str] = None,
|
||||
) -> BackgroundSettlementContext:
|
||||
return BackgroundSettlementContext(
|
||||
interaction_id=INTERACTION_ID,
|
||||
custom_llm_provider="gemini",
|
||||
model="gemini-2.5-flash",
|
||||
model_group=model_group,
|
||||
deployment=deployment,
|
||||
litellm_call_id="original-create-call-id",
|
||||
call_type="acreate_interaction",
|
||||
attribution=SettlementKeyAttribution(
|
||||
|
|
@ -117,7 +119,7 @@ class _FakeRowStore:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rebuilt_logging_bills_with_original_attribution_and_request_id():
|
||||
logging_obj = rebuild_logging_for_settlement(_settlement_context())
|
||||
logging_obj = rebuild_logging_for_settlement(_settlement_context(deployment="gemini/gemini-2.5-flash"))
|
||||
response = _response("completed", with_usage=True)
|
||||
|
||||
await logging_obj.async_log_background_interaction_completion(result=response)
|
||||
|
|
@ -134,6 +136,7 @@ async def test_rebuilt_logging_bills_with_original_attribution_and_request_id():
|
|||
assert payload["team_id"] == "team-456"
|
||||
assert payload["api_key"]
|
||||
assert payload["spend"] > 0
|
||||
assert payload["model"] == "gemini/gemini-2.5-flash"
|
||||
|
||||
|
||||
def test_rebuilt_logging_start_time_is_naive_for_duration_math():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue