feat: mcp

Author: XiaJunjie2020

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 @@ def filter_sys_question(self):
 
     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')

backend/apps/mcp/mcp.py

@@ -12,7 +12,7 @@
 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")