mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: mcp tools
This commit is contained in:
@@ -8,7 +8,7 @@ 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
|
||||
from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion, ChatToken
|
||||
from apps.chat.task.llm import LLMService
|
||||
from apps.datasource.crud.datasource import get_table_schema
|
||||
from apps.datasource.models.datasource import CoreDatasource
|
||||
@@ -58,15 +58,15 @@ async def delete(session: SessionDep, chart_id: int):
|
||||
|
||||
|
||||
@router.post("/mcp_start", operation_id="mcp_start")
|
||||
async def mcp_start(session: SessionDep, token: str, create_chat_obj: CreateChat):
|
||||
user = await get_current_user(session, token)
|
||||
return await start_chat(session, user, create_chat_obj)
|
||||
async def mcp_start(session: SessionDep, chat: ChatToken):
|
||||
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, token: str, request_question: ChatQuestion):
|
||||
user = await get_current_user(session, token)
|
||||
return await stream_sql(session, user, request_question)
|
||||
async def mcp_question(session: SessionDep, chat: ChatToken):
|
||||
user = await get_current_user(session, chat.token)
|
||||
return await stream_sql(session, user, chat)
|
||||
|
||||
|
||||
@router.post("/start")
|
||||
|
||||
@@ -138,5 +138,9 @@ class AiModelQuestion(BaseModel):
|
||||
|
||||
|
||||
class ChatQuestion(AiModelQuestion):
|
||||
question: str
|
||||
chat_id: int
|
||||
question: str = ''
|
||||
chat_id: int = 0
|
||||
|
||||
|
||||
class ChatToken(ChatQuestion):
|
||||
token:str = ''
|
||||
Reference in New Issue
Block a user