diff --git a/.env.example b/.env.example index d211818c..9b11faa8 100644 --- a/.env.example +++ b/.env.example @@ -26,3 +26,5 @@ SENTRY_DSN= # Configure these with your own Docker registry images DOCKER_IMAGE_BACKEND=backend DOCKER_IMAGE_FRONTEND=frontend + +MCP_IMAGE_PATH=/opt/sqlbot/images diff --git a/backend/apps/api.py b/backend/apps/api.py index c20d842c..9702c247 100644 --- a/backend/apps/api.py +++ b/backend/apps/api.py @@ -1,10 +1,11 @@ from fastapi import APIRouter -from apps.system.api import login, user, aimodel -from apps.settings.api import terminology -from apps.datasource.api import datasource from apps.chat.api import chat from apps.dashboard.api import dashboard_api +from apps.datasource.api import datasource +from apps.settings.api import terminology +from apps.system.api import login, user, aimodel +from apps.mcp import mcp api_router = APIRouter() api_router.include_router(login.router) @@ -12,5 +13,9 @@ api_router.include_router(user.router) api_router.include_router(aimodel.router) api_router.include_router(terminology.router) api_router.include_router(datasource.router) +# api_router.include_router(row_permission.router) +# api_router.include_router(column_permission.router) api_router.include_router(chat.router) api_router.include_router(dashboard_api.router) +api_router.include_router(mcp.router) + diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 8229b56d..32f338a7 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -1,20 +1,14 @@ import traceback -from typing import List import orjson from fastapi import APIRouter, HTTPException from fastapi.responses import StreamingResponse -from sqlmodel import select from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \ - delete_chat, list_records -from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion, ChatMcp -from apps.chat.task.llm import LLMService -from apps.datasource.crud.datasource import get_table_schema -from apps.datasource.models.datasource import CoreDatasource -from apps.system.crud.user import get_user_info -from apps.system.models.system_model import AiModelDetail -from common.core.deps import SessionDep, CurrentUser, get_current_user + delete_chat +from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion +from apps.chat.task.llm import LLMService, run_task +from common.core.deps import SessionDep, CurrentUser router = APIRouter(tags=["Data Q&A"], prefix="/chat") @@ -57,18 +51,6 @@ async def delete(session: SessionDep, chart_id: int): ) -@router.post("/mcp_start", operation_id="mcp_start") -async def mcp_start(session: SessionDep, chat: ChatMcp): - user = await get_current_user(session, chat.token) - return create_chat(session, user, CreateChat(), False) - - -@router.post("/mcp_question", operation_id="mcp_question") -async def mcp_question(session: SessionDep, chat: ChatMcp): - user = await get_current_user(session, chat.token) - return await stream_sql(session, user, chat) - - @router.post("/start") async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat): try: @@ -80,7 +62,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat ) -@router.post("/question", operation_id="question") +@router.post("/question") async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion): """Stream SQL analysis results @@ -93,107 +75,17 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que Streaming response with analysis results """ - chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first() - if not chat: - raise HTTPException( - status_code=400, - detail=f"Chat with id {request_question.chat_id} not found" - ) - ds: CoreDatasource | None = None - if chat.datasource: - # Get available datasource - ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first() - if not ds: - raise HTTPException( - status_code=500, - detail="No available datasource configuration found" - ) - - request_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL' - - # Get available AI model - aimodel = session.exec(select(AiModelDetail).where( - AiModelDetail.status == True, - AiModelDetail.api_key.is_not(None) - )).first() - if not aimodel: + try: + llm_service = LLMService(session, current_user, request_question) + llm_service.init_record() + except Exception as e: + traceback.print_exc() raise HTTPException( status_code=500, - detail="No available AI model configuration found" + detail=str(e) ) - history_records: List[ChatRecord] = list(filter(lambda r: True if r.first_chat != True else False, - list_records(session=session, current_user=current_user, - chart_id=request_question.chat_id))) - # get schema - if ds: - request_question.db_schema = get_table_schema(session=session, ds=ds) - - db_user = get_user_info(session=session, user_id=current_user.id) - request_question.lang = db_user.language - - llm_service = LLMService(request_question, aimodel, history_records, - CoreDatasource(**ds.model_dump()) if ds else None) - - llm_service.init_record(session=session, current_user=current_user) - - def run_task(): - try: - # return id - yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n' - - # select datasource if datasource is none - if not ds: - ds_res = llm_service.select_datasource(session=session) - for chunk in ds_res: - yield orjson.dumps({'content': chunk, 'type': 'datasource-result'}).decode() + '\n\n' - yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name, - 'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n' - - llm_service.chat_question.db_schema = get_table_schema(session=session, ds=llm_service.ds) - - # generate sql - sql_res = llm_service.generate_sql(session=session) - full_sql_text = '' - for chunk in sql_res: - full_sql_text += chunk - yield orjson.dumps({'content': chunk, 'type': 'sql-result'}).decode() + '\n\n' - yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n' - - # filter sql - print(full_sql_text) - sql = llm_service.check_save_sql(session=session, res=full_sql_text) - print(sql) - yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n' - - # execute sql - result = llm_service.execute_sql(sql=sql) - llm_service.save_sql_data(session=session, data_obj=result) - yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n' - - # generate chart - chart_res = llm_service.generate_chart(session=session) - full_chart_text = '' - for chunk in chart_res: - full_chart_text += chunk - yield orjson.dumps({'content': chunk, 'type': 'chart-result'}).decode() + '\n\n' - yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n' - - # filter chart - print(full_chart_text) - chart = llm_service.check_save_chart(session=session, res=full_chart_text) - print(chart) - yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n' - - llm_service.finish(session=session) - yield orjson.dumps({'type': 'finish'}).decode() + '\n\n' - - except Exception as e: - traceback.print_exc() - llm_service.save_error(session=session, message=str(e)) - yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n' - - return StreamingResponse(run_task(), media_type="text/event-stream") + return StreamingResponse(run_task(llm_service, session), media_type="text/event-stream") @router.post("/record/{chart_record_id}/{action_type}") @@ -217,35 +109,9 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch detail=f"Chat record with id {chart_record_id} has not generated chart, do not support to analyze it" ) - chat = session.query(Chat).filter(Chat.id == record.chat_id).first() - if not chat: - raise HTTPException( - status_code=400, - detail=f"Chat with id {record.chart_id} not found" - ) + request_question = ChatQuestion(chat_id=record.chat_id, question='') - if chat.create_by != current_user.id: - raise HTTPException( - status_code=401, - detail=f"You cannot use the chat with id {record.chart_id}" - ) - - # Get available AI model - aimodel = session.exec(select(AiModelDetail).where( - AiModelDetail.status == True, - AiModelDetail.api_key.is_not(None) - )).first() - if not aimodel: - raise HTTPException( - status_code=500, - detail="No available AI model configuration found" - ) - - request_question = ChatQuestion(chat_id=chat.id, question='') - db_user = get_user_info(session=session, user_id=current_user.id) - request_question.lang = db_user.language - - llm_service = LLMService(request_question, aimodel) + llm_service = LLMService(session, current_user, request_question) llm_service.set_record(record) def run_task(): @@ -261,14 +127,14 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch elif action_type == 'predict': # generate predict - analysis_res = llm_service.generate_predict(session=session) + analysis_res = llm_service.generate_predict() full_text = '' for chunk in analysis_res: yield orjson.dumps({'content': chunk, 'type': 'predict-result'}).decode() + '\n\n' full_text += chunk yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n' - _data = llm_service.check_save_predict_data(session=session, res=full_text) + _data = llm_service.check_save_predict_data(res=full_text) yield orjson.dumps({'type': 'predict', 'content': _data}).decode() + '\n\n' yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n' diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 8cc9d80c..5418beb6 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -2,7 +2,9 @@ import logging import warnings from typing import Any, List, Union, Dict +import numpy as np import orjson +import pandas as pd from langchain_community.utilities import SQLDatabase from langchain_core.language_models import BaseLLM from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, AIMessageChunk @@ -13,10 +15,12 @@ from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_llm_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_select_datasource_message_and_answer, list_records from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat +from apps.datasource.crud.datasource import get_table_schema from apps.datasource.models.datasource import CoreDatasource from apps.db.db import exec_sql +from apps.system.crud.user import get_user_info from apps.system.models.system_model import AiModelDetail from common.core.deps import SessionDep, CurrentUser @@ -34,13 +38,49 @@ class LLMService: 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 - def __init__(self, chat_question: ChatQuestion, aimodel: AiModelDetail, history_records: List[ChatRecord] = [], - ds: CoreDatasource = None): - self.ds = ds + def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question: ChatQuestion): + + self.session = session + self.current_user = current_user + + chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first() + if not chat: + raise Exception(f"Chat with id {chat_question.chat_id} not found") + ds: CoreDatasource | None = None + if chat.datasource: + # Get available datasource + ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first() + if not ds: + raise Exception("No available datasource configuration found") + + chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL' + + # Get available AI model + aimodel = self.session.exec(select(AiModelDetail).where( + AiModelDetail.status == True, + AiModelDetail.api_key.is_not(None) + )).first() + if not aimodel and aimodel[0]: + raise Exception("No available AI model configuration found") + + history_records: List[ChatRecord] = list( + map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False, + list_records(session=self.session, current_user=current_user, + chart_id=chat_question.chat_id)))) + # get schema + if ds: + chat_question.db_schema = get_table_schema(session=self.session, ds=ds) + + db_user = get_user_info(session=self.session, user_id=current_user.id) + chat_question.lang = db_user.language + + self.ds = CoreDatasource(**ds.model_dump()) if ds else None self.chat_question = chat_question - self.chat_question.ai_modal_id = aimodel.id - self.config = get_llm_config(aimodel) + self.chat_question.ai_modal_id = aimodel[0].id + self.config = get_llm_config(aimodel[0]) # Create LLM instance through factory llm_instance = LLMFactory.create_llm(self.config) @@ -112,8 +152,8 @@ class LLMService: _msg = AIMessage(content=last_chart_message['content']) self.chart_message.append(_msg) - def init_record(self, session: SessionDep, current_user: CurrentUser) -> ChatRecord: - self.record = save_question(session=session, current_user=current_user, question=self.chat_question) + 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): @@ -183,7 +223,7 @@ class LLMService: in analysis_msg]).decode()) - def generate_predict(self, session: SessionDep): + def generate_predict(self): fields = self.get_fields_from_chart() self.chat_question.fields = orjson.dumps(fields).decode() @@ -196,7 +236,7 @@ class LLMService: 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=session, record_id=self.record.id, answer='', + 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, @@ -218,7 +258,7 @@ class LLMService: continue predict_msg.append(AIMessage(full_predict_text)) - self.record = save_full_predict_message_and_answer(session=session, record_id=self.record.id, + self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id, answer=full_predict_text, data='', full_message=orjson.dumps(history_msg + [{'type': msg.type, @@ -226,10 +266,10 @@ class LLMService: in predict_msg]).decode()) - def select_datasource(self, session: SessionDep): + def select_datasource(self): datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = [] datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question())) - _ds_list = session.exec(select(CoreDatasource).options( + _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: @@ -241,7 +281,7 @@ class LLMService: 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=session, record_id=self.record.id, + 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, @@ -252,7 +292,6 @@ class LLMService: full_text = '' res = self.llm.stream(datasource_msg) for chunk in res: - print(chunk) if isinstance(chunk, dict): full_text += chunk['content'] yield chunk['content'] @@ -273,7 +312,7 @@ class LLMService: if data['id'] and data['id'] != 0: _datasource = data['id'] - _ds = session.query(CoreDatasource).filter(CoreDatasource.id == _datasource).first() + _ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == _datasource).first() if not _ds: _datasource = None raise Exception(f"Datasource configuration with id {_datasource} not found") @@ -281,14 +320,14 @@ class LLMService: self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL' _engine_type = self.chat_question.engine # save chat - _chat = session.query(Chat).filter(Chat.id == self.record.chat_id).first() + _chat = self.session.query(Chat).filter(Chat.id == self.record.chat_id).first() _chat.datasource = _datasource _chat.engine_type = _ds.type_name - session.add(_chat) - session.flush() - session.refresh(_chat) - session.commit() + self.session.add(_chat) + self.session.flush() + self.session.refresh(_chat) + self.session.commit() elif data['fail']: raise Exception(data['fail']) @@ -298,7 +337,7 @@ class LLMService: except Exception as e: _error = e - self.record = save_full_select_datasource_message_and_answer(session=session, record_id=self.record.id, + self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id, answer=full_text, datasource=_datasource, engine_type=_engine_type, full_message=orjson.dumps(history_msg + @@ -312,10 +351,10 @@ class LLMService: if _error: raise _error - def generate_sql(self, session: SessionDep): + 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=session, record_id=self.record.id, + 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()) @@ -333,16 +372,16 @@ class LLMService: continue self.sql_message.append(AIMessage(full_sql_text)) - self.record = save_full_sql_message_and_answer(session=session, record_id=self.record.id, + self.record = save_full_sql_message_and_answer(session=self.session, record_id=self.record.id, answer=full_sql_text, full_message=orjson.dumps( [{'type': msg.type, 'content': msg.content} for msg in self.sql_message]).decode()) - def generate_chart(self, session: SessionDep): + 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=session, record_id=self.record.id, + 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()) @@ -359,13 +398,13 @@ class LLMService: continue self.chart_message.append(AIMessage(full_chart_text)) - self.record = save_full_chart_message_and_answer(session=session, record_id=self.record.id, + self.record = save_full_chart_message_and_answer(session=self.session, record_id=self.record.id, answer=full_chart_text, full_message=orjson.dumps( [{'type': msg.type, 'content': msg.content} for msg in self.chart_message]).decode()) - def check_save_sql(self, session: SessionDep, res: str) -> str: + def check_save_sql(self, res: str) -> str: json_str = extract_nested_json(res) data = orjson.loads(json_str) @@ -385,13 +424,13 @@ class LLMService: if sql.strip() == '': raise Exception("SQL query is empty") - save_sql(session=session, sql=sql, record_id=self.record.id) + save_sql(session=self.session, sql=sql, record_id=self.record.id) self.chat_question.sql = sql return sql - def check_save_chart(self, session: SessionDep, res: str) -> Dict[str, Any]: + def check_save_chart(self, res: str) -> Dict[str, Any]: json_str = extract_nested_json(res) data = orjson.loads(json_str) @@ -413,30 +452,30 @@ class LLMService: if error: raise Exception(message) - save_chart(session=session, chart=orjson.dumps(chart).decode(), record_id=self.record.id) + save_chart(session=self.session, chart=orjson.dumps(chart).decode(), record_id=self.record.id) return chart - def check_save_predict_data(self, session: SessionDep, res: str) -> Dict[str, Any]: + def check_save_predict_data(self, res: str) -> Dict[str, Any]: json_str = extract_nested_json(res) if not json_str: json_str = '' - save_predict_data(session=session, record_id=self.record.id, data=json_str) + save_predict_data(session=self.session, record_id=self.record.id, data=json_str) return json_str - def save_error(self, session: SessionDep, message: str): - return save_error_message(session=session, record_id=self.record.id, message=message) + 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, session: SessionDep, data_obj: Dict[str, Any]): - return save_sql_exec_data(session=session, record_id=self.record.id, + 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, session: SessionDep): - return finish_record(session=session, record_id=self.record.id) + def finish(self): + return finish_record(session=self.session, record_id=self.record.id) def execute_sql(self, sql: str): """Execute SQL query @@ -503,3 +542,109 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str: error_msg = f"SQL execution failed: {str(e)}" logging.error(error_msg) raise RuntimeError(error_msg) + + +def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True): + try: + # return id + if in_chat: + yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n' + + # select datasource if datasource is none + if not llm_service.ds: + ds_res = llm_service.select_datasource() + + for chunk in ds_res: + print(chunk) + if in_chat: + yield orjson.dumps({'content': chunk, 'type': 'datasource-result'}).decode() + '\n\n' + if in_chat: + yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name, + 'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n' + + llm_service.chat_question.db_schema = get_table_schema(session=session, ds=llm_service.ds) + + # generate sql + sql_res = llm_service.generate_sql() + full_sql_text = '' + for chunk in sql_res: + full_sql_text += chunk + if in_chat: + yield orjson.dumps({'content': chunk, 'type': 'sql-result'}).decode() + '\n\n' + if in_chat: + yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n' + + # filter sql + print(full_sql_text) + sql = llm_service.check_save_sql(res=full_sql_text) + print(sql) + if in_chat: + yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n' + else: + yield f'```sql\n{sql}\n```\n\n' + + # execute sql + result = llm_service.execute_sql(sql=sql) + llm_service.save_sql_data(data_obj=result) + if in_chat: + yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n' + + # generate chart + chart_res = llm_service.generate_chart() + full_chart_text = '' + for chunk in chart_res: + full_chart_text += chunk + if in_chat: + yield orjson.dumps({'content': chunk, 'type': 'chart-result'}).decode() + '\n\n' + if in_chat: + yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n' + + # filter chart + print(full_chart_text) + chart = llm_service.check_save_chart(res=full_chart_text) + print(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 = llm_service.finish() + if in_chat: + yield orjson.dumps({'type': 'finish'}).decode() + '\n\n' + else: + # todo generate picture + if chart['type'] != 'table': + yield '# todo generate chart picture' + + yield f'![{chart["type"]}](https://sqlbot.fit2cloud.cn/images/111.png)' + except Exception as e: + llm_service.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)}。' diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py index 08f2e6f7..b0921673 100644 --- a/backend/apps/datasource/api/datasource.py +++ b/backend/apps/datasource/api/datasource.py @@ -18,7 +18,7 @@ router = APIRouter(tags=["datasource"], prefix="/datasource") path = "/opt/sqlbot/data/excel" -@router.get("/list", operation_id="get_datasource_list") +@router.get("/list") async def datasource_list(session: SessionDep): return get_datasource_list(session=session) diff --git a/backend/apps/mcp/__init__.py b/backend/apps/mcp/__init__.py new file mode 100644 index 00000000..a14a79d9 --- /dev/null +++ b/backend/apps/mcp/__init__.py @@ -0,0 +1,2 @@ +# Author: Junjun +# Date: 2025/7/1 diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py new file mode 100644 index 00000000..a9c266b1 --- /dev/null +++ b/backend/apps/mcp/mcp.py @@ -0,0 +1,63 @@ +# Author: Junjun +# Date: 2025/7/1 + +from datetime import timedelta +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse +from fastapi.security import OAuth2PasswordRequestForm + +from apps.chat.api.chat import create_chat +from apps.chat.models.chat_model import ChatMcp, CreateChat +from apps.chat.task.llm import LLMService, run_task +from apps.datasource.crud.datasource import get_datasource_list +from apps.system.crud.user import authenticate +from apps.system.models.system_model import AiModelDetail +from common.core.config import settings +from common.core.deps import SessionDep, get_current_user +from common.core.schemas import Token +from common.core.security import create_access_token + +router = APIRouter(tags=["mcp"], prefix="/mcp") + + +@router.post("/access_token", operation_id="access_token") +def local_login( + session: SessionDep, + form_data: Annotated[OAuth2PasswordRequestForm, Depends()] +) -> Token: + user = authenticate(session=session, account=form_data.username, password=form_data.password) + if not user: + raise HTTPException(status_code=400, detail="Incorrect account or password") + access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) + user_dict = user.to_dict() + return Token(access_token=create_access_token( + user_dict, expires_delta=access_token_expires + )) + + +@router.get("/ds_list", operation_id="get_datasource_list") +async def datasource_list(session: SessionDep): + return get_datasource_list(session=session) + + +@router.get("/model_list", operation_id="get_model_list") +async def get_model_list(session: SessionDep): + return session.query(AiModelDetail).all() + + +@router.post("/mcp_start", operation_id="mcp_start") +async def mcp_start(session: SessionDep, chat: ChatMcp): + user = await get_current_user(session, chat.token) + return create_chat(session, user, CreateChat(), False) + + +@router.post("/mcp_question", operation_id="mcp_question") +async def mcp_question(session: SessionDep, chat: ChatMcp): + user = await get_current_user(session, chat.token) + + llm_service = LLMService(session, user, chat) + llm_service.init_record() + + return StreamingResponse(run_task(llm_service, session, False), media_type="text/event-stream") diff --git a/backend/apps/system/api/aimodel.py b/backend/apps/system/api/aimodel.py index a3bf6f21..eabc1fd6 100644 --- a/backend/apps/system/api/aimodel.py +++ b/backend/apps/system/api/aimodel.py @@ -88,8 +88,3 @@ async def status(session: SessionDep, dto: model_status): session.add(term) session.commit() return {"message": f"AiModel with IDs {ids} updated successfully."} - - -@router.get("/list", operation_id="get_model_list") -async def get_model_list(session: SessionDep): - return session.query(AiModelDetail).all() diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 078f682e..b54fa12c 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -79,6 +79,7 @@ class Settings(BaseSettings): SMTP_PASSWORD: str | None = None EMAILS_FROM_EMAIL: EmailStr | None = None EMAILS_FROM_NAME: EmailStr | None = None + MCP_IMAGE_PATH: str @model_validator(mode="after") def _set_default_emails_from(self) -> Self: diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py index 33b621fe..abf414bb 100644 --- a/backend/common/utils/whitelist.py +++ b/backend/common/utils/whitelist.py @@ -21,14 +21,11 @@ wlist = [ "*.ttf", "*.eot", "*.otf", - "/mcp", - "/mcp/message", - "/datasource/list", - "/system/aimodel/list", - "/chat/mcp_question", - "/chat/mcp_start", + "/mcp*", "/system/license", - "/system/config/key" + "/system/config/key", + "/images/*", + "/sse", ] class WhitelistChecker: diff --git a/backend/main.py b/backend/main.py index 4131bb3b..85d14420 100644 --- a/backend/main.py +++ b/backend/main.py @@ -12,6 +12,7 @@ from common.core.response_middleware import ResponseMiddleware, exception_handle from alembic.config import Config from alembic import command from fastapi_mcp import FastApiMCP +from fastapi.staticfiles import StaticFiles import sqlbot_xpack def run_migrations(): @@ -39,6 +40,8 @@ app = FastAPI( generate_unique_id_function=custom_generate_unique_id, lifespan=lifespan ) +# mcp server, images path +app.mount("/images", StaticFiles(directory=settings.MCP_IMAGE_PATH), name="images") mcp_app = FastAPI() @@ -48,7 +51,7 @@ mcp = FastApiMCP( description="SQLBot MCP Server", describe_all_responses=True, describe_full_response_schema=True, - include_operations=["get_datasource_list", "get_model_list", "mcp_question", "mcp_start"] + include_operations=["get_datasource_list", "get_model_list", "mcp_question", "mcp_start", "access_token"] ) mcp.mount(mcp_app) diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 5e4981f8..2f894061 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -35,14 +35,15 @@ dependencies = [ "oracledb (>=3.1.1,<4.0.0)", "pyyaml (>=6.0.2,<7.0.0)", "fastapi-mcp (>=0.3.4,<0.4.0)", - "sqlbot-xpack==0.0.3", + "tabulate>=0.9.0", + "sqlbot-xpack==0.0.3.2", ] [[tool.uv.index]] name = "default" url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" default = true -[[tool.uv.index]] +[[tool.uv.index]] name = "testpypi" url = "https://test.pypi.org/simple" explicit = true diff --git a/frontend/src/router/watch.ts b/frontend/src/router/watch.ts index 98224f6b..075140db 100644 --- a/frontend/src/router/watch.ts +++ b/frontend/src/router/watch.ts @@ -38,9 +38,19 @@ const loadXpackStatic = () => { return Promise.resolve() } const url = '/xpack_static/license-generator.umd.js' - return request.loadRemoteScript(url, 'sqlbot_xpack_static', () => { - // eslint-disable-next-line @typescript-eslint/ban-ts-comment - // @ts-ignore - LicenseGenerator?.init(import.meta.env.VITE_API_BASE_URL) + return new Promise((resolve, reject) => { + request + .loadRemoteScript(url, 'sqlbot_xpack_static', () => { + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + LicenseGenerator?.init(import.meta.env.VITE_API_BASE_URL).then(() => { + resolve(true) + }) + }) + .catch((error) => { + console.error('Failed to load xpack_static script:', error) + ElMessage.error('Failed to load license generator script') + reject(error) + }) }) }