Getting PaperQA to handle None tokens

Author: jamesbraza

src/paperqa/types.py

@@ -394,14 +394,14 @@ def add_tokens(self, result: LLMResult | Message) -> None:
                 prompt_count=result.info["usage"][0],
                 completion_count=result.info["usage"][1],
             )
+
+        prompt_count = result.prompt_count or 0
+        completion_count = result.completion_count or 0
         if result.model not in self.token_counts:
-            self.token_counts[result.model] = [
-                result.prompt_count,
-                result.completion_count,
-            ]
+            self.token_counts[result.model] = [prompt_count, completion_count]
         else:
-            self.token_counts[result.model][0] += result.prompt_count
-            self.token_counts[result.model][1] += result.completion_count
+            self.token_counts[result.model][0] += prompt_count
+            self.token_counts[result.model][1] += completion_count
 
         self.cost += result.cost
 

tests/test_paperqa.py

@@ -447,7 +447,9 @@ def accum(x) -> None:
     assert isinstance(first_id, UUID)
     assert completion.text
     assert completion.seconds_to_first_token > 0
+    assert completion.prompt_count is not None
     assert completion.prompt_count > 0
+    assert completion.completion_count is not None
     assert completion.completion_count > 0
     assert completion.model == "babbage-002"
     assert str(completion) == "".join(outputs)
@@ -461,7 +463,9 @@ def accum(x) -> None:
     assert completion.text
     assert completion.seconds_to_first_token == 0
     assert completion.seconds_to_last_token > 0
+    assert completion.prompt_count is not None
     assert completion.prompt_count > 0
+    assert completion.completion_count is not None
     assert completion.completion_count > 0
     try:
         assert completion.model == "babbage-002"
@@ -492,7 +496,9 @@ def accum(x) -> None:
         callbacks=[accum],
     )
     assert completion.seconds_to_first_token > 0
+    assert completion.prompt_count is not None
     assert completion.prompt_count > 0
+    assert completion.completion_count is not None
     assert completion.completion_count > 0
     assert str(completion) == "".join(outputs)
     assert isinstance(completion.text, str)