diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index a4aa3a8a9f7..d1a3fccaac5 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -83,7 +83,7 @@ class BackgroundInteractionCreateContext(BaseModel): model_config = ConfigDict(frozen=True) - model: str + model: str | None call_type: str litellm_call_id: str function_id: str @@ -275,7 +275,7 @@ def _rebuild_logging_obj(create_context: BackgroundInteractionCreateContext) -> from litellm.litellm_core_utils.litellm_logging import Logging logging_obj: Final = Logging( - model=create_context.model, + model=create_context.model, # pyright: ignore[reportArgumentType] # function_setup builds the live object with the same None for an agent-only create messages=None, stream=False, call_type=create_context.call_type, diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 9dbb2c2f279..40d63090054 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -9,6 +9,7 @@ import pytest from litellm.interactions.background_cost_polling import ( _create_context, _poll_intervals, + _rebuild_logging_obj, BackgroundInteractionPollContext, InMemoryBackgroundSettlementStore, maybe_schedule_background_interaction_cost_polling, @@ -270,6 +271,37 @@ async def test_schedule_creates_poll_task_for_in_progress_create(): await task +@pytest.mark.asyncio +async def test_schedule_registers_an_agent_only_create_that_names_no_model(): + logging_obj = LitellmLogging( + model=None, + messages=None, + stream=False, + call_type="acreate_interaction", + start_time=time.time(), + litellm_call_id="bg-agent-call-id", + function_id="bg-agent-fn-id", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}, custom_llm_provider="gemini") + store = InMemoryBackgroundSettlementStore() + + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=store, + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + pending = await store.pending("interactions/bg-abc") + assert pending is not None + assert pending.create_context.model is None + assert _rebuild_logging_obj(pending.create_context).model is None + + @pytest.mark.asyncio @pytest.mark.parametrize( "response,create_kwargs",