From f556e1d2e3534afa2d1f0ed492dbee0af78b0a9d Mon Sep 17 00:00:00 2001 From: ulleo Date: Fri, 11 Jul 2025 17:32:51 +0800 Subject: [PATCH] feat: save token usage --- .../versions/023_modify_chat_record.py | 75 +++++++++++++++++++ backend/apps/chat/curd/chat.py | 31 ++++++-- backend/apps/chat/models/chat_model.py | 12 +-- backend/apps/chat/task/llm.py | 32 +++++++- 4 files changed, 137 insertions(+), 13 deletions(-) create mode 100644 backend/alembic/versions/023_modify_chat_record.py diff --git a/backend/alembic/versions/023_modify_chat_record.py b/backend/alembic/versions/023_modify_chat_record.py new file mode 100644 index 00000000..cbe236c8 --- /dev/null +++ b/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 ### diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index cd143ae0..0c399bda 100644 --- a/backend/apps/chat/curd/chat.py +++ b/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,7 +274,7 @@ 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() @@ -275,6 +282,9 @@ def save_full_predict_message_and_answer(session: SessionDep, record_id: int, an 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) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 9502a21d..9832a432 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/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) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 6c94d05f..674d644c 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -9,7 +9,7 @@ import pandas as pd 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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 @@ class LLMService: 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