|
1 | 1 | # Author: Junjun |
2 | 2 | # Date: 2025/7/1 |
3 | | - |
| 3 | +import json |
4 | 4 | from datetime import timedelta |
5 | 5 |
|
6 | 6 | import jwt |
|
12 | 12 | from sqlmodel import select |
13 | 13 |
|
14 | 14 | from apps.chat.api.chat import create_chat |
15 | | -from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion |
| 15 | +from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion, McpAssistant, ChatQuestion |
16 | 16 | from apps.chat.task.llm import LLMService |
17 | 17 | from apps.system.crud.user import authenticate |
18 | 18 | from apps.system.crud.user import get_db_user |
19 | 19 | from apps.system.models.system_model import UserWsModel |
20 | 20 | from apps.system.models.user import UserModel |
21 | | -from apps.system.schemas.system_schema import BaseUserDTO |
| 21 | +from apps.system.schemas.system_schema import BaseUserDTO, AssistantHeader |
22 | 22 | from apps.system.schemas.system_schema import UserInfoDTO |
23 | 23 | from common.core import security |
24 | 24 | from common.core.config import settings |
@@ -111,3 +111,29 @@ async def mcp_question(session: SessionDep, chat: McpQuestion): |
111 | 111 | llm_service.init_record() |
112 | 112 |
|
113 | 113 | return StreamingResponse(llm_service.run_task(False), media_type="text/event-stream") |
| 114 | + |
| 115 | + |
| 116 | +@router.post("/mcp_assistant", operation_id="mcp_assistant") |
| 117 | +async def mcp_assistant(session: SessionDep, chat: McpAssistant): |
| 118 | + session_user = BaseUserDTO(**{ |
| 119 | + "id": -1, "account": 'sqlbot-mcp-assistant', "oid": 1, "assistant_id": -1, "password": '', "language": "zh-CN" |
| 120 | + }) |
| 121 | + # session_user: UserModel = get_db_user(session=session, user_id=1) |
| 122 | + # session_user.oid = 1 |
| 123 | + c = create_chat(session, session_user, CreateChat(origin=1), False) |
| 124 | + |
| 125 | + # build assistant param |
| 126 | + configuration = {"endpoint": chat.url} |
| 127 | + # authorization = [{"key": "x-de-token", |
| 128 | + # "value": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1aWQiOjEsIm9pZCI6MSwiZXhwIjoxNzU4NTEyMDA2fQ.3NR-pgnADLdXZtI3dXX5-LuxfGYRvYD9kkr2de7KRP0", |
| 129 | + # "target": "header"}] |
| 130 | + mcp_assistant_header = AssistantHeader(id=1, name='mcp_assist', domain='', type=1, |
| 131 | + configuration=json.dumps(configuration), |
| 132 | + certificate=chat.authorization) |
| 133 | + |
| 134 | + # assistant question |
| 135 | + mcp_chat = ChatQuestion(chat_id=c.id, question=chat.question) |
| 136 | + # ask |
| 137 | + llm_service = await LLMService.create(session_user, mcp_chat, mcp_assistant_header) |
| 138 | + llm_service.init_record() |
| 139 | + return llm_service.run_task(False) |
0 commit comments