Merge remote-tracking branch 'origin/main'

backend/apps/chat/curd/chat.py

@@ -341,7 +341,7 @@ def save_full_select_datasource_message_and_answer(session: SessionDep, record_i
                                                    engine_type: str = None, token_usage: dict = None) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
-    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record = session.get(ChatRecord, record_id)
     record.full_select_datasource_message = full_message
     record.datasource_select_answer = answer
 

backend/apps/chat/models/chat_model.py

@@ -151,6 +151,7 @@ class AiModelQuestion(BaseModel):
 class ChatQuestion(AiModelQuestion):
     question: str = ''
     chat_id: int = 0
+    assistant_certificate: Optional[str] = None
 
 
 class ChatMcp(ChatQuestion):

backend/apps/chat/task/llm.py

@@ -30,7 +30,7 @@ from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.crud.row_permission import transFilterTree
 from apps.datasource.models.datasource import CoreDatasource, CoreTable
 from apps.db.db import exec_sql
-from apps.system.crud.assistant import get_assistant_ds
+from apps.system.crud.assistant import AssistantOutDs, get_assistant_ds
 from common.core.config import settings
 from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
 from common.utils.utils import extract_nested_json
@@ -52,6 +52,8 @@ class LLMService:
     session: SessionDep
     current_user: CurrentUser
     current_assistant: Optional[CurrentAssistant] = None
+    assistant_certificate: str
+    out_ds_instance: Optional[AssistantOutDs] = None
     change_title: bool = False
 
     def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question: ChatQuestion,
@@ -60,6 +62,8 @@ class LLMService:
         self.session = session
         self.current_user = current_user
         self.current_assistant = current_assistant
+        if chat_question.assistant_certificate:
+            self.assistant_certificate = chat_question.assistant_certificate
         # chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
         chat_id = chat_question.chat_id
         chat: Chat = self.session.get(Chat, chat_id)
@@ -346,7 +350,7 @@ class LLMService:
         datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
         datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
         if self.current_assistant:
-            _ds_list = get_assistant_ds(session=self.session, assistant=self.current_assistant)
+            _ds_list = get_assistant_ds(llm_service=self)
         else:
             _ds_list = self.session.exec(select(CoreDatasource).options(
                 load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all()
@@ -398,15 +402,19 @@ class LLMService:
 
             if data['id'] and data['id'] != 0:
                 _datasource = data['id']
-                _ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == _datasource).first()
-                if not _ds:
-                    _datasource = None
-                    raise Exception(f"Datasource configuration with id {_datasource} not found")
-                self.ds = CoreDatasource(**_ds.model_dump())
+                if self.current_assistant.type == 1:
+                    _ds = self.out_ds_instance.get_ds(data['id'])
+                    self.ds = _ds
+                else:
+                    _ds = self.session.get(CoreDatasource, _datasource)
+                    if not _ds:
+                        _datasource = None
+                        raise Exception(f"Datasource configuration with id {_datasource} not found")
+                    self.ds = CoreDatasource(**_ds.model_dump())
                 self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
                 _engine_type = self.chat_question.engine
                 # save chat
-                _chat = self.session.query(Chat).filter(Chat.id == self.record.chat_id).first()
+                _chat = self.session.get(Chat, self.record.chat_id)
                 _chat.datasource = _datasource
                 _chat.engine_type = _ds.type_name
 
@@ -727,7 +735,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, 'type': 'datasource'}).decode() + '\n\n'
 
-            llm_service.chat_question.db_schema = 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() 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/system/crud/assistant.py

@@ -1,8 +1,11 @@
 
 
 import json
+from typing import Optional
 from fastapi import FastAPI
+import requests
 from sqlmodel import Session, select
+from apps.chat.task.llm import LLMService
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
@@ -11,6 +14,7 @@ from common.core.sqlbot_cache import cache
 from common.core.db import engine
 from starlette.middleware.cors import CORSMiddleware
 from common.core.config import settings
