diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1f780ffcc3..fd11cf93a54 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1428,7 +1428,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1521,6 +1521,14 @@ class Usage(SafeAttributeModel, CompletionUsage): "cache_creation_input_tokens" ] + # Handle server_tool_use - convert dict to ServerToolUse object if needed + _server_tool_use: Optional[ServerToolUse] = None + if server_tool_use is not None: + if isinstance(server_tool_use, dict): + _server_tool_use = ServerToolUse(**server_tool_use) + elif isinstance(server_tool_use, ServerToolUse): + _server_tool_use = server_tool_use + super().__init__( prompt_tokens=prompt_tokens or 0, completion_tokens=completion_tokens or 0, @@ -1529,8 +1537,8 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) - if server_tool_use is not None: - self.server_tool_use = server_tool_use + if _server_tool_use is not None: + self.server_tool_use = _server_tool_use else: # maintain openai compatibility in usage object if possible del self.server_tool_use