feat: Q&A Chat

Author: ulleo

backend/alembic/env.py

@@ -25,7 +25,8 @@
 # target_metadata = None
 
 # from apps.system.models.user import SQLModel  # noqa
-from apps.settings.models.setting_models import SQLModel 
+# from apps.settings.models.setting_models import SQLModel
+from apps.chat.models.chat_model import SQLModel
 from common.core.config import settings # noqa
 
 target_metadata = SQLModel.metadata

backend/apps/chat/api/chat.py

@@ -1,37 +1,68 @@
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 from sqlmodel import select
+
+from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, save_question, save_answer
+from apps.chat.models.chat_model import CreateChat, ChatRecord
 from apps.chat.schemas.chat_base_schema import LLMConfig
 from apps.chat.schemas.chat_schema import ChatQuestion
 from apps.chat.schemas.llm import AgentService
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AiModelDetail
-from common.core.deps import SessionDep
+from common.core.deps import SessionDep, CurrentUser
 import json
 import asyncio
 
 router = APIRouter(tags=["Data Q&A"], prefix="/chat")
 
 
+@router.get("/list")
+async def chats(session: SessionDep, current_user: CurrentUser):
+    return list_chats(session, current_user)
+
+
+@router.get("/get/{chart_id}")
+async def list_chat(session: SessionDep, current_user: CurrentUser, chart_id: int):
+    try:
+        return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
+
+
+@router.post("/start")
+async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat):
+    try:
+        return create_chat(session, current_user, create_chat_obj)
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail=str(e)
+        )
+
+
 @router.post("/question")
-async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion):
+async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
     """Stream SQL analysis results
     
     Args:
         session: Database session
-        requestQuestion: User question model
+        current_user: CurrentUser
+        request_question: User question model
         
     Returns:
         Streaming response with analysis results
     """
-    question = requestQuestion.question
-    
+    question = request_question.question
+
     # Get available AI model
     aimodel = session.exec(select(AiModelDetail).where(
-        AiModelDetail.status == True, 
+        AiModelDetail.status == True,
         AiModelDetail.api_key.is_not(None)
     )).first()
-    
+
     # Get available datasource
     ds = session.exec(select(CoreDatasource).where(
         CoreDatasource.status == 'Success'
@@ -42,13 +73,22 @@ async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion):
             status_code=400,
             detail="No available AI model configuration found"
         )
-        
+
     if not ds:
         raise HTTPException(
             status_code=400,
             detail="No available datasource configuration found"
         )
