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:
mateo-berri 2026-07-15 10:50:04 -07:00
parent b3e25df019
commit 47c57cf138
4 changed files with 14 additions and 2 deletions

View file

@ -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),

View file

@ -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,

View file

@ -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"

View file

@ -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():