Add tokens counter

Palm1r · Oct 11, 2024 · 47f192a · 47f192a
1 parent fb0935f
commit 47f192a
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 2 deletions.
diff --git a/chatview/ChatModel.cpp b/chatview/ChatModel.cpp
@@ -101,6 +101,7 @@ void ChatModel::addMessage(const QString &content, ChatRole role)
     m_totalTokens += tokenCount;
     endInsertRows();
     trim();
+    emit totalTokensChanged();
 }
 
 void ChatModel::clear()
@@ -135,4 +136,9 @@ QJsonArray ChatModel::prepareMessagesForRequest() const
     return messages;
 }
 
+int ChatModel::totalTokens() const
+{
+    return m_totalTokens;
+}
+
 } // namespace QodeAssist::Chat
diff --git a/chatview/ChatModel.hpp b/chatview/ChatModel.hpp
@@ -29,7 +29,7 @@ class ChatModel : public QAbstractListModel
 {
     Q_OBJECT
     QML_ELEMENT
-    QML_UNCREATABLE("ChatModel no need to create")
+    Q_PROPERTY(int totalTokens READ totalTokens NOTIFY totalTokensChanged FINAL)
 public:
     enum Roles { RoleType = Qt::UserRole, Content };
 
@@ -57,6 +57,11 @@ class ChatModel : public QAbstractListModel
     void setSystemPrompt(const QString &prompt);
     QJsonArray prepareMessagesForRequest() const;
 
+    int totalTokens() const;
+
+signals:
+    void totalTokensChanged();
+
 private:
     void trim();
     int estimateTokenCount(const QString &text) const;

diff --git a/chatview/qml/RootItem.qml b/chatview/qml/RootItem.qml
@@ -12,8 +12,23 @@ ChatRootView {
         color: root.backgroundColor
     }
 
+    Rectangle {
+        width: tokensCounter.implicitWidth + 4
+        height: childrenRect.height
+        color: "lightgreen"
+        radius: 4
+
+        Text {
+            id: tokensCounter
+
+            anchors.horizontalCenter: parent.horizontalCenter
+            text: root.chatModel.totalTokens
+        }
+    }
+
     ColumnLayout {
         anchors.fill: parent
+        anchors.topMargin: 30
         spacing: 10
 
         ListView {
@@ -28,7 +43,7 @@ ChatRootView {
                 width: ListView.view.width
                 height: messageText.height + 20
                 color: model.roleType === ChatModel.User ? "#e6e6e6" : "#f0f0f0"
-                radius: 10
+                radius: 8
 
                 Text {
                     id: messageText