ka1kuk commited on
Commit
5484efc
1 Parent(s): d1edbcf

Update messagers/message_outputer.py

Browse files
Files changed (1) hide show
  1. messagers/message_outputer.py +5 -2
messagers/message_outputer.py CHANGED
@@ -1,5 +1,6 @@
1
  import json
2
  import time
 
3
 
4
  class OpenaiStreamOutputer:
5
  """
@@ -86,8 +87,10 @@ class OpenaiStreamOutputer:
86
  ]
87
 
88
  # Update token counts
 
 
89
  data["usage"]["prompt_tokens"] += tokens_count
90
- data["usage"]["completion_tokens"] += len(content.split())
91
- data["usage"]["total_tokens"] = data["usage"]["prompt_tokens"] + data["usage"]["completion_tokens"]
92
 
93
  return self.data_to_string(data, content_type)
 
1
  import json
2
  import time
3
+ import tiktoken
4
 
5
  class OpenaiStreamOutputer:
6
  """
 
87
  ]
88
 
89
  # Update token counts
90
+ tokens_used = len(content.split())
91
+ self.token_manager.add_tokens(tokens_used) # hypothetical method
92
  data["usage"]["prompt_tokens"] += tokens_count
93
+ data["usage"]["completion_tokens"] += tokens_used
94
+ data["usage"]["total_tokens"] = self.token_manager.get_total_tokens() # hypothetical method
95
 
96
  return self.data_to_string(data, content_type)