From dfdc5a45a631612ccbb3b2055d1127dba37bd8b7 Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Wed, 14 May 2025 13:11:18 +0800 Subject: [PATCH] perf: Data Q&A --- backend/apps/chat/api/chat.py | 76 ++++++-- backend/apps/chat/schemas/chat_base_schema.py | 1 + backend/apps/chat/schemas/llm.py | 162 +++++++++++++++++- backend/apps/db/db.py | 28 +++ backend/pyproject.toml | 9 +- frontend/src/api/chat.ts | 9 +- frontend/src/utils/request.ts | 19 +- frontend/src/views/chat/index.vue | 151 +++++++++++----- 8 files changed, 389 insertions(+), 66 deletions(-) diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 3a6d723c..08afe6c6 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -1,43 +1,65 @@ from fastapi import APIRouter, HTTPException +from fastapi.responses import StreamingResponse from sqlmodel import select from apps.chat.schemas.chat_base_schema import LLMConfig from apps.chat.schemas.chat_schema import ChatQuestion -from apps.chat.schemas.llm import LLMService +from apps.chat.schemas.llm import AgentService, LLMService +from apps.datasource.models.datasource import CoreDatasource from apps.system.models.system_modle import AiModelDetail from common.core.deps import SessionDep -# from sse_starlette.sse import EventSourceResponse +from sse_starlette.sse import EventSourceResponse +import json +import asyncio + router = APIRouter(tags=["Data Q&A"], prefix="/chat") @router.post("/question") async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion): + """Stream SQL analysis results + + Args: + session: Database session + requestQuestion: User question model + + Returns: + Streaming response with analysis results + """ question = requestQuestion.question - # Use OpenAI model - """ openai_config = LLMConfig( - model_type="openai", - model_name="gpt-4", - api_key="your-api-key", - additional_params={"temperature": 0.7} - ) - openai_service = LLMService(openai_config) """ - - aimodel = session.exec(select(AiModelDetail).where(AiModelDetail.status == True, AiModelDetail.api_key.is_not(None))).first() + # Get available AI model + aimodel = session.exec(select(AiModelDetail).where( + AiModelDetail.status == True, + AiModelDetail.api_key.is_not(None) + )).first() + # Get available datasource + ds = session.exec(select(CoreDatasource).where( + CoreDatasource.status == 'Success' + )).first() + if not aimodel: raise HTTPException( status_code=400, detail="No available AI model configuration found" ) + + if not ds: + raise HTTPException( + status_code=400, + detail="No available datasource configuration found" + ) # Use Tongyi Qianwen tongyi_config = LLMConfig( - model_type="tongyi", + model_type="openai", model_name=aimodel.name, api_key=aimodel.api_key, + api_base_url=aimodel.endpoint, additional_params={"temperature": aimodel.temperature} ) - llm_service = LLMService(tongyi_config) + # llm_service = LLMService(tongyi_config) + llm_service = AgentService(tongyi_config, ds) # Use Custom VLLM model """ vllm_config = LLMConfig( @@ -49,5 +71,27 @@ async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion): } ) vllm_service = LLMService(vllm_config) """ - result = llm_service.generate_sql(question) - return result \ No newline at end of file + """ result = llm_service.generate_sql(question) + return result """ + + async def event_generator(): + try: + async for chunk in llm_service.async_generate(question): + data = json.loads(chunk.replace('data: ', '')) + + if data['type'] in ['final', 'tool_result']: + content = data['content'] + for char in content: + yield f"data: {json.dumps({'type': 'char', 'content': char})}\n\n" + await asyncio.sleep(0.05) + + if 'html' in data: + yield f"data: {json.dumps({'type': 'html', 'content': data['html']})}\n\n" + else: + yield chunk + + except Exception as e: + yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n" + + #return EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"}) + return StreamingResponse(event_generator(), media_type="text/event-stream") \ No newline at end of file diff --git a/backend/apps/chat/schemas/chat_base_schema.py b/backend/apps/chat/schemas/chat_base_schema.py index b5cb540e..5f0fcd72 100644 --- a/backend/apps/chat/schemas/chat_base_schema.py +++ b/backend/apps/chat/schemas/chat_base_schema.py @@ -36,6 +36,7 @@ class OpenAILLM(BaseLLM): return ChatOpenAI( model=self.config.model_name, api_key=self.config.api_key, + base_url=self.config.api_base_url, **self.config.additional_params ) diff --git a/backend/apps/chat/schemas/llm.py b/backend/apps/chat/schemas/llm.py index fff9a2a2..3771f258 100644 --- a/backend/apps/chat/schemas/llm.py +++ b/backend/apps/chat/schemas/llm.py @@ -1,10 +1,17 @@ from langchain_community.utilities import SQLDatabase -# from langchain_community.agent_toolkits import create_sql_agent -from langchain_community.llms import Tongyi +from langgraph.prebuilt import create_react_agent from langchain_core.prompts import ChatPromptTemplate from apps.chat.schemas.chat_base_schema import LLMConfig, LLMFactory +from apps.datasource.models.datasource import CoreDatasource +from apps.db.db import exec_sql, get_uri from common.core.config import settings import warnings +from langchain.tools import Tool +from functools import partial +import logging +from typing import AsyncGenerator +import json +import asyncio warnings.filterwarnings("ignore") @@ -31,3 +38,154 @@ class LLMService: schema = self.db.get_table_info() return chain.invoke({"schema": schema, "question": question}) + +class AgentService: + def __init__(self, config: LLMConfig, ds: CoreDatasource): + # Initialize database connection + self.ds = ds + db_uri = get_uri(ds) + self.db = SQLDatabase.from_uri(db_uri) + # self.db = SQLDatabase.from_uri(str(settings.SQLALCHEMY_DATABASE_URI)) + + # Create LLM instance through factory + llm_instance = LLMFactory.create_llm(config) + self.llm = llm_instance.llm + + # Create a partial function of execute_sql with preset ds parameter + # bound_execute_sql = partial(execute_sql, self.ds) + bound_execute_sql = partial(execute_sql_with_db, self.db) + + # Wrap as Tool object + tools = [ + Tool( + name="execute_sql", + func=bound_execute_sql, + description="""A tool for executing SQL queries. + Input: SQL query statement (string) + Output: Query results + Example: "SELECT * FROM table_name LIMIT 5" + """ + ) + ] + + self.agent_executor = create_react_agent(self.llm, tools) + + system_prompt = """ + You are an intelligent agent capable of data analysis. When users input their data analysis requirements, + you need to first convert the requirements into executable SQL, then execute the SQL through tools to return results, + and finally summarize the SQL query results. When all tasks are completed, you need to generate an HTML format data analysis report. + + You can analyze requirements step by step to determine the final SQL query to generate. + To improve SQL generation accuracy, please evaluate the accuracy of the SQL after generation, + if there are issues, regenerate the SQL. + When SQL execution fails, you need to correct the SQL based on the error message and try to execute again. + + ### Tools ### + execute_sql: Can execute SQL by passing in SQL statements and return execution results + """ + user_prompt = """ + Below is the database information I need to query: + {schema} + + My requirement is: {question} + """ + # Define prompt template + self.prompt = ChatPromptTemplate.from_messages([ + ("system", system_prompt), + ("human", user_prompt) + ]) + + def generate_sql(self, question: str) -> str: + chain = self.prompt | self.agent_executor + schema = self.db.get_table_info() + return chain.invoke({"schema": schema, "question": question}) + + async def async_generate(self, question: str) -> AsyncGenerator[str, None]: + + chain = self.prompt | self.agent_executor + schema = self.db.get_table_info() + + async for chunk in chain.astream({"schema": schema, "question": question}): + if not isinstance(chunk, dict): + continue + + if "agent" in chunk: + messages = chunk["agent"].get("messages", []) + for msg in messages: + if tool_calls := msg.additional_kwargs.get("tool_calls"): + for tool_call in tool_calls: + response = { + "type": "tool_call", + "tool": tool_call["function"]["name"], + "args": tool_call["function"]["arguments"] + } + yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + + if content := msg.content: + html_start = content.find("```html") + html_end = content.find("```", html_start + 6) + if html_start != -1 and html_end != -1: + html_content = content[html_start + 7:html_end].strip() + response = { + "type": "final", + "content": content.split("```html")[0].strip(), + "html": html_content + } + else: + response = { + "type": "final", + "content": content + } + yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + + if "tools" in chunk: + messages = chunk["tools"].get("messages", []) + for msg in messages: + response = { + "type": "tool_result", + "tool": msg.name, + "content": msg.content + } + yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + + await asyncio.sleep(0.1) + + yield f"data: {json.dumps({'type': 'complete'})}\n\n" + +def execute_sql(ds: CoreDatasource, sql: str) -> str: + """Execute SQL query + + Args: + ds: Data source instance + sql: SQL query statement + + Returns: + Query results + """ + print(f"Executing SQL on ds_id {ds.id}: {sql}") + return exec_sql(ds, sql) + +def execute_sql_with_db(db: SQLDatabase, sql: str) -> str: + """Execute SQL query using SQLDatabase + + Args: + db: SQLDatabase instance + sql: SQL query statement + + Returns: + str: Query results formatted as string + """ + try: + # Execute query + result = db.run(sql) + + if not result: + return "Query executed successfully but returned no results." + + # Format results + return str(result) + + except Exception as e: + error_msg = f"SQL execution failed: {str(e)}" + logging.error(error_msg) + raise RuntimeError(error_msg) \ No newline at end of file diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 90b374d3..8bd7d4a8 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -5,8 +5,18 @@ import urllib.parse from typing import Any import json from apps.datasource.utils.utils import aes_decrypt +from common.core.deps import SessionDep +def get_uri(ds: CoreDatasource): + conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) + db_url: str + if ds.type == "mysql": + db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}" + else: + raise 'The datasource type not support.' + return db_url + def get_session(ds: CoreDatasource): conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) db_url: str @@ -88,3 +98,21 @@ def exec_sql(ds: CoreDatasource, sql: str): result.close() if session is not None: session.close() + +def exec_sql(ds: CoreDatasource, sql: str): + ds = session.get(CoreDatasource, id) + session = get_session(ds) + result = session.execute(text(sql)) + try: + columns = result.keys()._keys + res = result.fetchall() + result_list = [ + {columns[i]: value for i, value in enumerate(tuple_item)} + for tuple_item in res + ] + return {"fields": columns, "data": result_list} + finally: + if result is not None: + result.close() + if session is not None: + session.close() \ No newline at end of file diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 19a770a3..692132b8 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -19,10 +19,11 @@ dependencies = [ "sentry-sdk[fastapi]<2.0.0,>=1.40.6", "pyjwt<3.0.0,>=2.8.0", "pycryptodome (>=3.22.0,<4.0.0)", - "langchain>=0.1.0,<0.2.0", - "langchain-core>=0.1.10,<0.2.0", - "langchain-openai>=0.1.0,<0.2.0", - "langchain-community>=0.0.19,<0.1.0", + "langchain>=0.3,<0.4", + "langchain-core>=0.3,<0.4", + "langchain-openai>=0.3,<0.4", + "langchain-community>=0.3,<0.4", + "langgraph>=0.3,<0.4", "vllm>=0.8.5", "dashscope>=1.14.0,<2.0.0", "sse-starlette>=1.8.0,<2.0.0", diff --git a/frontend/src/api/chat.ts b/frontend/src/api/chat.ts index c96d9d64..3a5ec56d 100644 --- a/frontend/src/api/chat.ts +++ b/frontend/src/api/chat.ts @@ -2,8 +2,13 @@ import { request } from '@/utils/request' export const questionApi = { pager: (pageNumber: number, pageSize: number) => request.get(`/chat/question/pager/${pageNumber}/${pageSize}`), - // add: (data: any, progress: any) => request.post('/chat/question', data, { responseType: 'stream', onDownloadProgress: progress }), - add: (data: any) => request.post('/chat/question', data), + /* add: (data: any) => new Promise((resolve, reject) => { + request.post('/chat/question', data, { responseType: 'stream', timeout: 0, onDownloadProgress: p => { + resolve(p) + }}).catch(e => reject(e)) + }), */ + // add: (data: any) => request.post('/chat/question', data), + add: (data: any) => request.fetchStream('/chat/question', data), edit: (data: any) => request.put('/chat/question', data), delete: (id: number) => request.delete(`/chat/question/${id}`), query: (id: number) => request.get(`/chat/question/${id}`) diff --git a/frontend/src/utils/request.ts b/frontend/src/utils/request.ts index c314c492..ad83d3f6 100644 --- a/frontend/src/utils/request.ts +++ b/frontend/src/utils/request.ts @@ -49,7 +49,7 @@ class HttpService { this.cancelTokenSource = axios.CancelToken.source() this.instance = axios.create({ baseURL: import.meta.env.VITE_API_BASE_URL, - timeout: 50000, + timeout: 100000, headers: { 'Content-Type': 'application/json', ...config?.headers @@ -79,7 +79,6 @@ class HttpService { return Promise.reject(error) } ) - // Response interceptor this.instance.interceptors.response.use( (response: AxiosResponse) => { @@ -207,6 +206,22 @@ class HttpService { return this.request({ ...config, method: 'POST', url, data }) } + public fetchStream(url: string ,data?: any): Promise { + const token = wsCache.get('user.token') + const heads: any = { + 'Content-Type': 'application/json', + } + if (token) { + heads['X-SQLBOT-TOKEN'] = `Bearer ${token}` + } + const real_url = import.meta.env.VITE_API_BASE_URL + return fetch(real_url + url, { + method: 'POST', + headers: heads, + body: JSON.stringify(data), + }) + } + // PUT request public put( url: string, diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index b97139c9..2cd6482f 100644 --- a/frontend/src/views/chat/index.vue +++ b/frontend/src/views/chat/index.vue @@ -48,10 +48,17 @@ v-for="(msg, index) in currentChat.messages" :key="index" class="message" - :class="msg.role" + :class="[msg.role, { 'is-typing': msg.isTyping }]" > - -
+
+
+ {{ msg.content }} +
@@ -92,10 +99,10 @@ import { ref, computed, nextTick, watch } from 'vue' import { Plus, Search, More, Position } from '@element-plus/icons-vue' import { questionApi } from '@/api/chat' - interface ChatMessage { role: 'user' | 'assistant' content: string + isTyping?: boolean } interface Chat { @@ -135,7 +142,7 @@ const switchChat = (chatId: string) => { } const scrollToBottom = () => { - nextTick(() => { + /* nextTick(() => { if (messagesRef.value) { const scrollHeight = messagesRef.value.scrollHeight const animateScroll = () => { @@ -148,20 +155,30 @@ const scrollToBottom = () => { } animateScroll() } - }) + }) */ } -const sendMessage = () => { +const updateMessageContent = (index: number, content: string) => { + if (currentChat.value) { + currentChat.value.messages[index] = { + ...currentChat.value.messages[index], + content + } + } +} + +const sendMessage = async () => { if (!inputMessage.value.trim()) return let chat = currentChat.value - // if (!chat) return if (!chat) { chat = { - id: null, - title: 'test', + id: Date.now().toString(), + title: 'New Chat', messages: [] - } as unknown as Chat + } + chatHistory.value.unshift(chat) + currentChatId.value = chat.id } chat.messages.push({ @@ -169,23 +186,67 @@ const sendMessage = () => { content: inputMessage.value }) - const param = { - question: inputMessage.value - } - inputMessage.value = '' - chatHistory.value.unshift(chat) - currentChatId.value = chat.id - - - // const index = chat.messages.length - 1 - questionApi.add(param).then((result: string) => { - // chat.messages[index]['content'] += result - chat.messages.push({ - role: 'assistant', - content: result - }) - scrollToBottom() + const defaultMsg = 'Thinking...' + chat.messages.push({ + role: 'assistant', + content: defaultMsg, + isTyping: true }) + + const userInput = inputMessage.value + inputMessage.value = '' + + const index = chat.messages.length - 1 + + try { + const response = await questionApi.add({ question: userInput }) + const reader = response.body.getReader() + const decoder = new TextDecoder() + + while (true) { + const { done, value } = await reader.read() + if (done) { + chat.messages[index].isTyping = false + break + } + + const chunk = decoder.decode(value) + const lines = chunk.split('\n\n') + + for (const line of lines) { + if (line.startsWith('data: ')) { + const data = JSON.parse(line.replace('data:', '').trim()) + let realContent = data.content + + switch (data.type) { + case 'html': + realContent = '\n' + data.content + break + + case 'error': + realContent = '\nError: ' + data.content + break + + default: + if (data.content) { + const currentContent = chat.messages[index].content + const newContent = currentContent.startsWith(defaultMsg) + ? realContent + : currentContent + realContent + + updateMessageContent(index, newContent) + await nextTick() + scrollToBottom() + } + } + } + } + } + } catch (error) { + console.error('Error:', error) + updateMessageContent(index, chat.messages[index].content + '\nError: Failed to get response') + chat.messages[index].isTyping = false + } } watch(() => currentChat.value?.messages, () => { @@ -231,7 +292,7 @@ const handleCtrlEnter = (e: KeyboardEvent) => {