feat: generate sql with history

Author: ulleo

backend/alembic/versions/1f077c30e476_009_modify_chat.py

@@ -0,0 +1,64 @@
+"""009_modify_chat
+
+Revision ID: 1f077c30e476
+Revises: 35d925df4568
+Create Date: 2025-05-30 16:11:08.020715
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '1f077c30e476'
+down_revision = '35d925df4568'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('chat_record')
+    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('finish_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('sql_answer', sa.Text(), nullable=True),
+    sa.Column('sql', sa.Text(), nullable=True),
+    sa.Column('sql_exec_result', sa.Text(), nullable=True),
+    sa.Column('data', sa.Text(), nullable=True),
+    sa.Column('chart_answer', sa.Text(), nullable=True),
+    sa.Column('chart', sa.Text(), nullable=True),
+    sa.Column('full_sql_message', sa.Text(), nullable=True),
+    sa.Column('full_chart_message', sa.Text(), nullable=True),
+    sa.Column('finish', sa.Boolean(), nullable=True),
+    sa.Column('error', 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.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 ###

backend/apps/ai_model/__init__.py

@@ -0,0 +1 @@
+# todo

backend/apps/ai_model/llm.py

@@ -3,6 +3,10 @@
 from abc import ABC, abstractmethod
 from langchain_core.language_models import BaseLLM as LangchainBaseLLM
 from langchain_openai import ChatOpenAI
+
+from apps.system.models.system_model import AiModelDetail
+
+
 # from langchain_community.llms import Tongyi, VLLM
 
 class LLMConfig(BaseModel):
@@ -84,4 +88,16 @@ def create_llm(cls, config: LLMConfig) -> BaseLLM:
     @classmethod
     def register_llm(cls, model_type: str, llm_class: Type[BaseLLM]):
         """Register new model type"""
-        cls._llm_types[model_type] = llm_class
+        cls._llm_types[model_type] = llm_class
+
+
+#  todo
+def get_llm_config(aimodel: AiModelDetail) -> LLMConfig:
+    config = LLMConfig(
+        model_type="openai",
+        model_name=aimodel.name,
+        api_key=aimodel.api_key,
+        api_base_url=aimodel.endpoint,
+        additional_params={"temperature": aimodel.temperature}
+    )
+    return config

backend/apps/chat/schemas/chat_base_schema.py renamed to backend/apps/ai_model/model_factory.py

@@ -0,0 +1 @@
+# todo

backend/apps/ai_model/openai/__init__.py

@@ -1,19 +1,19 @@
+import asyncio
+import json
+from typing import List
+
 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, 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
+    delete_chat, list_records
+from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion
+from apps.chat.task.llm import LLMService
+from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AiModelDetail
 from common.core.deps import SessionDep, CurrentUser
-import json
-import asyncio
 
 router = APIRouter(tags=["Data Q&A"], prefix="/chat")
 
@@ -96,14 +96,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
             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)
-        )
+    request_question.engine = ds.type_name
 
     # Get available AI model
     aimodel = session.exec(select(AiModelDetail).where(
@@ -112,93 +105,54 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     )).first()
     if not aimodel:
         raise HTTPException(
-            status_code=400,
+            status_code=500,
             detail="No available AI model configuration found"
         )
 
-    # Use Tongyi Qianwen
-    tongyi_config = LLMConfig(
-        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 = AgentService(tongyi_config, ds)
-
-    # Use Custom VLLM model
-    """ vllm_config = LLMConfig(
-        model_type="vllm",
-        model_name="your-model-path",
-        additional_params={
-            "max_new_tokens": 200,
-            "temperature": 0.3
-        }
-    )
-    vllm_service = LLMService(vllm_config) """
-    """ result = llm_service.generate_sql(question)
-    return result """
-
+    history_records: List[ChatRecord] = list_records(session=session, current_user=current_user,
+                                                     chart_id=request_question.chat_id)
     # 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, schema_str):
-                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)
-
-                    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"
-
-        try:
-            save_answer(session=session, id=record.id, answer=all_text)
-        except Exception as e:
-            raise HTTPException(
-                status_code=500,
-                detail=str(e)
-            )
+    request_question.db_schema = get_table_schema(session=session, ds=ds)
+    llm_service = LLMService(request_question, history_records, ds, aimodel)
+
+    llm_service.init_record(session=session, current_user=current_user)
+
+    def run_task():
+        sql_res = llm_service.generate_sql(session=session)
+        for chunk in sql_res:
+            yield json.dumps({'content': chunk, 'type': 'sql'}) + '\n\n'
+        yield json.dumps({'type': 'info', 'msg': 'sql generated'}) + '\n\n'
+
+    # async def event_generator():
+    #     all_text = ''
+    #     try:
+    #         async for chunk in llm_service.async_generate(question, request_question.db_schema):
+    #             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)
+    #
+    #                 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"
+    #
+    #     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")
+    return StreamingResponse(run_task(), media_type="text/event-stream")

backend/apps/ai_model/openai/llm.py

@@ -1,12 +1,9 @@
 import datetime
-import json
 from typing import List
 
-from sqlalchemy import text, and_
-from sqlmodel import select
+from sqlalchemy import and_
 
-from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat
-from apps.chat.schemas.chat_schema import ChatQuestion
+from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
 from common.core.deps import SessionDep, CurrentUser
 
@@ -66,6 +63,12 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
     return chat_info
 
 
+def list_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> List[ChatRecord]:
+    record_list = session.query(ChatRecord).filter(
+        and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
+    return record_list
+
+
 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")
@@ -125,11 +128,33 @@ def save_question(session: SessionDep, current_user: CurrentUser, question: Chat
     return result
 
 
-def save_full_question(session: SessionDep, id: int, full_question: str) -> ChatRecord:
+def save_full_sql_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord:
+    return save_full_sql_message_and_answer(session=session, record_id=record_id, full_message=full_message, answer='')
+
+
+def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer: str, full_message: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.full_sql_message = full_message
+    record.sql_answer = answer
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
+def save_full_chart_message(session: SessionDep, id: int, full_message: str) -> ChatRecord:
     if not id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == id).first()
-    record.full_question = full_question
+    record.full_chart_message = full_message
 
     result = ChatRecord(**record.model_dump())