mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
1034 lines
50 KiB
Python
1034 lines
50 KiB
Python
import concurrent
|
|
import json
|
|
import traceback
|
|
import warnings
|
|
from concurrent.futures import ThreadPoolExecutor, Future
|
|
from typing import Any, List, Optional, Union, Dict
|
|
|
|
import numpy as np
|
|
import orjson
|
|
import pandas as pd
|
|
import requests
|
|
import sqlparse
|
|
from langchain.chat_models.base import BaseChatModel
|
|
from langchain_community.utilities import SQLDatabase
|
|
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk
|
|
from sqlalchemy import and_, cast, or_
|
|
from sqlalchemy import select
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlbot_xpack.permissions.api.permission import transRecord2DTO
|
|
from sqlbot_xpack.permissions.models.ds_permission import DsPermission, PermissionDTO
|
|
from sqlbot_xpack.permissions.models.ds_rules import DsRules
|
|
|
|
from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config
|
|
from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \
|
|
save_error_message, save_sql_exec_data, save_full_chart_message, save_full_chart_message_and_answer, save_chart, \
|
|
finish_record, save_full_analysis_message_and_answer, save_full_predict_message_and_answer, save_predict_data, \
|
|
save_full_select_datasource_message_and_answer, save_full_recommend_question_message_and_answer, \
|
|
get_old_questions, save_analysis_predict_record, list_base_records, rename_chat
|
|
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat
|
|
from apps.datasource.crud.datasource import get_table_schema
|
|
from apps.datasource.crud.row_permission import transFilterTree
|
|
from apps.datasource.models.datasource import CoreDatasource, CoreTable
|
|
from apps.db.db import exec_sql
|
|
from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds
|
|
from apps.system.schemas.system_schema import AssistantOutDsSchema
|
|
from common.core.config import settings
|
|
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
|
|
from common.utils.utils import SQLBotLogUtil, extract_nested_json
|
|
|
|
warnings.filterwarnings("ignore")
|
|
|
|
base_message_count_limit = 5
|
|
|
|
executor = ThreadPoolExecutor(max_workers=200)
|
|
|
|
|
|
class LLMService:
|
|
ds: CoreDatasource
|
|
chat_question: ChatQuestion
|
|
record: ChatRecord
|
|
config: LLMConfig
|
|
llm: BaseChatModel
|
|
sql_message: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
chart_message: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
history_records: List[ChatRecord] = []
|
|
session: SessionDep
|
|
current_user: CurrentUser
|
|
current_assistant: Optional[CurrentAssistant] = None
|
|
out_ds_instance: Optional[AssistantOutDs] = None
|
|
change_title: bool = False
|
|
|
|
chunk_list: List[str] = []
|
|
future: Future
|
|
|
|
def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question: ChatQuestion,
|
|
current_assistant: Optional[CurrentAssistant] = None):
|
|
self.chunk_list = []
|
|
self.session = session
|
|
self.current_user = current_user
|
|
self.current_assistant = current_assistant
|
|
# chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
|
|
chat_id = chat_question.chat_id
|
|
chat: Chat = self.session.get(Chat, chat_id)
|
|
if not chat:
|
|
raise Exception(f"Chat with id {chat_id} not found")
|
|
ds: CoreDatasource | AssistantOutDsSchema | None = None
|
|
if chat.datasource:
|
|
# Get available datasource
|
|
# ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
|
|
if current_assistant and current_assistant.type == 1:
|
|
self.out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
|
|
ds = self.out_ds_instance.get_ds(chat.datasource)
|
|
if not ds:
|
|
raise Exception("No available datasource configuration found")
|
|
chat_question.engine = ds.type
|
|
chat_question.db_schema = self.out_ds_instance.get_db_schema(ds.id)
|
|
else:
|
|
ds = self.session.get(CoreDatasource, chat.datasource)
|
|
if not ds:
|
|
raise Exception("No available datasource configuration found")
|
|
chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
|
|
chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
|
|
|
|
history_records: List[ChatRecord] = list(
|
|
map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False,
|
|
list_base_records(session=self.session,
|
|
current_user=current_user,
|
|
chart_id=chat_id))))
|
|
self.change_title = len(history_records) == 0
|
|
|
|
chat_question.lang = current_user.language
|
|
|
|
self.ds = (ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None
|
|
self.chat_question = chat_question
|
|
self.config = get_default_config()
|
|
self.chat_question.ai_modal_id = self.config.model_id
|
|
|
|
# Create LLM instance through factory
|
|
llm_instance = LLMFactory.create_llm(self.config)
|
|
self.llm = llm_instance.llm
|
|
|
|
self.history_records = history_records
|
|
|
|
self.init_messages()
|
|
|
|
def is_running(self, timeout=0.5):
|
|
try:
|
|
r = concurrent.futures.wait([self.future], timeout)
|
|
if len(r.not_done) > 0:
|
|
return True
|
|
else:
|
|
return False
|
|
except Exception as e:
|
|
return True
|
|
|
|
def init_messages(self):
|
|
# self.agent_executor = create_react_agent(self.llm)
|
|
last_sql_messages = list(
|
|
filter(lambda r: True if r.full_sql_message is not None and r.full_sql_message.strip() != '' else False,
|
|
self.history_records))
|
|
last_sql_message_str = "[]" if last_sql_messages is None or len(last_sql_messages) == 0 else last_sql_messages[
|
|
-1].full_sql_message
|
|
|
|
last_chart_messages = list(
|
|
filter(
|
|
lambda r: True if r.full_chart_message is not None and r.full_chart_message.strip() != '' else False,
|
|
self.history_records))
|
|
last_chart_message_str = "[]" if last_chart_messages is None or len(last_chart_messages) == 0 else \
|
|
last_chart_messages[-1].full_chart_message
|
|
|
|
last_sql_messages: List[dict[str, Any]] = orjson.loads(last_sql_message_str)
|
|
|
|
# todo maybe can configure
|
|
count_limit = 0 - base_message_count_limit
|
|
|
|
self.sql_message = []
|
|
if last_sql_messages is None or len(last_sql_messages) == 0:
|
|
# add sys prompt
|
|
self.sql_message.append(SystemMessage(content=self.chat_question.sql_sys_question()))
|
|
else:
|
|
# limit count
|
|
for last_sql_message in last_sql_messages:
|
|
if last_sql_message['type'] == 'system':
|
|
_msg = SystemMessage(content=last_sql_message['content'])
|
|
self.sql_message.append(_msg)
|
|
break
|
|
for last_sql_message in last_sql_messages[count_limit:]:
|
|
_msg: BaseMessage
|
|
if last_sql_message['type'] == 'human':
|
|
_msg = HumanMessage(content=last_sql_message['content'])
|
|
self.sql_message.append(_msg)
|
|
elif last_sql_message['type'] == 'ai':
|
|
_msg = AIMessage(content=last_sql_message['content'])
|
|
self.sql_message.append(_msg)
|
|
|
|
last_chart_messages: List[dict[str, Any]] = orjson.loads(last_chart_message_str)
|
|
|
|
self.chart_message = []
|
|
if last_chart_messages is None or len(last_chart_messages) == 0:
|
|
# add sys prompt
|
|
self.chart_message.append(SystemMessage(content=self.chat_question.chart_sys_question()))
|
|
else:
|
|
# limit count
|
|
for last_chart_message in last_chart_messages:
|
|
if last_chart_message['type'] == 'system':
|
|
_msg = SystemMessage(content=last_chart_message['content'])
|
|
self.chart_message.append(_msg)
|
|
break
|
|
for last_chart_message in last_chart_messages:
|
|
_msg: BaseMessage
|
|
if last_chart_message['type'] == 'human':
|
|
_msg = HumanMessage(content=last_chart_message['content'])
|
|
self.chart_message.append(_msg)
|
|
elif last_chart_message['type'] == 'ai':
|
|
_msg = AIMessage(content=last_chart_message['content'])
|
|
self.chart_message.append(_msg)
|
|
|
|
def init_record(self) -> ChatRecord:
|
|
self.record = save_question(session=self.session, current_user=self.current_user, question=self.chat_question)
|
|
return self.record
|
|
|
|
def get_record(self):
|
|
return self.record
|
|
|
|
def set_record(self, record: ChatRecord):
|
|
self.record = record
|
|
|
|
def get_fields_from_chart(self):
|
|
chart_info = orjson.loads(self.record.chart)
|
|
fields = []
|
|
if chart_info.get('columns') and len(chart_info.get('columns')) > 0:
|
|
for column in chart_info.get('columns'):
|
|
column_str = column.get('value')
|
|
if column.get('value') != column.get('name'):
|
|
column_str = column_str + '(' + column.get('name') + ')'
|
|
fields.append(column_str)
|
|
if chart_info.get('axis'):
|
|
for _type in ['x', 'y', 'series']:
|
|
if chart_info.get('axis').get(_type):
|
|
column = chart_info.get('axis').get(_type)
|
|
column_str = column.get('value')
|
|
if column.get('value') != column.get('name'):
|
|
column_str = column_str + '(' + column.get('name') + ')'
|
|
fields.append(column_str)
|
|
return fields
|
|
|
|
def generate_analysis(self):
|
|
fields = self.get_fields_from_chart()
|
|
|
|
self.chat_question.fields = orjson.dumps(fields).decode()
|
|
self.chat_question.data = orjson.dumps(orjson.loads(self.record.data).get('data')).decode()
|
|
analysis_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question()))
|
|
analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question()))
|
|
|
|
history_msg = []
|
|
if self.record.full_analysis_message and self.record.full_analysis_message.strip() != '':
|
|
history_msg = orjson.loads(self.record.full_analysis_message)
|
|
|
|
self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='',
|
|
full_message=orjson.dumps(history_msg +
|
|
[{'type': msg.type,
|
|
'content': msg.content} for msg
|
|
in
|
|
analysis_msg]).decode())
|
|
full_thinking_text = ''
|
|
full_analysis_text = ''
|
|
res = self.llm.stream(analysis_msg)
|
|
token_usage = {}
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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 +
|
|
[{'type': msg.type,
|
|
'content': msg.content} for msg
|
|
in
|
|
analysis_msg]).decode())
|
|
|
|
def generate_predict(self):
|
|
fields = self.get_fields_from_chart()
|
|
|
|
self.chat_question.fields = orjson.dumps(fields).decode()
|
|
self.chat_question.data = orjson.dumps(orjson.loads(self.record.data).get('data')).decode()
|
|
predict_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
predict_msg.append(SystemMessage(content=self.chat_question.predict_sys_question()))
|
|
predict_msg.append(HumanMessage(content=self.chat_question.predict_user_question()))
|
|
|
|
history_msg = []
|
|
if self.record.full_predict_message and self.record.full_predict_message.strip() != '':
|
|
history_msg = orjson.loads(self.record.full_predict_message)
|
|
|
|
self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id, answer='',
|
|
data='',
|
|
full_message=orjson.dumps(history_msg +
|
|
[{'type': msg.type,
|
|
'content': msg.content} for msg
|
|
in
|
|
predict_msg]).decode())
|
|
full_thinking_text = ''
|
|
full_predict_text = ''
|
|
res = self.llm.stream(predict_msg)
|
|
token_usage = {}
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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='',
|
|
full_message=orjson.dumps(history_msg +
|
|
[{'type': msg.type,
|
|
'content': msg.content} for msg
|
|
in
|
|
predict_msg]).decode())
|
|
|
|
def generate_recommend_questions_task(self):
|
|
|
|
# get schema
|
|
if self.ds and not self.chat_question.db_schema:
|
|
self.chat_question.db_schema = self.out_ds_instance.get_db_schema(
|
|
self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session,
|
|
current_user=self.current_user, ds=self.ds)
|
|
|
|
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
|
|
|
|
old_questions = list(map(lambda q: q[0].strip(), get_old_questions(self.session, self.record.datasource)))
|
|
guess_msg.append(
|
|
HumanMessage(content=self.chat_question.guess_user_question(orjson.dumps(old_questions).decode())))
|
|
|
|
self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id,
|
|
full_message=orjson.dumps([{'type': msg.type,
|
|
'content': msg.content}
|
|
for msg
|
|
in
|
|
guess_msg]).decode())
|
|
full_thinking_text = ''
|
|
full_guess_text = ''
|
|
token_usage = {}
|
|
res = self.llm.stream(guess_msg)
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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,
|
|
'content': msg.content}
|
|
for msg
|
|
in
|
|
guess_msg]).decode())
|
|
yield {'recommended_question': self.record.recommended_question}
|
|
|
|
def select_datasource(self):
|
|
datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
|
|
if self.current_assistant:
|
|
_ds_list = get_assistant_ds(llm_service=self)
|
|
else:
|
|
oid: str = self.current_user.oid
|
|
stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(
|
|
CoreDatasource.oid == oid)
|
|
_ds_list = [
|
|
{
|
|
"id": ds.id,
|
|
"name": ds.name,
|
|
"description": ds.description
|
|
}
|
|
for ds in self.session.exec(stmt)
|
|
]
|
|
""" _ds_list = self.session.exec(select(CoreDatasource).options(
|
|
load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all() """
|
|
_ds_list_dict = []
|
|
for _ds in _ds_list:
|
|
_ds_list_dict.append(_ds)
|
|
datasource_msg.append(
|
|
HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
|
|
|
|
history_msg = []
|
|
if self.record.full_select_datasource_message and self.record.full_select_datasource_message.strip() != '':
|
|
history_msg = orjson.loads(self.record.full_select_datasource_message)
|
|
|
|
self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
|
|
answer='',
|
|
full_message=orjson.dumps(history_msg +
|
|
[{'type': msg.type,
|
|
'content': msg.content}
|
|
for msg
|
|
in
|
|
datasource_msg]).decode())
|
|
full_thinking_text = ''
|
|
full_text = ''
|
|
token_usage = {}
|
|
res = self.llm.stream(datasource_msg)
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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)
|
|
|
|
_error: Exception | None = None
|
|
_datasource: int | None = None
|
|
_engine_type: str | None = None
|
|
try:
|
|
data = orjson.loads(json_str)
|
|
|
|
if data['id'] and data['id'] != 0:
|
|
_datasource = data['id']
|
|
_chat = self.session.get(Chat, self.record.chat_id)
|
|
_chat.datasource = _datasource
|
|
if self.current_assistant and self.current_assistant.type == 1:
|
|
_ds = self.out_ds_instance.get_ds(data['id'])
|
|
self.ds = _ds
|
|
self.chat_question.engine = _ds.type
|
|
_engine_type = self.chat_question.engine
|
|
_chat.engine_type = _ds.type
|
|
else:
|
|
_ds = self.session.get(CoreDatasource, _datasource)
|
|
if not _ds:
|
|
_datasource = None
|
|
raise Exception(f"Datasource configuration with id {_datasource} not found")
|
|
self.ds = CoreDatasource(**_ds.model_dump())
|
|
self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
|
|
_engine_type = self.chat_question.engine
|
|
_chat.engine_type = _ds.type_name
|
|
# save chat
|
|
self.session.add(_chat)
|
|
self.session.flush()
|
|
self.session.refresh(_chat)
|
|
self.session.commit()
|
|
|
|
elif data['fail']:
|
|
raise Exception(data['fail'])
|
|
else:
|
|
raise Exception('No available datasource configuration found')
|
|
|
|
except Exception as e:
|
|
_error = e
|
|
|
|
self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
|
|
answer=orjson.dumps({'content': full_text,
|
|
'reasoning_content': full_thinking_text}).decode(),
|
|
datasource=_datasource,
|
|
engine_type=_engine_type,
|
|
full_message=orjson.dumps(history_msg +
|
|
[{'type': msg.type,
|
|
'content': msg.content}
|
|
for msg
|
|
in
|
|
datasource_msg]).decode())
|
|
self.init_messages()
|
|
|
|
if _error:
|
|
raise _error
|
|
|
|
def generate_sql(self):
|
|
# append current question
|
|
self.sql_message.append(HumanMessage(self.chat_question.sql_user_question()))
|
|
self.record = save_full_sql_message(session=self.session, record_id=self.record.id,
|
|
full_message=orjson.dumps(
|
|
[{'type': msg.type, 'content': msg.content} for msg in
|
|
self.sql_message]).decode())
|
|
full_thinking_text = ''
|
|
full_sql_text = ''
|
|
token_usage = {}
|
|
res = self.llm.stream(self.sql_message)
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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(
|
|
[{'type': msg.type, 'content': msg.content} for msg in
|
|
self.sql_message]).decode())
|
|
|
|
def generate_filter(self, sql: str, tables: List):
|
|
table_list = self.session.query(CoreTable).filter(
|
|
and_(CoreTable.ds_id == self.ds.id, CoreTable.table_name.in_(tables))
|
|
).all()
|
|
|
|
filters = []
|
|
for table in table_list:
|
|
row_permissions = self.session.query(DsPermission).filter(
|
|
and_(DsPermission.table_id == table.id, DsPermission.type == 'row')).all()
|
|
res: List[PermissionDTO] = []
|
|
if row_permissions is not None:
|
|
for permission in row_permissions:
|
|
# check permission and user in same rules
|
|
obj = self.session.query(DsRules).filter(
|
|
and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)),
|
|
or_(DsRules.user_list.op('@>')(cast([f'{self.current_user.id}'], JSONB)),
|
|
DsRules.user_list.op('@>')(cast([self.current_user.id], JSONB))))
|
|
).first()
|
|
if obj is not None:
|
|
res.append(transRecord2DTO(self.session, permission))
|
|
wheres = transFilterTree(self.session, res, self.ds)
|
|
filters.append({"table": table.table_name, "filter": wheres})
|
|
|
|
filter = json.dumps(filters, ensure_ascii=False)
|
|
|
|
self.chat_question.sql = sql
|
|
self.chat_question.filter = filter
|
|
msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
|
msg.append(SystemMessage(content=self.chat_question.filter_sys_question()))
|
|
msg.append(HumanMessage(content=self.chat_question.filter_user_question()))
|
|
|
|
history_msg = []
|
|
# if self.record.full_analysis_message and self.record.full_analysis_message.strip() != '':
|
|
# history_msg = orjson.loads(self.record.full_analysis_message)
|
|
|
|
# self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='',
|
|
# full_message=orjson.dumps(history_msg +
|
|
# [{'type': msg.type,
|
|
# 'content': msg.content} for msg
|
|
# in
|
|
# msg]).decode())
|
|
full_thinking_text = ''
|
|
full_filter_text = ''
|
|
res = self.llm.stream(msg)
|
|
token_usage = {}
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
full_filter_text += chunk.content
|
|
# yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
|
|
get_token_usage(chunk, token_usage)
|
|
|
|
msg.append(AIMessage(full_filter_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 +
|
|
# [{'type': msg.type,
|
|
# 'content': msg.content} for msg
|
|
# in
|
|
# analysis_msg]).decode())
|
|
SQLBotLogUtil.info(full_filter_text)
|
|
return full_filter_text
|
|
|
|
def generate_chart(self):
|
|
# append current question
|
|
self.chart_message.append(HumanMessage(self.chat_question.chart_user_question()))
|
|
self.record = save_full_chart_message(session=self.session, record_id=self.record.id,
|
|
full_message=orjson.dumps(
|
|
[{'type': msg.type, 'content': msg.content} for msg in
|
|
self.chart_message]).decode())
|
|
full_thinking_text = ''
|
|
full_chart_text = ''
|
|
token_usage = {}
|
|
res = self.llm.stream(self.chart_message)
|
|
for chunk in res:
|
|
SQLBotLogUtil.info(chunk)
|
|
reasoning_content_chunk = ''
|
|
if 'reasoning_content' in chunk.additional_kwargs:
|
|
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
|
# else:
|
|
# reasoning_content_chunk = chunk.get('reasoning_content')
|
|
if reasoning_content_chunk is None:
|
|
reasoning_content_chunk = ''
|
|
full_thinking_text += reasoning_content_chunk
|
|
|
|
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(
|
|
[{'type': msg.type, 'content': msg.content} for msg in
|
|
self.chart_message]).decode())
|
|
|
|
def check_save_sql(self, res: str) -> str:
|
|
|
|
json_str = extract_nested_json(res)
|
|
data = orjson.loads(json_str)
|
|
|
|
sql = ''
|
|
message = ''
|
|
error = False
|
|
|
|
if data['success']:
|
|
sql = data['sql']
|
|
else:
|
|
message = data['message']
|
|
error = True
|
|
|
|
if error:
|
|
raise Exception(message)
|
|
if sql.strip() == '':
|
|
raise Exception("SQL query is empty")
|
|
|
|
save_sql(session=self.session, sql=sql, record_id=self.record.id)
|
|
|
|
self.chat_question.sql = sql
|
|
|
|
return sql
|
|
|
|
def check_save_chart(self, res: str) -> Dict[str, Any]:
|
|
|
|
json_str = extract_nested_json(res)
|
|
data = orjson.loads(json_str)
|
|
|
|
chart: Dict[str, Any] = {}
|
|
message = ''
|
|
error = False
|
|
|
|
if data['type'] and data['type'] != 'error':
|
|
# todo type check
|
|
chart = data
|
|
elif data['type'] == 'error':
|
|
message = data['reason']
|
|
error = True
|
|
else:
|
|
message = 'Chart is empty'
|
|
error = True
|
|
|
|
if error:
|
|
raise Exception(message)
|
|
|
|
save_chart(session=self.session, chart=orjson.dumps(chart).decode(), record_id=self.record.id)
|
|
|
|
return chart
|
|
|
|
def check_save_predict_data(self, res: str) -> bool:
|
|
|
|
json_str = extract_nested_json(res)
|
|
|
|
if not json_str:
|
|
json_str = ''
|
|
|
|
save_predict_data(session=self.session, record_id=self.record.id, data=json_str)
|
|
|
|
if json_str == '':
|
|
return False
|
|
|
|
return True
|
|
|
|
def save_error(self, message: str):
|
|
return save_error_message(session=self.session, record_id=self.record.id, message=message)
|
|
|
|
def save_sql_data(self, data_obj: Dict[str, Any]):
|
|
return save_sql_exec_data(session=self.session, record_id=self.record.id,
|
|
data=orjson.dumps(data_obj).decode())
|
|
|
|
def finish(self):
|
|
return finish_record(session=self.session, record_id=self.record.id)
|
|
|
|
def execute_sql(self, sql: str):
|
|
"""Execute SQL query
|
|
|
|
Args:
|
|
ds: Data source instance
|
|
sql: SQL query statement
|
|
|
|
Returns:
|
|
Query results
|
|
"""
|
|
SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}")
|
|
return exec_sql(self.ds, sql)
|
|
|
|
def pop_chunk(self):
|
|
try:
|
|
chunk = self.chunk_list.pop(0)
|
|
return chunk
|
|
except IndexError as e:
|
|
return None
|
|
|
|
def await_result(self):
|
|
while self.is_running():
|
|
while True:
|
|
chunk = self.pop_chunk()
|
|
if chunk is not None:
|
|
yield chunk
|
|
else:
|
|
break
|
|
while True:
|
|
chunk = self.pop_chunk()
|
|
if chunk is None:
|
|
break
|
|
yield chunk
|
|
|
|
def run_task_async(self, in_chat: bool = True):
|
|
self.future = executor.submit(self.run_task_cache, in_chat)
|
|
|
|
def run_task_cache(self, in_chat: bool = True):
|
|
for chunk in self.run_task(in_chat):
|
|
self.chunk_list.append(chunk)
|
|
|
|
def run_task(self, in_chat: bool = True):
|
|
try:
|
|
# return id
|
|
if in_chat:
|
|
yield orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
|
|
|
|
# return title
|
|
if self.change_title:
|
|
if self.chat_question.question or self.chat_question.question.strip() != '':
|
|
brief = rename_chat(session=self.session,
|
|
rename_object=RenameChat(id=self.get_record().chat_id,
|
|
brief=self.chat_question.question.strip()[:20]))
|
|
if in_chat:
|
|
yield orjson.dumps({'type': 'brief', 'brief': brief}).decode() + '\n\n'
|
|
|
|
# select datasource if datasource is none
|
|
if not self.ds:
|
|
ds_res = self.select_datasource()
|
|
|
|
for chunk in ds_res:
|
|
SQLBotLogUtil.info(chunk)
|
|
if in_chat:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'datasource-result'}).decode() + '\n\n'
|
|
if in_chat:
|
|
yield orjson.dumps({'id': self.ds.id, 'datasource_name': self.ds.name,
|
|
'engine_type': self.ds.type_name or self.ds.type,
|
|
'type': 'datasource'}).decode() + '\n\n'
|
|
|
|
self.chat_question.db_schema = self.out_ds_instance.get_db_schema(
|
|
self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session,
|
|
current_user=self.current_user,
|
|
ds=self.ds)
|
|
|
|
# generate sql
|
|
sql_res = self.generate_sql()
|
|
full_sql_text = ''
|
|
for chunk in sql_res:
|
|
full_sql_text += chunk.get('content')
|
|
if in_chat:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'sql-result'}).decode() + '\n\n'
|
|
if in_chat:
|
|
yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
|
|
|
|
# filter sql
|
|
SQLBotLogUtil.info(full_sql_text)
|
|
|
|
# todo row permission
|
|
sql_json_str = extract_nested_json(full_sql_text)
|
|
data = orjson.loads(sql_json_str)
|
|
|
|
sql = ''
|
|
message = ''
|
|
error = False
|
|
if data['success']:
|
|
sql = data['sql']
|
|
else:
|
|
message = data['message']
|
|
error = True
|
|
if error:
|
|
raise Exception(message)
|
|
if sql.strip() == '':
|
|
raise Exception("SQL query is empty")
|
|
|
|
sql_result = self.generate_filter(data.get('sql'), data.get('tables')) # maybe no sql and tables
|
|
SQLBotLogUtil.info(sql_result)
|
|
sql = self.check_save_sql(res=sql_result)
|
|
# sql = llm_service.check_save_sql(res=full_sql_text)
|
|
|
|
SQLBotLogUtil.info(sql)
|
|
format_sql = sqlparse.format(sql, reindent=True)
|
|
if in_chat:
|
|
yield orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n'
|
|
else:
|
|
yield f'```sql\n{format_sql}\n```\n\n'
|
|
|
|
# execute sql
|
|
result = self.execute_sql(sql=sql)
|
|
self.save_sql_data(data_obj=result)
|
|
if in_chat:
|
|
yield orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n'
|
|
|
|
# generate chart
|
|
chart_res = self.generate_chart()
|
|
full_chart_text = ''
|
|
for chunk in chart_res:
|
|
full_chart_text += chunk.get('content')
|
|
if in_chat:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'chart-result'}).decode() + '\n\n'
|
|
if in_chat:
|
|
yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
|
|
|
|
# filter chart
|
|
SQLBotLogUtil.info(full_chart_text)
|
|
chart = self.check_save_chart(res=full_chart_text)
|
|
SQLBotLogUtil.info(chart)
|
|
if in_chat:
|
|
yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
|
|
else:
|
|
data = []
|
|
_fields = {}
|
|
if chart.get('columns'):
|
|
for _column in chart.get('columns'):
|
|
if _column:
|
|
_fields[_column.get('value')] = _column.get('name')
|
|
if chart.get('axis'):
|
|
if chart.get('axis').get('x'):
|
|
_fields[chart.get('axis').get('x').get('value')] = chart.get('axis').get('x').get('name')
|
|
if chart.get('axis').get('y'):
|
|
_fields[chart.get('axis').get('y').get('value')] = chart.get('axis').get('y').get('name')
|
|
if chart.get('axis').get('series'):
|
|
_fields[chart.get('axis').get('series').get('value')] = chart.get('axis').get('series').get(
|
|
'name')
|
|
_fields_list = []
|
|
_fields_skip = False
|
|
for _data in result.get('data'):
|
|
_row = []
|
|
for field in result.get('fields'):
|
|
_row.append(_data.get(field))
|
|
if not _fields_skip:
|
|
_fields_list.append(field if not _fields.get(field) else _fields.get(field))
|
|
data.append(_row)
|
|
_fields_skip = True
|
|
df = pd.DataFrame(np.array(data), columns=_fields_list)
|
|
markdown_table = df.to_markdown(index=False)
|
|
yield markdown_table + '\n\n'
|
|
|
|
record = self.finish()
|
|
if in_chat:
|
|
yield orjson.dumps({'type': 'finish'}).decode() + '\n\n'
|
|
else:
|
|
# todo generate picture
|
|
if chart['type'] != 'table':
|
|
yield '### generated chart picture\n\n'
|
|
image_url = request_picture(self.record.chat_id, self.record.id, chart, result)
|
|
SQLBotLogUtil.info(image_url)
|
|
yield f'![{chart["type"]}]({image_url})'
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.save_error(message=str(e))
|
|
if in_chat:
|
|
yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
|
|
else:
|
|
yield f'> ❌ **ERROR**\n\n> \n\n> {str(e)}。'
|
|
|
|
def run_recommend_questions_task_async(self):
|
|
self.future = executor.submit(self.run_recommend_questions_task_cache)
|
|
|
|
def run_recommend_questions_task_cache(self):
|
|
for chunk in self.run_recommend_questions_task():
|
|
self.chunk_list.append(chunk)
|
|
|
|
def run_recommend_questions_task(self):
|
|
res = self.generate_recommend_questions_task()
|
|
|
|
for chunk in res:
|
|
if chunk.get('recommended_question'):
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('recommended_question'), 'type': 'recommended_question'}).decode() + '\n\n'
|
|
else:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'recommended_question_result'}).decode() + '\n\n'
|
|
|
|
def run_analysis_or_predict_task_async(self, action_type: str, base_record: ChatRecord):
|
|
self.set_record(save_analysis_predict_record(self.session, base_record, action_type))
|
|
self.future = executor.submit(self.run_analysis_or_predict_task_cache, action_type)
|
|
|
|
def run_analysis_or_predict_task_cache(self, action_type: str):
|
|
for chunk in self.run_analysis_or_predict_task(action_type):
|
|
self.chunk_list.append(chunk)
|
|
|
|
def run_analysis_or_predict_task(self, action_type: str):
|
|
try:
|
|
|
|
yield orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
|
|
|
|
if action_type == 'analysis':
|
|
# generate analysis
|
|
analysis_res = self.generate_analysis()
|
|
for chunk in analysis_res:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'analysis-result'}).decode() + '\n\n'
|
|
yield orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
|
|
|
|
yield orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
|
|
|
|
elif action_type == 'predict':
|
|
# generate predict
|
|
analysis_res = self.generate_predict()
|
|
full_text = ''
|
|
for chunk in analysis_res:
|
|
yield orjson.dumps(
|
|
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
|
'type': 'predict-result'}).decode() + '\n\n'
|
|
full_text += chunk.get('content')
|
|
yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
|
|
|
|
_data = self.check_save_predict_data(res=full_text)
|
|
if _data:
|
|
yield orjson.dumps({'type': 'predict-success'}).decode() + '\n\n'
|
|
else:
|
|
yield orjson.dumps({'type': 'predict-failed'}).decode() + '\n\n'
|
|
|
|
yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
|
|
|
|
self.finish()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.save_error(message=str(e))
|
|
yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
|
|
finally:
|
|
# end
|
|
pass
|
|
|
|
|
|
def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
|
|
"""Execute SQL query using SQLDatabase
|
|
|
|
Args:
|
|
db: SQLDatabase instance
|
|
sql: SQL query statement
|
|
|
|
Returns:
|
|
str: Query results formatted as string
|
|
"""
|
|
try:
|
|
# Execute query
|
|
result = db.run(sql)
|
|
|
|
if not result:
|
|
return "Query executed successfully but returned no results."
|
|
|
|
# Format results
|
|
return str(result)
|
|
|
|
except Exception as e:
|
|
error_msg = f"SQL execution failed: {str(e)}"
|
|
SQLBotLogUtil.exception(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
|
|
def request_picture(chat_id: int, record_id: int, chart: dict, data: dict):
|
|
file_name = f'c_{chat_id}_r_{record_id}'
|
|
|
|
columns = chart.get('columns') if chart.get('columns') else []
|
|
x = None
|
|
y = None
|
|
series = None
|
|
if chart.get('axis'):
|
|
x = chart.get('axis').get('x')
|
|
y = chart.get('axis').get('y')
|
|
series = chart.get('axis').get('series')
|
|
|
|
axis = []
|
|
for v in columns:
|
|
axis.append({'name': v.get('name'), 'value': v.get('value')})
|
|
if x:
|
|
axis.append({'name': x.get('name'), 'value': x.get('value'), 'type': 'x'})
|
|
if y:
|
|
axis.append({'name': y.get('name'), 'value': y.get('value'), 'type': 'y'})
|
|
if series:
|
|
axis.append({'name': series.get('name'), 'value': series.get('value'), 'type': 'series'})
|
|
|
|
request_obj = {
|
|
"path": (settings.MCP_IMAGE_PATH if settings.MCP_IMAGE_PATH[-1] == '/' else (
|
|
settings.MCP_IMAGE_PATH + '/')) + file_name,
|
|
"type": chart['type'],
|
|
"data": orjson.dumps(data.get('data') if data.get('data') else []).decode(),
|
|
"axis": orjson.dumps(axis).decode(),
|
|
}
|
|
|
|
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
|