Update messagers/message_outputer.py
Browse files
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"] +=
|
91 |
-
data["usage"]["total_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)
|