-    
+
+    record: ChatRecord
+    try:
+        record = save_question(session=session, current_user=current_user, question=request_question)
+    except Exception as e1:
+        raise HTTPException(
+            status_code=400,
+            detail=str(e1)
+        )
+
     # Use Tongyi Qianwen
     tongyi_config = LLMConfig(
         model_type="openai",
@@ -72,25 +112,37 @@ async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion):
     vllm_service = LLMService(vllm_config) """
     """ result = llm_service.generate_sql(question)
     return result """
-    
+
     async def event_generator():
+        all_text = ''
         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']
+                    print('--  ' + content)
+                    all_text += content
                     for char in content:
                         yield f"data: {json.dumps({'type': 'char', 'content': char})}\n\n"
-                        await asyncio.sleep(0.05) 
-                    
+                        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:
+            all_text = 'Exception:' + str(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")
+
+        try:
+            save_answer(session=session, id=record.id, answer=all_text)
+        except Exception as e:
+            raise HTTPException(
+                status_code=500,
+                detail=str(e)
+            )
+
+    # return EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"})
+    return StreamingResponse(event_generator(), media_type="text/event-stream")

backend/apps/chat/curd/__init__.py

@@ -0,0 +1,132 @@
+import datetime
+import json
+from typing import List
+
+from sqlalchemy import text, and_
+from sqlmodel import select
+
+from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo
+from apps.chat.schemas.chat_schema import ChatQuestion
+from apps.datasource.models.datasource import CoreDatasource
+from common.core.deps import SessionDep, CurrentUser
+
+
+def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
+    chart_list = session.query(Chat).filter(Chat.create_by == current_user.id).order_by(
+        Chat.create_time.desc()).all()
+    return chart_list
+
+
+def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> ChatInfo:
+    chat = session.query(Chat).filter(Chat.id == chart_id).first()
+    if not chat:
+        raise Exception(f"Chat with id {chart_id} not found")
+
+    chat_info = ChatInfo(**chat.model_dump())
+
+    ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
+    if not ds:
+        chat_info.datasource_exists = False
+        chat_info.datasource_name = 'Datasource not exist'
+    else:
+        chat_info.datasource_exists = True
+        chat_info.datasource_name = ds.name
+
+    record_list = session.query(ChatRecord).filter(
+        and_(Chat.create_by == current_user.id, ChatRecord.id == chart_id)).order_by(ChatRecord.create_time).all()
+
+    chat_info.records = record_list
+
+    return chat_info
+
+
+def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat) -> ChatInfo:
+    if not create_chat_obj.datasource:
+        raise Exception("Datasource cannot be None")
+
+    if not create_chat_obj.question or create_chat_obj.question.strip() == '':
+        raise Exception("Question cannot be Empty")
+
+    chat = Chat(create_time=datetime.datetime.now(),
+                create_by=current_user.id,
+                brief=create_chat_obj.question.strip()[:20],
+                datasource=create_chat_obj.datasource)
+
+    ds = session.query(CoreDatasource).filter(CoreDatasource.id == create_chat_obj.datasource).first()
+
+    if not ds:
+        raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")
+
+    chat.engine_type = ds.type_name
+
+    chat_info = ChatInfo(**chat.model_dump())
+
+    session.add(chat)
+    session.flush()
+    session.refresh(chat)
+    chat_info.id = chat.id
+    session.commit()
+
+    return chat_info
+
+
+def save_question(session: SessionDep, current_user: CurrentUser, question: ChatQuestion) -> ChatRecord:
+    if not question.chat_id:
+        raise Exception("ChatId cannot be None")
+    if not question.question or question.question.strip() == '':
+        raise Exception("Question cannot be Empty")
+
+    chat = session.query(Chat).filter(Chat.id == question.chat_id).first()
+    if not chat:
+        raise Exception(f"Chat with id {question.chat_id} not found")
+
+    record = ChatRecord()
+    record.question = question.question
+    record.chat_id = chat.id
+    record.create_time = datetime.datetime.now()
+    record.create_by = current_user.id
+    record.datasource = chat.datasource
+    record.engine_type = chat.engine_type
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+    result.id = record.id
+    session.commit()
+
+    return result
+
+def save_full_question(session: SessionDep, id: int, full_question: str) -> ChatRecord:
+    if not id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(Chat.id == id).first()
+    record.full_question = full_question
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+def save_answer(session: SessionDep, id: int, answer: str) -> ChatRecord:
+    if not id:
+        raise Exception("Record id cannot be None")
+
+    record = session.query(ChatRecord).filter(Chat.id == id).first()
+    record.answer = answer
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result

backend/apps/chat/curd/chat.py

@@ -0,0 +1,49 @@
+from sqlmodel import SQLModel, Field
+from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity
+from datetime import datetime
+from pydantic import BaseModel
+from typing import List, Optional
+
+
+class Chat(SQLModel, table=True):
+    __tablename__ = "chat"
+    id: Optional[int] = Field(sa_column=Column(Integer, Identity(always=True), primary_key=True))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
+    brief: str = Field(max_length=64, nullable=True)
+    chat_type: str = Field(max_length=20, default="chat")  # chat, datasource
+    datasource: int = Field(sa_column=Column(Integer, nullable=False))
+    engine_type: str = Field(max_length=64)
+
+
+class ChatRecord(SQLModel, table=True):
+    __tablename__ = "chat_record"
+    id: Optional[int] = Field(sa_column=Column(Integer, Identity(always=True), primary_key=True))
+    chat_id: int = Field(sa_column=Column(Integer))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
+    datasource: int = Field(sa_column=Column(Integer, nullable=False))
+    engine_type: str = Field(max_length=64)
+    question: str = Field(sa_column=Column(Text, nullable=True))
+    full_question: str = Field(sa_column=Column(Text, nullable=True))
+    answer: str = Field(sa_column=Column(Text, nullable=True))
+    run_time: float = Field(default=0)
+
+
+class CreateChat(BaseModel):
+    id: int = None
+    question: str = ''
+    datasource: int = None
+
+
+class ChatInfo(BaseModel):
+    id: Optional[int] = None
+    create_time: datetime = None
+    create_by: int = None
+    brief: str = ''
+    chat_type: str = "chat"
+    datasource: int = None
+    engine_type: str = ''
+    datasource_name: str = ''
+    datasource_exists: bool = True
+    records: List[ChatRecord] = []

backend/apps/chat/models/__init__.py

@@ -3,4 +3,5 @@
 
 
 class ChatQuestion(BaseModel):
-    question: str
+    question: str
+    chat_id: int

backend/apps/chat/models/chat_model.py

@@ -36,7 +36,7 @@ dependencies = [
     "oracledb (>=3.1.1,<4.0.0)"
 ]
 [[tool.uv.index]]
-url = "https://pypi.tuna.tsinghua.edu.cn/simple"
+url = "https://mirrors.aliyun.com/pypi/simple"
 default = true
 
 [tool.uv]

backend/apps/chat/schemas/chat_schema.py

@@ -13,7 +13,8 @@
     "snowflake-id": "^1.1.0",
     "vue": "^3.5.13",
     "vue-router": "^4.5.0",
-    "web-storage-cache": "^1.1.1"
+    "web-storage-cache": "^1.1.1",
+    "dayjs": "^1.11.13"
   },
   "devDependencies": {
     "@element-plus/icons-vue": "^2.3.1",