From 9fa388c0ea8c565f9d42041aacb6de002d4f5ea1 Mon Sep 17 00:00:00 2001 From: junjun Date: Thu, 31 Jul 2025 10:01:40 +0800 Subject: [PATCH] feat: mcp --- backend/apps/chat/models/chat_model.py | 24 ++++++++++++++++-------- backend/apps/mcp/mcp.py | 7 ++++--- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index c1f431d7..c909d6a4 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -111,7 +111,8 @@ class AiModelQuestion(BaseModel): sub_query: Optional[list[dict]] = None def sql_sys_question(self): - return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, lang=self.lang) + return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, + lang=self.lang) def sql_user_question(self): return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question, @@ -153,22 +154,29 @@ class AiModelQuestion(BaseModel): def filter_user_question(self): return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter) - - def dynamic_sys_question(self): + + def dynamic_sys_question(self): return get_dynamic_template()['system'].format(lang=self.lang, engine=self.engine) - - def dynamic_user_question(self): + + def dynamic_user_question(self): return get_dynamic_template()['user'].format(sql=self.sql, sub_query=self.sub_query) + class ChatQuestion(AiModelQuestion): - question: str = Body(description='用户提问') - chat_id: int = Body(description='会话ID') + question: str + chat_id: int class ChatMcp(ChatQuestion): - token: str = Body(description='token') + token: str class ChatStart(BaseModel): username: str = Body(description='用户名') password: str = Body(description='密码') + + +class McpQuestion(BaseModel): + question: str = Body(description='用户提问') + chat_id: int = Body(description='会话ID') + token: str = Body(description='token') diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 70373038..19afffe7 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from sqlmodel import select from apps.chat.api.chat import create_chat -from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart +from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion from apps.chat.task.llm import LLMService from apps.system.crud.user import authenticate from apps.system.crud.user import get_db_user @@ -76,7 +76,7 @@ async def mcp_start(session: SessionDep, chat: ChatStart): @router.post("/mcp_question", operation_id="mcp_question") -async def mcp_question(session: SessionDep, chat: ChatMcp): +async def mcp_question(session: SessionDep, chat: McpQuestion): try: payload = jwt.decode( chat.token, settings.SECRET_KEY, algorithms=[security.ALGORITHM] @@ -105,8 +105,9 @@ async def mcp_question(session: SessionDep, chat: ChatMcp): if session_user.status != 1: raise HTTPException(status_code=400, detail="Inactive user") + mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question) # ask - llm_service = LLMService(session_user, chat) + llm_service = LLMService(session_user, mcp_chat) llm_service.init_record() return StreamingResponse(llm_service.run_task(False), media_type="text/event-stream")