fix - missing total token count calculation for maintain_format = Fal…

…se (getomni-ai#32) Co-authored-by: Pradyumna Singh Rathore <pradyumna.singhrathore@halliburton.com>
namtho7078 · Sep 12, 2024 · 3cd8b4a · 3cd8b4a
1 parent f0ffec3
commit 3cd8b4a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -157,5 +157,7 @@ dmypy.json
 # Ignore Cython debug files
 cython_debug/
 
-# Ignore custom files
+
+# Ignore dev files and folders
 _test_*.py
+dev_scripts
diff --git a/py_zerox/zerox/core/zerox.py b/py_zerox/zerox/core/zerox.py
@@ -123,6 +123,10 @@ async def zerox(
 
         aggregated_markdown = [result[0] for result in results if isinstance(result[0], str)]
 
+        ## add token usage
+        input_token_count += sum([result[1] for result in results])
+        output_token_count += sum([result[2] for result in results])
+
     # Write the aggregated markdown to a file
     if output_dir:
         result_file_path = os.path.join(output_dir, f"{file_name}.md")