Merge remote-tracking branch 'origin/main'

# Conflicts:
#	frontend/components.d.ts

Author: ziyujiahao

backend/alembic/versions/ff653d5df198_007_add_chat.py

@@ -0,0 +1,52 @@
+"""006_add_chat
+
+Revision ID: ff653d5df198
+Revises: 0a6f11be9be4
+Create Date: 2025-05-21 15:44:29.763091
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = 'ff653d5df198'
+down_revision = 'e6276ddab06e'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('chat',
+    sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
+    sa.Column('create_time', sa.DateTime(timezone=True), nullable=True),
+    sa.Column('create_by', sa.BigInteger(), nullable=True),
+    sa.Column('brief', sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True),
+    sa.Column('chat_type', sqlmodel.sql.sqltypes.AutoString(length=20), nullable=False),
+    sa.Column('datasource', sa.Integer(), nullable=False),
+    sa.Column('engine_type', sqlmodel.sql.sqltypes.AutoString(length=64), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_table('chat_record',
+    sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
+    sa.Column('chat_id', sa.Integer(), nullable=True),
+    sa.Column('create_time', sa.DateTime(timezone=True), nullable=True),
+    sa.Column('create_by', sa.BigInteger(), nullable=True),
+    sa.Column('datasource', sa.Integer(), nullable=False),
+    sa.Column('engine_type', sqlmodel.sql.sqltypes.AutoString(length=64), nullable=False),
+    sa.Column('question', sa.Text(), nullable=True),
+    sa.Column('full_question', sa.Text(), nullable=True),
+    sa.Column('answer', sa.Text(), nullable=True),
+    sa.Column('run_time', sa.Float(), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('chat_record')
+    op.drop_table('chat')
+    # ### end Alembic commands ###

backend/apps/chat/api/chat.py

@@ -2,11 +2,13 @@
 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.curd.chat import list_chats, get_chat_with_records, create_chat, save_question, save_answer, rename_chat, \
+    delete_chat
+from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat
 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.crud.datasource import get_table_obj_by_ds
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AiModelDetail
 from common.core.deps import SessionDep, CurrentUser
@@ -32,6 +34,28 @@ async def list_chat(session: SessionDep, current_user: CurrentUser, chart_id: in
         )
 
 
+@router.post("/rename")
+async def rename(session: SessionDep, chat: RenameChat):
+    try:
+        return rename_chat(session=session, rename_object=chat)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
+
+
+@router.get("/delete/{chart_id}")
+async def delete(session: SessionDep, chart_id: int):
+    try:
+        return delete_chat(session=session, chart_id=chart_id)
+    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:
@@ -57,26 +81,18 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     """
     question = request_question.question
 
-    # 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:
+    chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first()
+    if not chat:
         raise HTTPException(
             status_code=400,
-            detail="No available AI model configuration found"
+            detail=f"Chat with id {request_question.chart_id} not found"
         )
 
+    # Get available datasource
+    ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
     if not ds:
         raise HTTPException(
-            status_code=400,
+            status_code=500,
             detail="No available datasource configuration found"
         )
 
@@ -89,6 +105,17 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
             detail=str(e1)
         )
 
+    # Get available AI model
+    aimodel = session.exec(select(AiModelDetail).where(
+        AiModelDetail.status == True,
+        AiModelDetail.api_key.is_not(None)
+    )).first()
+    if not aimodel:
+        raise HTTPException(
+            status_code=400,
+            detail="No available AI model configuration found"
+        )
+
     # Use Tongyi Qianwen
     tongyi_config = LLMConfig(
         model_type="openai",
@@ -113,10 +140,39 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     """ result = llm_service.generate_sql(question)
     return result """
 
+    # get schema
+    schema_str = ""
+    table_objs = get_table_obj_by_ds(session=session, ds=ds)
+    db_name = table_objs[0].schema
+    schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
+    for obj in table_objs:
+        schema_str += f"# Table: {db_name}.{obj.table.table_name}"
+        table_comment = ''
+        if obj.table.custom_comment:
+            table_comment = obj.table.custom_comment.strip()
+        if table_comment == '':
+            schema_str += '\n[\n'
+        else:
+            schema_str += f", {table_comment}\n[\n"
+
+        field_list = []
+        for field in obj.fields:
+            field_comment = ''
+            if field.custom_comment:
+                field_comment = field.custom_comment.strip()
+            if field_comment == '':
+                field_list.append(f"({field.field_name}:{field.field_type})")
+            else:
+                field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})")
+        schema_str += ",\n".join(field_list)
+        schema_str += '\n]\n'
+
+    print(schema_str)
+
     async def event_generator():
         all_text = ''
         try:
-            async for chunk in llm_service.async_generate(question):
+            async for chunk in llm_service.async_generate(question, schema_str):
                 data = json.loads(chunk.replace('data: ', ''))
 
                 if data['type'] in ['final', 'tool_result']:

backend/apps/chat/curd/chat.py

@@ -5,7 +5,7 @@
 from sqlalchemy import text, and_
 from sqlmodel import select
 
-from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo
+from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat
 from apps.chat.schemas.chat_schema import ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
 from common.core.deps import SessionDep, CurrentUser
@@ -17,6 +17,32 @@ def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
     return chart_list
 
 
