diff --git a/src/codemodder/codetf/v3/codetf.py b/src/codemodder/codetf/v3/codetf.py index f5ec09b8..b4036472 100644 --- a/src/codemodder/codetf/v3/codetf.py +++ b/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