perf: Assistant dynamic ds

Author: fit2cloud-chenyw

backend/apps/chat/task/llm.py

@@ -748,7 +748,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
                 yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
                                     'engine_type': llm_service.ds.type_name or llm_service.ds.type, 'type': 'datasource'}).decode() + '\n\n'
 
-            llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema() if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
+            llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema(llm_service.ds.id) if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
 
         # generate sql
         sql_res = llm_service.generate_sql()

backend/apps/db/db.py

@@ -11,26 +11,32 @@
 from apps.db.engine import get_engine_config
 from decimal import Decimal
 
+from apps.system.crud.assistant import get_ds_engine
+from apps.system.schemas.system_schema import AssistantOutDsSchema
 
-def get_uri(ds: CoreDatasource):
+
+def get_uri(ds: CoreDatasource) -> str:
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
+    return get_uri_from_config(ds.type, conf)
+
+def get_uri_from_config(type: str,conf: DatasourceConf) -> str:
     db_url: str
-    if ds.type == "mysql":
+    if type == "mysql":
         if conf.extraJdbc is not None and conf.extraJdbc != '':
             db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
         else:
             db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
-    elif ds.type == "sqlServer":
+    elif type == "sqlServer":
         if conf.extraJdbc is not None and conf.extraJdbc != '':
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
         else:
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
-    elif ds.type == "pg" or ds.type == "excel":
+    elif type == "pg" or type == "excel":
         if conf.extraJdbc is not None and conf.extraJdbc != '':
             db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
         else:
             db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
-    elif ds.type == "oracle":
+    elif type == "oracle":
         if conf.mode == "service_name":
             if conf.extraJdbc is not None and conf.extraJdbc != '':
                 db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}?service_name={urllib.parse.quote(conf.database)}&{conf.extraJdbc}"
@@ -62,8 +68,8 @@ def get_engine(ds: CoreDatasource) -> Engine:
     return engine
 
 
-def get_session(ds: CoreDatasource):
-    engine = get_engine(ds)
+def get_session(ds: CoreDatasource | AssistantOutDsSchema):
+    engine = get_engine(ds) if isinstance(ds, CoreDatasource) else get_ds_engine(ds)
     session_maker = sessionmaker(bind=engine)
     session = session_maker()
     return session
@@ -238,7 +244,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
             session.close()
 
 
-def exec_sql(ds: CoreDatasource, sql: str):
+def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str):
     session = get_session(ds)
     result = session.execute(text(sql))
     try:

backend/apps/system/crud/assistant.py

@@ -4,8 +4,10 @@
 from typing import Optional
 from fastapi import FastAPI
 import requests
+from sqlalchemy import Engine, create_engine
 from sqlmodel import Session, select
-from apps.datasource.models.datasource import CoreDatasource
+from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
+
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import AssistantOutDsSchema, UserInfoDTO
@@ -92,7 +94,7 @@ def __init__(self, assistant: AssistantModel, certificate: Optional[str] = None)
         self.get_ds_from_api(certificate)
 
     #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")    
-    async def get_ds_from_api(self, certificate: Optional[str] = None):
+    def get_ds_from_api(self, certificate: Optional[str] = None):
         config: dict[any] = json.loads(self.assistant.configuration)
         endpoint: str = config['endpoint']
         certificateList: list[any] = json.loads(certificate)
@@ -106,7 +108,7 @@ async def get_ds_from_api(self, certificate: Optional[str] = None):
 
         res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10)
         if res.status_code == 200:
