token counts

src/chatdbg/assistant/assistant.py

@@ -213,7 +213,7 @@ def run(self, prompt, client_print=print):
             )
             client_print()
             client_print(f"[Cost: ~${cost:.2f} USD]")
-            return run.usage.total_tokens, cost, elapsed_time
+            return run.usage.total_tokens,run.usage.prompt_tokens, run.usage.completion_tokens, cost, elapsed_time
         except OpenAIError as e:
             client_print(f"*** OpenAI Error: {e}")
             sys.exit(-1)

src/chatdbg/chatdbg_pdb.py

@@ -471,8 +471,8 @@ def client_print(line=""):
         full_prompt = truncate_proportionally(full_prompt)
 
         self._log.push_chat(arg, full_prompt)
-        tokens, cost, time = self._assistant.run(full_prompt, client_print)
-        self._log.pop_chat(tokens, cost, time)
+        total_tokens, prompt_tokens, completion_tokens, cost, time = self._assistant.run(full_prompt, client_print)
+        self._log.pop_chat(total_tokens, prompt_tokens, completion_tokens, cost, time)
 
     def do_mark(self, arg):
         marks = [ 'Full', 'Partial', 'Wrong', 'None', '?' ]

src/chatdbg/ipdb_util/logging.py

@@ -106,9 +106,11 @@ def push_chat(self, line, full_prompt):
             }
         }
 
-    def pop_chat(self, tokens, cost, time):
+    def pop_chat(self, total_tokens, prompt_tokens, completion_tokens, cost, time):
         self.chat_step['stats'] = {
-            'tokens' : tokens,
+            'tokens' : total_tokens,
+            'prompt' : prompt_tokens,
+            'completion' : completion_tokens,
             'cost' : cost, 
             'time' : time
         }