From 962e75eb70579edcae36c2e9ccfbaf2f1b312318 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 27 Oct 2023 09:32:10 -0700
Subject: [PATCH] (feat) create a usage class in ModelResponse, use it for
 anthropic

---
 litellm/utils.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/litellm/utils.py b/litellm/utils.py
index 2d7c233462..6754116454 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -146,6 +146,16 @@ class Choices(OpenAIObject):
         else:
             self.message = message
 
+class Usage(OpenAIObject):
+    def __init__(self, prompt_tokens=None, completion_tokens=None, total_tokens=None, **params):
+        super(Usage, self).__init__(**params)
+        if prompt_tokens:
+            self.prompt_tokens = prompt_tokens
+        if completion_tokens:
+            self.completion_tokens = completion_tokens
+        if total_tokens:
+            self.total_tokens = total_tokens
+
 class StreamingChoices(OpenAIObject):
     def __init__(self, finish_reason=None, index=0, delta: Optional[Delta]=None, **params):
         super(StreamingChoices, self).__init__(**params)
@@ -180,15 +190,11 @@ class ModelResponse(OpenAIObject):
         else:
             self._response_ms = None
         self.model = model
-        self.usage = (
-            usage
-            if usage
-            else {
-                "prompt_tokens": None,
-                "completion_tokens": None,
-                "total_tokens": None,
-            }
-        )
+        if usage:
+            self.usage = usage
+        else:
+            print("using the usage class")
+            self.usage = Usage()
         super(ModelResponse, self).__init__(**params)
 
     def to_dict_recursive(self):