perf: Data Q&A

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
+    """ 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")

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
         )
     

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)

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()

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",

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}`)

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<any> {
+    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<T = any>(
     url: string,

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 }]"
           >
-            <!-- <div class="message-content">{{ msg.content }}</div> -->
-            <div class="message-content" v-html="formatMessage(msg.content)"></div>
+            <div
+              v-if="msg.role === 'assistant'"
+              class="message-content"
+              :class="{ 'typing': msg.isTyping }"
+              v-html="formatMessage(msg.content)"
+            ></div>
+            <div v-else class="message-content">
+              {{ msg.content }}
+            </div>
           </div>
         </template>
         <div v-else class="empty-state">
@@ -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) => {
 
 <style lang="less" scoped>
 .chat-container {
-  height: 100%;
+  height: calc(100vh - 61px); 
   display: flex;
   background-color: var(--white);
   border-radius: var(--border-radius);
@@ -305,11 +366,16 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
     flex: 1;
     display: flex;
     flex-direction: column;
+    height: 100%;
+    position: relative;
+    overflow: hidden;
     
     .chat-messages {
-      flex: 1;
+      // flex: 1;
+      height: calc(100% - 188px);
       overflow-y: auto;
       padding: 24px;
+      scroll-behavior: smooth;
       
       &::-webkit-scrollbar {
         width: 6px;
@@ -317,11 +383,11 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
       }
       
       &::-webkit-scrollbar-thumb {
-        background: #ccc;
+        background: rgba(0, 0, 0, 0.2);
         border-radius: 3px;
         
         &:hover {
-          background: #999;
+          background: rgba(0, 0, 0, 0.3);
         }
       }
       
@@ -335,6 +401,7 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
 
       .message {
         margin-bottom: 24px;
+        transition: all 0.3s ease;
 
         &.user {
           text-align: right;
@@ -343,9 +410,11 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
           }
         }
 
-        &.assistant .message-content {
-          background-color: #f5f5f5;
-          text-align: left;
+        &.assistant {
+          .message-content {
+            background-color: #f5f5f5;
+            text-align: left;
+          }
         }
 
         .message-content {
@@ -353,6 +422,7 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
           padding: 12px 16px;
           border-radius: 8px;
           max-width: 80%;
+          transition: all 0.3s ease;
         }
       }
 
@@ -383,9 +453,14 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
     }
 
     .chat-input {
-      padding: 16px 24px;
-      border-top: 1px solid var(--el-border-color-lighter);
+      position: absolute;
+      bottom: 0;
+      left: 0;
+      right: 0;
       background: #fff;
+      border-top: 1px solid var(--el-border-color-lighter);
+      padding: 16px 24px;
+      z-index: 10;
 
       .input-wrapper {
         position: relative;
@@ -402,8 +477,6 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
         :deep(.el-textarea__inner) {
           padding: 16px;
           font-size: 14px;
-          // line-height: 1.6;
-          // min-height: 80px;
           resize: none;
           border: none;
           background: transparent;
@@ -427,8 +500,6 @@ const handleCtrlEnter = (e: KeyboardEvent) => {
           padding: 8px;
         }
       }
-
-    
     }
   }
 }