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

perf: Simple assistant

Browse Source
This commit is contained in:
fit2cloud-chenyw
2025-07-21 19:20:24 +08:00
parent 7137a0071e
commit fbcb147b8b
3 changed files with 33 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/api/chat.py
+7 -7
View File
@@ -7,7 +7,7 @@ from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat,
delete_chat, get_chat_chart_data, get_chat_predict_data
from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion
from apps.chat.task.llm import LLMService, run_task, run_analysis_or_predict_task, run_recommend_questions_task
from common.core.deps import SessionDep, CurrentUser
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
router = APIRouter(tags=["Data Q&A"], prefix="/chat")
@@ -94,7 +94,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser):
@router.post("/recommend_questions/{chat_record_id}")
async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int, current_assistant: CurrentAssistant):
try:
record = session.query(ChatRecord).get(chat_record_id)
if not record:
@@ -104,7 +104,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch
)
request_question = ChatQuestion(chat_id=record.chat_id, question=record.question if record.question else '')
llm_service = LLMService(session, current_user, request_question)
llm_service = LLMService(session, current_user, request_question, current_assistant)
llm_service.set_record(record)
except Exception as e:
traceback.print_exc()
@@ -117,7 +117,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch
@router.post("/question")
async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion, current_assistant: CurrentAssistant):
"""Stream SQL analysis results
Args:
@@ -130,7 +130,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
"""
try:
llm_service = LLMService(session, current_user, request_question)
llm_service = LLMService(session, current_user, request_question, current_assistant)
llm_service.init_record()
except Exception as e:
traceback.print_exc()
@@ -143,7 +143,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
@router.post("/record/{chat_record_id}/{action_type}")
async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str):
async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str, current_assistant: CurrentAssistant):
if action_type != 'analysis' and action_type != 'predict':
raise HTTPException(
status_code=404,
@@ -166,7 +166,7 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
request_question = ChatQuestion(chat_id=record.chat_id, question='')
try:
llm_service = LLMService(session, current_user, request_question)
llm_service = LLMService(session, current_user, request_question, current_assistant)
except Exception as e:
traceback.print_exc()
raise HTTPException(
backend/apps/chat/task/llm.py
+13 -3
View File
@@ -353,11 +353,21 @@ class LLMService:
if self.current_assistant:
_ds_list = get_assistant_ds(llm_service=self)
else:
_ds_list = self.session.exec(select(CoreDatasource).options(
load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all()
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({'id': _ds[0].id, 'name': _ds[0].name, 'description': _ds[0].description})
_ds_list_dict.append(_ds)
datasource_msg.append(
HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
backend/apps/system/crud/assistant.py
+13 -3
View File
@@ -33,12 +33,22 @@ def get_assistant_ds(llm_service) -> list[dict]:
config: dict[any] = json.loads(configuration)
oid: str = config['oid']
stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
private_list:list[int] = config['private_list']
private_list:list[int] = config.get('private_list') or None
if private_list:
stmt.where(~CoreDatasource.id.in_(private_list))
db_ds_list = session.exec(stmt).all()
db_ds_list = session.exec(stmt)
result_list = [
{
"id": ds.id,
"name": ds.name,
"description": ds.description
}
for ds in db_ds_list
]
# filter private ds if offline
return db_ds_list
return result_list
out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant, llm_service.assistant_certificate)
llm_service.out_ds_instance = out_ds_instance
dslist = out_ds_instance.get_simple_ds_list()