diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 459dd050..9977eae7 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -18,9 +18,9 @@ async def chats(session: SessionDep, current_user: CurrentUser): @router.get("/get/{chart_id}") -async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int): +async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int, current_assistant: CurrentAssistant): try: - return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user) + return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user, current_assistant=current_assistant) except Exception as e: raise HTTPException( status_code=500, diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index d90ad855..220ad6d0 100644 --- a/backend/apps/chat/curd/chat.py +++ b/backend/apps/chat/curd/chat.py @@ -8,7 +8,8 @@ from sqlalchemy.orm import load_only 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 +from apps.system.crud.assistant import AssistantOutDsFactory +from common.core.deps import CurrentAssistant, SessionDep, CurrentUser from common.utils.utils import extract_nested_json @@ -65,14 +66,19 @@ def get_chat_predict_data(session: SessionDep, chart_record_id: int): return '' -def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> ChatInfo: - chat = session.query(Chat).filter(Chat.id == chart_id).first() +def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser, current_assistant: CurrentAssistant) -> ChatInfo: + chat = session.get(Chat, chart_id) if not chat: raise Exception(f"Chat with id {chart_id} not found") chat_info = ChatInfo(**chat.model_dump()) - ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first() + if current_assistant and current_assistant.type == 1: + out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant) + ds = out_ds_instance.get_ds(chat.datasource) + else: + ds = session.get(CoreDatasource, chat.datasource) if chat.datasource else None + if not ds: chat_info.datasource_exists = False chat_info.datasource_name = 'Datasource not exist' diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 4194e75f..b4786662 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -31,7 +31,8 @@ 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 AssistantOutDs, get_assistant_ds +from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds +from apps.system.schemas.system_schema import AssistantOutDsSchema from common.core.config import settings from common.core.deps import CurrentAssistant, SessionDep, CurrentUser from common.utils.utils import extract_nested_json @@ -67,15 +68,23 @@ class LLMService: chat: Chat = self.session.get(Chat, chat_id) if not chat: raise Exception(f"Chat with id {chat_id} not found") - ds: CoreDatasource | None = None + ds: CoreDatasource | AssistantOutDsSchema | None = None if chat.datasource: # Get available datasource # ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first() - ds = self.session.get(CoreDatasource, chat.datasource) - if not ds: - raise Exception("No available datasource configuration found") - - chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL' + if current_assistant and current_assistant.type == 1: + self.out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant) + ds = self.out_ds_instance.get_ds(chat.datasource) + if not ds: + raise Exception("No available datasource configuration found") + chat_question.engine = ds.type + chat_question.db_schema = self.out_ds_instance.get_db_schema(ds.id) + else: + ds = self.session.get(CoreDatasource, chat.datasource) + if not ds: + raise Exception("No available datasource configuration found") + chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL' + chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds) history_records: List[ChatRecord] = list( map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False, @@ -84,13 +93,9 @@ class LLMService: chart_id=chat_id)))) self.change_title = len(history_records) == 0 - # get schema - if ds: - chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds) - chat_question.lang = current_user.language - self.ds = CoreDatasource(**ds.model_dump()) if ds else None + self.ds = (ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None self.chat_question = chat_question self.config = get_default_config() self.chat_question.ai_modal_id = self.config.model_id @@ -298,7 +303,7 @@ class LLMService: # get schema if self.ds and not self.chat_question.db_schema: - self.chat_question.db_schema = get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds) + self.chat_question.db_schema = self.out_ds_instance.get_db_schema(self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds) guess_msg: List[Union[BaseMessage, dict[str, Any]]] = [] guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question())) diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 45f451e1..fddf2247 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -15,6 +15,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 common.utils.utils import string_to_numeric_hash @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id") async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None: @@ -100,7 +101,7 @@ class AssistantOutDs: header = {} cookies = {} for item in certificateList: - if item['target'] == 'head': + if item['target'] == 'header': header[item['key']] = item['value'] if item['target'] == 'cookie': cookies[item['key']] = item['value'] @@ -111,8 +112,8 @@ class AssistantOutDs: if result_json.get('code') == 0: temp_list = result_json.get('data', []) self.ds_list = [ - AssistantOutDsSchema(**{**item, "id": idx}) - for idx, item in enumerate(temp_list, start=1) + self.convert2schema(item) + for item in temp_list ] return self.ds_list @@ -151,6 +152,14 @@ class AssistantOutDs: raise Exception("Datasource list is not found.") raise Exception(f"Datasource with id {ds_id} not found.") + def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema: + id_marker: str = '' + attr_list = ['name', 'type', 'host', 'port', 'user', 'dataBase', 'schema'] + for attr in attr_list: + if attr in ds_dict: + id_marker += str(ds_dict.get(attr, '')) + '--sqlbot--' + id = string_to_numeric_hash(id_marker) + return AssistantOutDsSchema(**{**ds_dict, "id": id}) class AssistantOutDsFactory: @staticmethod diff --git a/backend/common/utils/utils.py b/backend/common/utils/utils.py index e069b9de..f690b27f 100644 --- a/backend/common/utils/utils.py +++ b/backend/common/utils/utils.py @@ -1,5 +1,7 @@ +import hashlib import logging from datetime import datetime, timedelta, timezone +from typing import Optional import jwt import orjson @@ -70,3 +72,7 @@ def extract_nested_json(text): if len(results) > 0 and results[0]: return results[0] return None + +def string_to_numeric_hash(text: str, bits: Optional[int] = 64) -> int: + hash_bytes = hashlib.sha256(text.encode()).digest() + return int.from_bytes(hash_bytes[:bits//8], byteorder='big')