-            result_json: dict[any] = json.loads(res.json())
+            result_json: dict[any] = json.loads(res.text)
             if result_json.get('code') == 0:
                 temp_list = result_json.get('data', [])
                 self.ds_list = [
@@ -144,20 +146,35 @@ def get_db_schema(self, ds_id: int) -> str:
     def get_ds(self, ds_id: int):
         if self.ds_list:
             for ds in self.ds_list:
-                if ds['id'] == ds_id:
+                if ds.id == ds_id:
                     return ds
         else:
             raise Exception("Datasource list is not found.")
         raise Exception(f"Datasource with id {ds_id} not found.")
-    def get_ds_engine(self, ds_id: int):
-        ds = self.get_ds(ds_id)
-        ds_type =  ds.get('type') if ds else None
-        if not ds_type:
-            raise Exception(f"Datasource with id {ds_id} not found or type is not defined.")
-        return ds_type
+
 
 class AssistantOutDsFactory:
     @staticmethod
     def get_instance(assistant: AssistantModel, certificate: Optional[str] = None) -> AssistantOutDs:
         return AssistantOutDs(assistant, certificate)
+
+def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
+    timeout: int = 30
+    connect_args = {"connect_timeout": timeout}
+    conf = DatasourceConf(
+        host=ds.host, 
+        port=ds.port, 
+        username=ds.user,
+        password=ds.password,   
+        database=ds.dataBase,
+        driver='',
+        extraJdbc=ds.extraParams,
+        dbSchema=ds.schema or ''
+    )
+    from apps.db.db import get_uri_from_config
+    uri = get_uri_from_config(ds.type, conf)
+    if ds.type == "pg" and ds.schema:
+        connect_args.update({"options": f"-c search_path={ds.schema}"})
+    engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10)
+    return engine
 

backend/apps/system/schemas/system_schema.py

@@ -125,29 +125,14 @@ class AssistantOutDsBase(BaseModel):
     comment: Optional[str] = None
     description: Optional[str] = None
 
-    def __init__(self, id: Optional[int] = None, name: str = '', type: Optional[str] = None,
-                 type_name: Optional[str] = None, comment: Optional[str] = None):
-        super().__init__(id=id, name=name, type=type, type_name=type_name, comment=comment)
 
 class AssistantOutDsSchema(AssistantOutDsBase):
     host: Optional[str] = None
     port: Optional[int] = None
+    dataBase: Optional[str] = None
     user: Optional[str] = None
     password: Optional[str] = None
     schema: Optional[str] = None
+    extraParams: Optional[str] = None
     tables: Optional[list[AssistantTableSchema]] = None
-
-    def __init__(self, id: int, name: str, comment: Optional[str] = None, type: Optional[str] = None,
-                 type_name: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None,
-                 username: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None):
-        self.id = id
-        self.name = name
-        self.comment = comment
-        self.type = type
-        self.type_name = type_name
-        self.host = host
-        self.port = port
-        self.username = username
-        self.password = password
-        self.schema = schema
-        self.description = comment
+    

backend/common/core/deps.py

@@ -19,8 +19,8 @@ async def get_current_user(request: Request) -> UserInfoDTO:
 
 CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)]
 
-async def get_current_assistant(request: Request) -> AssistantModel:
-    return request.state.assistant
+async def get_current_assistant(request: Request) -> AssistantModel | None:
+    return request.state.assistant if hasattr(request.state, "assistant") else None
 
 CurrentAssistant = Annotated[AssistantModel, Depends(get_current_assistant)]
 

frontend/src/views/chat/answer/ChartAnswer.vue

@@ -1,7 +1,9 @@
 <script setup lang="ts">
 import BaseAnswer from './BaseAnswer.vue'
 import { Chat, chatApi, ChatInfo, type ChatMessage, ChatRecord, questionApi } from '@/api/chat.ts'
+import { useAssistantStore } from '@/stores/assistant'
 import { computed, nextTick, ref } from 'vue'
+const assistantStore = useAssistantStore()
 const props = withDefaults(
   defineProps<{
     chatList?: Array<ChatInfo>
@@ -96,11 +98,15 @@ const sendMessage = async () => {
 
   try {
     const controller: AbortController = new AbortController()
-    const response = await questionApi.add({
+    const param = {
       question: currentRecord.question,
       chat_id: _currentChatId.value,
+      assistant_certificate: assistantStore.getCertificate,
       controller,
-    })
+    }
+    console.log(assistantStore.getCertificate)
+    console.log(param)
+    const response = await questionApi.add(param)
     const reader = response.body.getReader()
     const decoder = new TextDecoder()