+from deps import CurrentUser
 
 @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
 async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
@@ -20,21 +24,26 @@ async def get_assistant_info(*, session: Session, assistant_id: int) -> Assistan
 def get_assistant_user(*, id: int):
     return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="[email protected]")
 
-def get_assistant_ds(*, session: Session, assistant: AssistantModel):
+# def get_assistant_ds(*, session: Session, assistant: AssistantModel):
+def get_assistant_ds(llm_service: LLMService) -> list[dict]:
+    assistant: AssistantModel = llm_service.current_assistant
+    session: Session = llm_service.session
     type = assistant.type
     if type == 0:
-        stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description)
         configuration = assistant.configuration
         if configuration:
-            config = json.loads(configuration)
+            config: dict[any] = json.loads(configuration)
+            oid: str = config['oid']
+            stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
             private_list:list[int] = config['private_list']
-            if not private_list:
+            if private_list:
                 stmt.where(~CoreDatasource.id.in_(private_list))
         db_ds_list = session.exec(stmt).all()
         # filter private ds if offline
         return db_ds_list
-    out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant)
-    dslist = out_ds_instance.get_ds_list()
+    out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant, llm_service.assistant_certificate)
+    llm_service.out_ds_instance = out_ds_instance
+    dslist = out_ds_instance.get_simple_ds_list()
     # format?
     return dslist
 
@@ -66,16 +75,66 @@ def init_dynamic_cors(app: FastAPI):
 
 class AssistantOutDs:
     assistant: AssistantModel
-    def get_ds_list(self):
+    ds_list: Optional[list[dict]] = None
+    certificate: Optional[str] = None
+    def __init__(self, assistant: AssistantModel, certificate: Optional[str] = None):
+        self.assistant = assistant
+        self.ds_list = None
+        self.certificate = certificate
+        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):
         config: dict[any] = json.loads(self.assistant.configuration)
-        url: str = config['url']
+        endpoint: str = config['endpoint']
+        certificateList: list[any] = json.loads(certificate)
+        header = {}
+        cookies = {}
+        for item in certificateList:
+            if item['target'] == 'head':
+                header[item['key']] = item['value']
+            if item['target'] == 'cookie':
+                cookies[item['key']] = item['value']
+        
+        res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10)
+        if res.status_code == 200:
+            result_json: dict[any] = json.loads(res.json())
+            if result_json.get('code') == 0:
+                temp_list = result_json.get('data', [])
+                for idx, item in enumerate(temp_list, start=1):
+                    item["id"] = idx
+                self.ds_list = temp_list
+                return self.ds_list
+            else:
+                raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
+        else:
+            raise Exception(f"Failed to get datasource list from {endpoint}, status code: {res.status_code}")
+    
+    def get_simple_ds_list(self):
+        if self.ds_list:
+            return [{'id': ds['id'], 'name': ds['name'], 'description': ds['comment']} for ds in self.ds_list]
+        else:
+            raise Exception("Datasource list is not found.")
+       
+    def get_db_schema(self, ds_id: int):
         return None
+    def get_ds(self, ds_id: int):
+        if self.ds_list:
+            for ds in self.ds_list:
+                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:
-    _instance: AssistantOutDs = None
     @staticmethod
-    def get_instance(cls, assistant: AssistantModel) -> AssistantOutDs:
-        if not cls._instance:
-            cls._instance = AssistantOutDs(assistant)
-        return cls._instance
+    def get_instance(assistant: AssistantModel, certificate: Optional[str] = None) -> AssistantOutDs:
+        return AssistantOutDs(assistant, certificate)
     

backend/apps/system/schemas/auth.py

@@ -14,5 +14,6 @@ class CacheNamespace(Enum):
 class CacheName(Enum):
     USER_INFO = "user:info"
     ASSISTANT_INFO = "assistant:info"
+    ASSISTANT_DS = "assistant:ds"
     def __str__(self):
         return self.value