modified: app.py

new file:   migrate.py
	modified:   requirements.txt

app.py

@@ -5,6 +5,8 @@ import eventlet
 
 from openai import OpenAI
 
+import tiktoken
+
 import os
 import time
 
@@ -48,6 +50,11 @@ class Message(db.Model):
         self.room_id = room_id
         self.token_count = self._count_tokens()
 
+    def _count_tokens(self):
+        # Replace 'gpt-3.5-turbo' with the model you are using.
+        encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+        return len(encoding.encode(self.content))
+
     def _count_tokens(self):
         return len(self.content.split())
 
@@ -55,7 +62,6 @@ class Message(db.Model):
         return '<img src="data:image/jpeg;base64,' in self.content
 
 
-
 def get_room(room_name):
     """Utility function to get room from room name."""
     room = Room.query.filter_by(name=room_name).first()
@@ -152,7 +158,10 @@ def on_join(data):
     )
     emit(
         "message",
-        {"id": None, "content": f"Estimated {total_token_count} total tokens in conversation."},
+        {
+            "id": None,
+            "content": f"Estimated {total_token_count} total tokens in conversation.",
+        },
         room=request.sid,
     )
 

migrate.py

@@ -0,0 +1,12 @@
+from app import db, app
+from app import Message
+
+def update_token_counts():
+    with app.app_context():
+        messages = Message.query.all()
+        for message in messages:
+            message.token_count = message._count_tokens()
+            db.session.add(message)
+        db.session.commit()
+
+update_token_counts()

requirements.txt

@@ -2,6 +2,7 @@ flask
 flask-socketio
 eventlet
 openai
+tiktoken
 
 # sqlite
 Flask-SQLAlchemy