klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: save token usage

Browse Source
This commit is contained in:
ulleo
2025-07-11 17:32:51 +08:00
parent 682c8b4b24
commit f556e1d2e3
4 changed files with 137 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/alembic/versions/023_modify_chat_record.py
+75
View File
@@ -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
+25 -6
View File
@@ -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)
backend/apps/chat/models/chat_model.py
+6 -6
View File
@@ -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
+31 -1
View File
@@ -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