diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index c8b4948b..d90ad855 100644 --- a/backend/apps/chat/curd/chat.py +++ b/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 diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 3d43fbcd..199fae5a 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/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): diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index f17f601a..00e10889 100644 --- a/backend/apps/chat/task/llm.py +++ b/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() diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 6756b92e..de46446e 100644 --- a/backend/apps/system/crud/assistant.py +++ b/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="sqlbot-inner-assistant@sqlbot.com") -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) \ No newline at end of file diff --git a/backend/apps/system/schemas/auth.py b/backend/apps/system/schemas/auth.py index 70b87d74..f67b15f0 100644 --- a/backend/apps/system/schemas/auth.py +++ b/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 \ No newline at end of file