From f5c6316da6fdd1d97128430bf360e8e862f8655a Mon Sep 17 00:00:00 2001
From: trisongz <4735784+trisongz@users.noreply.github.com>
Date: Wed, 6 Mar 2024 23:58:55 -0600
Subject: [PATCH] update handling of usage

---
 async_openai/types/functions.py |  7 ++++++-
 async_openai/types/resources.py | 30 ++++++++++++++++++------------
 async_openai/version.py         |  2 +-
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/async_openai/types/functions.py b/async_openai/types/functions.py
index 8def67b..6e8e021 100644
--- a/async_openai/types/functions.py
+++ b/async_openai/types/functions.py
@@ -107,7 +107,12 @@ def _set_values_from_response(
         Sets the values from the response
         """
         if name: self.function_name = name
-        self.function_usage = response.usage
+        usage = response.usage
+        if isinstance(usage, dict): 
+            from async_openai.types.resources import Usage
+            usage = Usage(**usage)
+            
+        self.function_usage = usage
         if response.response_ms: self.function_duration = response.response_ms / 1000
         self.function_model = response.model
         if client_name: self.function_client_name = client_name
diff --git a/async_openai/types/resources.py b/async_openai/types/resources.py
index 6d56ce7..b3b3984 100644
--- a/async_openai/types/resources.py
+++ b/async_openai/types/resources.py
@@ -55,7 +55,8 @@ class Usage(BaseModel):
     completion_tokens: Optional[int] = 0
     total_tokens: Optional[int] = 0
 
-    @lazyproperty
+    # @lazyproperty
+    @property
     def consumption(self) -> int:
         """
         Gets the consumption
@@ -66,17 +67,22 @@ def update(self, usage: Union['Usage', Dict[str, int]]):
         """
         Updates the consumption
         """
-        for key in {
-            'prompt_tokens',
-            'completion_tokens',
-            'total_tokens',
-        }:
-            if not hasattr(self, key):
-                setattr(self, key, 0)
-            val = usage.get(key, 0) if isinstance(usage, dict) else getattr(usage, key, 0)
-            setattr(self, key, getattr(self, key) + val)
-
-
+        if isinstance(usage, Usage):
+            if usage.prompt_tokens: self.prompt_tokens += usage.prompt_tokens
+            if usage.completion_tokens: self.completion_tokens += usage.completion_tokens
+            if usage.total_tokens: self.total_tokens += usage.total_tokens
+            return
+        
+        if usage.get('prompt_tokens'): self.prompt_tokens += usage['prompt_tokens']
+        if usage.get('completion_tokens'): self.completion_tokens += usage['completion_tokens']
+        if usage.get('total_tokens'): self.total_tokens += usage['total_tokens']
+
+    def __iadd__(self, other: Union['Usage', Dict[str, int]]):
+        """
+        Adds the usage
+        """
+        self.update(other)
+        return self.consumption
 
 
 class BaseResource(BaseModel):
diff --git a/async_openai/version.py b/async_openai/version.py
index 2d510a4..d33e17e 100644
--- a/async_openai/version.py
+++ b/async_openai/version.py
@@ -1 +1 @@
-VERSION = '0.0.53rc0'
\ No newline at end of file
+VERSION = '0.0.53rc1'
\ No newline at end of file