Implement simple add for TokenUsage (#1002)

pixee · Feb 20, 2025 · 0e3722f · 0e3722f
1 parent b38681d
commit 0e3722f
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/codemodder/llm.py b/src/codemodder/llm.py
@@ -127,7 +127,13 @@ class TokenUsage:
     completion_tokens: int = 0
     prompt_tokens: int = 0
 
-    def __iadd__(self, other: Self) -> Self:
+    def __add__(self, other: TokenUsage) -> TokenUsage:
+        return TokenUsage(
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+        )
+
+    def __iadd__(self, other: TokenUsage) -> Self:
         self.completion_tokens += other.completion_tokens
         self.prompt_tokens += other.prompt_tokens
         return self