feat: save token usage

Author: ulleo

backend/alembic/versions/023_modify_chat_record.py

@@ -0,0 +1,75 @@
+"""023_modify_chat_record
+
+Revision ID: f535d09946f6
+Revises: e6b20ae73606
+Create Date: 2025-07-11 15:36:18.473133
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = 'f535d09946f6'
+down_revision = 'e6b20ae73606'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat_record', 'token_sql',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_chart',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_analysis',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_predict',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_recommended_question',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_select_datasource_question',
+               existing_type=sa.INTEGER(),
+               type_=sqlmodel.sql.sqltypes.AutoString(length=256),
+               existing_nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat_record', 'token_select_datasource_question',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_recommended_question',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_predict',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_analysis',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_chart',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'token_sql',
+               existing_type=sqlmodel.sql.sqltypes.AutoString(length=256),
+               type_=sa.INTEGER(),
+               existing_nullable=True)
+    # ### end Alembic commands ###

backend/apps/chat/curd/chat.py

@@ -229,13 +229,17 @@ def save_full_sql_message(session: SessionDep, record_id: int, full_message: str
     return save_full_sql_message_and_answer(session=session, record_id=record_id, full_message=full_message, answer='')
 
 
-def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer: str, full_message: str) -> ChatRecord:
+def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer: str, full_message: str,
+                                     token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.full_sql_message = full_message
     record.sql_answer = answer
 
+    if token_usage:
+        record.token_sql = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)
@@ -248,13 +252,16 @@ def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer
 
 
 def save_full_analysis_message_and_answer(session: SessionDep, record_id: int, answer: str,
-                                          full_message: str) -> ChatRecord:
+                                          full_message: str, token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.full_analysis_message = full_message
     record.analysis = answer
 
+    if token_usage:
+        record.token_analysis = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)
@@ -267,14 +274,17 @@ def save_full_analysis_message_and_answer(session: SessionDep, record_id: int, a
 
 
 def save_full_predict_message_and_answer(session: SessionDep, record_id: int, answer: str,
-                                         full_message: str, data: str) -> ChatRecord:
+                                         full_message: str, data: str, token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.full_predict_message = full_message
     record.predict = answer
     record.predict_data = data
 
+    if token_usage:
+        record.token_predict = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)
@@ -288,7 +298,7 @@ def save_full_predict_message_and_answer(session: SessionDep, record_id: int, an
 
 def save_full_select_datasource_message_and_answer(session: SessionDep, record_id: int, answer: str,
                                                    full_message: str, datasource: int = None,
-                                                   engine_type: str = None) -> ChatRecord:
+                                                   engine_type: str = None, token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
@@ -299,6 +309,9 @@ def save_full_select_datasource_message_and_answer(session: SessionDep, record_i
         record.datasource = datasource
         record.engine_type = engine_type
 
+    if token_usage:
+        record.token_select_datasource_question = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)
@@ -311,7 +324,7 @@ def save_full_select_datasource_message_and_answer(session: SessionDep, record_i
 
 
 def save_full_recommend_question_message_and_answer(session: SessionDep, record_id: int, answer: dict = None,
-                                                    full_message: str = '[]') -> ChatRecord:
+                                                    full_message: str = '[]', token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
@@ -329,6 +342,9 @@ def save_full_recommend_question_message_and_answer(session: SessionDep, record_
             pass
     record.recommended_question = json_str
 
+    if token_usage:
+        record.token_recommended_question = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)
@@ -363,13 +379,16 @@ def save_full_chart_message(session: SessionDep, record_id: int, full_message: s
 
 
 def save_full_chart_message_and_answer(session: SessionDep, record_id: int, answer: str,
-                                       full_message: str) -> ChatRecord:
+                                       full_message: str, token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.full_chart_message = full_message
     record.chart_answer = answer
 
+    if token_usage:
+        record.token_chart = orjson.dumps(token_usage).decode()
+
     result = ChatRecord(**record.model_dump())
 
     session.add(record)

backend/apps/chat/models/chat_model.py

@@ -49,17 +49,17 @@ class ChatRecord(SQLModel, table=True):
     recommended_question: str = Field(sa_column=Column(Text, nullable=True))
     datasource_select_answer: str = Field(sa_column=Column(Text, nullable=True))
     full_sql_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_sql: int = Field(default=0, nullable=True)
+    token_sql: str = Field(max_length=256, nullable=True)
     full_chart_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_chart: int = Field(default=0, nullable=True)
+    token_chart: str = Field(max_length=256, nullable=True)
     full_analysis_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_analysis: int = Field(default=0, nullable=True)
+    token_analysis: str = Field(max_length=256, nullable=True)
     full_predict_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_predict: int = Field(default=0, nullable=True)
+    token_predict: str = Field(max_length=256, nullable=True)
     full_recommended_question_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_recommended_question: int = Field(default=0, nullable=True)
+    token_recommended_question: str = Field(max_length=256, nullable=True)
     full_select_datasource_message: str = Field(sa_column=Column(Text, nullable=True))
-    token_select_datasource_question: int = Field(default=0, nullable=True)
+    token_select_datasource_question: str = Field(max_length=256, nullable=True)
     finish: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
     error: str = Field(sa_column=Column(Text, nullable=True))
     run_time: float = Field(default=0)

backend/apps/chat/task/llm.py

@@ -9,7 +9,7 @@
 import requests
 from langchain.chat_models.base import BaseChatModel
 from langchain_community.utilities import SQLDatabase
-from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
+from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk
 from sqlalchemy import select
 from sqlalchemy.orm import load_only
 
@@ -198,6 +198,7 @@ def generate_analysis(self):
         full_thinking_text = ''
         full_analysis_text = ''
         res = self.llm.stream(analysis_msg)
+        token_usage = {}
         for chunk in res:
             print(chunk)
             reasoning_content_chunk = ''
@@ -211,9 +212,11 @@ def generate_analysis(self):
 
             full_analysis_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
 
         analysis_msg.append(AIMessage(full_analysis_text))
         self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id,
+                                                            token_usage=token_usage,
                                                             answer=orjson.dumps({'content': full_analysis_text,
                                                                                  'reasoning_content': full_thinking_text}).decode(),
                                                             full_message=orjson.dumps(history_msg +
@@ -245,6 +248,7 @@ def generate_predict(self):
         full_thinking_text = ''
         full_predict_text = ''
         res = self.llm.stream(predict_msg)
+        token_usage = {}
         for chunk in res:
             print(chunk)
             reasoning_content_chunk = ''
@@ -258,9 +262,11 @@ def generate_predict(self):
 
             full_predict_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
 
         predict_msg.append(AIMessage(full_predict_text))
         self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id,
+                                                           token_usage=token_usage,
                                                            answer=orjson.dumps({'content': full_predict_text,
                                                                                 'reasoning_content': full_thinking_text}).decode(),
                                                            data='',
@@ -291,6 +297,7 @@ def generate_recommend_questions_task(self):
                                                                                                  guess_msg]).decode())
         full_thinking_text = ''
         full_guess_text = ''
+        token_usage = {}
         res = self.llm.stream(guess_msg)
         for chunk in res:
             print(chunk)
@@ -305,9 +312,11 @@ def generate_recommend_questions_task(self):
 
             full_guess_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
 
         guess_msg.append(AIMessage(full_guess_text))
         self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id,
+                                                                      token_usage=token_usage,
                                                                       answer={'content': full_guess_text,
                                                                               'reasoning_content': full_thinking_text},
                                                                       full_message=orjson.dumps([{'type': msg.type,
@@ -342,6 +351,7 @@ def select_datasource(self):
                                                                                                 datasource_msg]).decode())
         full_thinking_text = ''
         full_text = ''
+        token_usage = {}
         res = self.llm.stream(datasource_msg)
         for chunk in res:
             print(chunk)
@@ -356,6 +366,7 @@ def select_datasource(self):
 
             full_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
         datasource_msg.append(AIMessage(full_text))
 
         json_str = extract_nested_json(full_text)
@@ -418,6 +429,7 @@ def generate_sql(self):
                                                  self.sql_message]).decode())
         full_thinking_text = ''
         full_sql_text = ''
+        token_usage = {}
         res = self.llm.stream(self.sql_message)
         for chunk in res:
             print(chunk)
@@ -432,9 +444,11 @@ def generate_sql(self):
 
             full_sql_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
 
         self.sql_message.append(AIMessage(full_sql_text))
         self.record = save_full_sql_message_and_answer(session=self.session, record_id=self.record.id,
+                                                       token_usage=token_usage,
                                                        answer=orjson.dumps({'content': full_sql_text,
                                                                             'reasoning_content': full_thinking_text}).decode(),
                                                        full_message=orjson.dumps(
@@ -450,6 +464,7 @@ def generate_chart(self):
                                                    self.chart_message]).decode())
         full_thinking_text = ''
         full_chart_text = ''
+        token_usage = {}
         res = self.llm.stream(self.chart_message)
         for chunk in res:
             print(chunk)
@@ -464,9 +479,11 @@ def generate_chart(self):
 
             full_chart_text += chunk.content
             yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+            get_token_usage(chunk, token_usage)
 
         self.chart_message.append(AIMessage(full_chart_text))
         self.record = save_full_chart_message_and_answer(session=self.session, record_id=self.record.id,
+                                                         token_usage=token_usage,
                                                          answer=orjson.dumps({'content': full_chart_text,
                                                                               'reasoning_content': full_thinking_text}).decode(),
                                                          full_message=orjson.dumps(
@@ -740,6 +757,9 @@ def run_analysis_or_predict_task(llm_service: LLMService, action_type: str):
         traceback.print_exc()
         # llm_service.save_error(session=session, message=str(e))
         yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+    finally:
+        # end
+        pass
 
 
 def run_recommend_questions_task(llm_service: LLMService):
@@ -788,3 +808,13 @@ def request_picture(chat_id: int, record_id: int, chart: dict, data: dict):
     requests.post(url=settings.MCP_IMAGE_HOST, json=request_obj)
 
     return f'{(settings.SERVER_IMAGE_HOST if settings.SERVER_IMAGE_HOST[-1] == "/" else (settings.SERVER_IMAGE_HOST + "/"))}{file_name}.png'
+
+
+def get_token_usage(chunk: BaseMessageChunk, token_usage: dict = {}):
+    try:
+        if chunk.usage_metadata:
+            token_usage['input_tokens'] = chunk.usage_metadata.get('input_tokens')
+            token_usage['output_tokens'] = chunk.usage_metadata.get('output_tokens')
+            token_usage['total_tokens'] = chunk.usage_metadata.get('total_tokens')
+    except Exception:
+        pass