From 06a08671617d24bf660ea3bede3facf02f17eec9 Mon Sep 17 00:00:00 2001
From: Russell Ballestrini <russell.ballestrini@gmail.com>
Date: Thu, 7 Dec 2023 09:33:35 -0500
Subject: [PATCH] token count to message object 	modified:   app.py 	new
 file:   migrations/versions/190d5ef26e20_add_token_count_to_message.py

---
 app.py                                        | 29 ++++++++----
 ...190d5ef26e20_add_token_count_to_message.py | 45 +++++++++++++++++++
 2 files changed, 66 insertions(+), 8 deletions(-)
 create mode 100644 migrations/versions/190d5ef26e20_add_token_count_to_message.py

diff --git a/app.py b/app.py
index cbb1885..1831dd4 100644
--- a/app.py
+++ b/app.py
@@ -39,12 +39,21 @@ class Message(db.Model):
     id = db.Column(db.Integer, primary_key=True)
     username = db.Column(db.String(128), nullable=False)
     content = db.Column(db.String(1024), nullable=False)
+    token_count = db.Column(db.Integer)
     room_id = db.Column(db.Integer, db.ForeignKey("room.id"), nullable=False)
 
     def __init__(self, username, content, room_id):
         self.username = username
         self.content = content
         self.room_id = room_id
+        self.token_count = self._count_tokens()
+
+    def _count_tokens(self):
+        return len(self.content.split())
+
+    def is_base64_image(self):
+        return '<img src="data:image/jpeg;base64,' in self.content
+
 
 
 def get_room(room_name):
@@ -104,11 +113,16 @@ def on_join(data):
     # Fetch previous messages from the database
     previous_messages = Message.query.filter_by(room_id=room.id).all()
 
+    # count the number of tokens in this room.
+    total_token_count = 0
+
     # Send the history of messages only to the newly connected client.
     # The reason for using `request.sid` here is to target the specific session (or client) that
     # just connected, so only they receive the backlog of messages, rather than broadcasting
     # this information to all clients in the room.
     for message in previous_messages:
+        if not message.is_base64_image():
+            total_token_count += message.token_count
         emit(
             "previous_messages",
             {
@@ -136,6 +150,11 @@ def on_join(data):
         {"id": None, "content": f"{data['username']} has joined the room."},
         room=room.name,
     )
+    emit(
+        "message",
+        {"id": None, "content": f"Estimated {total_token_count} total tokens in conversation."},
+        room=request.sid,
+    )
 
 
 @socketio.on("message")
@@ -256,11 +275,8 @@ def chat_claude(username, room_name, message, model_name="anthropic.claude-v1"):
 
     chat_history = ""
 
-    def is_base64_image(content):
-        return '<img src="data:image/jpeg;base64,' in content
-
     for msg in reversed(all_messages):
-        if is_base64_image(msg.content):
+        if msg.is_base64_image():
             continue
         if msg.username in [
             "gpt-3.5-turbo",
@@ -410,9 +426,6 @@ def chat_gpt(username, room_name, message, model_name="gpt-3.5-turbo"):
             updated_room_data = {"id": room.id, "name": room.name, "title": room.title}
             socketio.emit("update_room_list", updated_room_data, room=None)
 
-        def is_base64_image(content):
-            return '<img src="data:image/jpeg;base64,' in content
-
         chat_history = [
             {
                 "role": "system"
@@ -427,7 +440,7 @@ def chat_gpt(username, room_name, message, model_name="gpt-3.5-turbo"):
                 "content": f"{msg.username}: {msg.content}",
             }
             for msg in reversed(last_messages)
-            if not is_base64_image(msg.content)
+            if not msg.is_base64_image()
         ]
 
         chat_history.append({"role": "user", "content": f"{message}\n\n{model_name}: "})
diff --git a/migrations/versions/190d5ef26e20_add_token_count_to_message.py b/migrations/versions/190d5ef26e20_add_token_count_to_message.py
new file mode 100644
index 0000000..5be3089
--- /dev/null
+++ b/migrations/versions/190d5ef26e20_add_token_count_to_message.py
@@ -0,0 +1,45 @@
+"""Add token_count to Message
+
+Revision ID: 190d5ef26e20
+Revises: a9e886c56482
+Create Date: 2023-12-07 08:55:50.378439
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+from sqlalchemy.orm import Session
+
+# revision identifiers, used by Alembic.
+revision = '190d5ef26e20'
+down_revision = 'a9e886c56482'
+branch_labels = None
+depends_on = None
+
+from app import Message
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('message', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('token_count', sa.Integer(), nullable=True))
+
+    # Use this binding to connect to the database
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    # Assuming the content column is not nullable and always has a value
+    messages = session.query(Message).all()
+    for message in messages:
+        message.token_count = len(message.content.split())
+        session.add(message)
+
+    session.commit()
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('message', schema=None) as batch_op:
+        batch_op.drop_column('token_count')
+
+    # ### end Alembic commands ###