diff --git a/src/codemodder/codetf.py b/src/codemodder/codetf.py
index 779d319b..5ccf4389 100644
--- a/src/codemodder/codetf.py
+++ b/src/codemodder/codetf.py
@@ -91,6 +91,8 @@ class AIMetadata(BaseModel):
     provider: Optional[str] = None
     model: Optional[str] = None
     tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    prompt_tokens: Optional[int] = None
 
 
 class Strategy(Enum):