Report token_usage in FixResult (#1081)

report token usage

Author: clavedeluna

src/codemodder/codetf/v3/codetf.py

@@ -5,6 +5,7 @@
 
 from pydantic import BaseModel, model_validator
 
+from codemodder.llm import TokenUsage
 from codemodder.logging import logger
 
 from ..common import Change, CodeTFWriter, Finding, FixQuality
@@ -119,6 +120,27 @@ def validate_fixMetadata(self):
                 raise ValueError("fixMetadata must be provided for fixed results")
         return self
 
+    def store_token_data(self, token_usage: TokenUsage):
+        if token_usage == TokenUsage() or not self.fixMetadata:
+            return
+
+        if not self.fixMetadata.generation.ai:
+            self.fixMetadata.generation.ai = AIMetadata(
+                total_tokens=token_usage.total,
+                completion_tokens=token_usage.completion_tokens,
+                prompt_tokens=token_usage.prompt_tokens,
+            )
+        else:
+            self.fixMetadata.generation.ai.completion_tokens = (
+                self.fixMetadata.generation.ai.completion_tokens or 0
+            ) + token_usage.completion_tokens
+            self.fixMetadata.generation.ai.prompt_tokens = (
+                self.fixMetadata.generation.ai.prompt_tokens or 0
+            ) + token_usage.prompt_tokens
+            self.fixMetadata.generation.ai.total_tokens = (
+                self.fixMetadata.generation.ai.total_tokens or 0
+            ) + token_usage.total
+
 
 class CodeTF(CodeTFWriter, BaseModel):
     run: Run