feat: chat use local db schema

Author: ulleo

backend/apps/chat/api/chat.py

@@ -4,10 +4,11 @@
 
 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
+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
@@ -80,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"
         )
 
@@ -112,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",
@@ -136,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/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]