+def rename_chat(session: SessionDep, rename_object: RenameChat) -> str:
+    chat = session.query(Chat).filter(Chat.id == rename_object.id).first()
+    if not chat:
+        raise Exception(f"Chat with id {rename_object.id} not found")
+
+    chat.brief = rename_object.brief.strip()[:20]
+    session.add(chat)
+    session.flush()
+    session.refresh(chat)
+
+    brief = chat.brief
+    session.commit()
+    return brief
+
+
+def delete_chat(session, chart_id) -> str:
+    chat = session.query(Chat).filter(Chat.id == chart_id).first()
+    if not chat:
+        return f'Chat with id {chart_id} has been deleted'
+
+    session.delete(chat)
+    session.commit()
+
+    return f'Chat with id {chart_id} has been deleted'
+
+
 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:
@@ -33,7 +59,7 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
         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()
+        and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
 
     chat_info.records = record_list
 
@@ -98,10 +124,11 @@ def save_question(session: SessionDep, current_user: CurrentUser, question: Chat
 
     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 = session.query(ChatRecord).filter(ChatRecord.id == id).first()
     record.full_question = full_question
 
     result = ChatRecord(**record.model_dump())
@@ -114,11 +141,12 @@ def save_full_question(session: SessionDep, id: int, full_question: str) -> Chat
 
     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 = session.query(ChatRecord).filter(ChatRecord.id == id).first()
     record.answer = answer
 
     result = ChatRecord(**record.model_dump())

backend/apps/chat/models/chat_model.py

@@ -36,6 +36,11 @@ class CreateChat(BaseModel):
     datasource: int = None
 
 
+class RenameChat(BaseModel):
+    id: int = None
+    brief: str = ''
+
+
 class ChatInfo(BaseModel):
     id: Optional[int] = None
     create_time: datetime = None

backend/apps/chat/schemas/llm.py

@@ -101,16 +101,17 @@ def generate_sql(self, question: str) -> str:
         schema = self.db.get_table_info()
         return chain.invoke({"schema": schema, "question": question})
 
-    async def async_generate(self, question: str) -> AsyncGenerator[str, None]:
+    async def async_generate(self, question: str, schema: str) -> AsyncGenerator[str, None]:
 
         chain = self.prompt | self.agent_executor
         # schema = self.db.get_table_info()
 
-        schema_engine = SchemaEngine(engine=self.db._engine)
-        mschema = schema_engine.mschema
-        mschema_str = mschema.to_mschema()
+        # schema_engine = SchemaEngine(engine=self.db._engine)
+        # mschema = schema_engine.mschema
+        # mschema_str = mschema.to_mschema()
 
-        async for chunk in chain.astream({"schema": mschema_str, "question": question}):
+        # async for chunk in chain.astream({"schema": mschema_str, "question": question}):
+        async for chunk in chain.astream({"schema": schema, "question": question}):
             if not isinstance(chunk, dict):
                 continue
 

backend/apps/datasource/crud/datasource.py

@@ -211,10 +211,12 @@ def preview(session: SessionDep, id: int, data: TableObj):
     return exec_sql(ds, sql)
 
 
-def get_tableobj_by_ds(session: SessionDep, id: int) -> List[TableAndFields]:
-    list: List = []
-    tables = session.query(CoreTable).filter(CoreTable.ds_id == id).all()
+def get_table_obj_by_ds(session: SessionDep, ds: CoreDatasource) -> List[TableAndFields]:
+    _list: List = []
+    tables = session.query(CoreTable).filter(CoreTable.ds_id == ds.id).all()
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
+    schema = conf.dbSchema if conf.dbSchema is not None and conf.dbSchema != "" else conf.database
     for table in tables:
         fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all()
-        list.append(TableAndFields(table=table, fields=fields))
-    return list
+        _list.append(TableAndFields(schema=schema, table=table, fields=fields))
+    return _list

backend/apps/datasource/models/datasource.py

@@ -111,9 +111,11 @@ def __init__(self, attr1, attr2, attr3):
 
 
 class TableAndFields:
-    def __init__(self, table, fields):
+    def __init__(self, schema, table, fields):
+        self.schema = schema
         self.table = table
         self.fields = fields
 
+    schema: str
     table: CoreTable
     fields: List[CoreField]

backend/pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "pymssql (>=2.3.4,<3.0.0)",
     "pandas (>=2.2.3,<3.0.0)",
     "openpyxl (>=3.1.5,<4.0.0)",
-    "psycopg2 (>=2.9.10,<3.0.0)",
+    "psycopg2-binary (>=2.9.10,<3.0.0)",
     "oracledb (>=3.1.1,<4.0.0)"
 ]
 [[tool.uv.index]]

frontend/src/api/chat.ts

@@ -16,6 +16,14 @@ export const questionApi = {
     query: (id: number) => request.get(`/chat/question/${id}`)
 }
 
+export interface ChatMessage {
+    role: 'user' | 'assistant'
+    create_time?: Date | string
+    content?: string | number
+    isTyping?: boolean
+    isWelcome?: boolean
+}
+
 export class ChatRecord {
     id?: number
     chat_id?: number
@@ -147,12 +155,22 @@ function startChat(data: any): Promise<ChatInfo> {
     return request.post('/chat/start', data)
 }
 
+function renameChat(chat_id: number | undefined, brief: string): Promise<string> {
+    return request.post('/chat/rename', {id: chat_id, brief: brief})
+}
+
+function deleteChat(id: number | undefined): Promise<string> {
+    return request.get(`/chat/delete/${id}`)
+}
+
 export const chatApi = {
     toChatRecord,
     toChatRecordList,
     toChatInfo,
     toChatInfoList,
     list,
     get,
-    startChat
+    startChat,
+    renameChat,
+    deleteChat
 }