diff --git a/pilot/connections/conn_spark.py b/pilot/connections/conn_spark.py index d89ca9399..5feabd8c1 100644 --- a/pilot/connections/conn_spark.py +++ b/pilot/connections/conn_spark.py @@ -100,7 +100,7 @@ class SparkConnect(BaseConnect): def get_charset(self): return "UTF-8" - def get_db_list(self): + def get_db_list(self, user_id: str = None): return ["default"] def get_db_names(self): diff --git a/pilot/connections/manages/connect_config_db.py b/pilot/connections/manages/connect_config_db.py index 7443d18ad..20c8148bf 100644 --- a/pilot/connections/manages/connect_config_db.py +++ b/pilot/connections/manages/connect_config_db.py @@ -18,6 +18,7 @@ class ConnectConfigEntity(Base): db_user = Column(String(255), nullable=True, comment="db user") db_pwd = Column(String(255), nullable=True, comment="db password") comment = Column(Text, nullable=True, comment="db comment") + user_id = Column(String(255), nullable=True, comment="user id") __table_args__ = ( UniqueConstraint("db_name", name="uk_db"), diff --git a/pilot/connections/manages/connect_storage_duckdb.py b/pilot/connections/manages/connect_storage_duckdb.py index bdf1602c1..1df4ffd6c 100644 --- a/pilot/connections/manages/connect_storage_duckdb.py +++ b/pilot/connections/manages/connect_storage_duckdb.py @@ -21,7 +21,7 @@ class DuckdbConnectConfig: if not result: # create config table self.connect.execute( - "CREATE TABLE connect_config (id integer primary key, db_name VARCHAR(100) UNIQUE, db_type VARCHAR(50), db_path VARCHAR(255) NULL, db_host VARCHAR(255) NULL, db_port INTEGER NULL, db_user VARCHAR(255) NULL, db_pwd VARCHAR(255) NULL, comment TEXT NULL)" + "CREATE TABLE connect_config (id integer primary key, db_name VARCHAR(100) UNIQUE, db_type VARCHAR(50), db_path VARCHAR(255) NULL, db_host VARCHAR(255) NULL, db_port INTEGER NULL, db_user VARCHAR(255) NULL, db_pwd VARCHAR(255) NULL, comment TEXT NULL, user_id VARCHAR(127) NULL)" ) self.connect.execute("CREATE SEQUENCE seq_id START 1;") @@ -34,12 +34,13 @@ class DuckdbConnectConfig: db_user: str, db_pwd: str, comment: str = "", + user_id: str = "", ): try: cursor = self.connect.cursor() cursor.execute( - "INSERT INTO connect_config(id, db_name, db_type, db_path, db_host, db_port, db_user, db_pwd, comment)VALUES(nextval('seq_id'),?,?,?,?,?,?,?,?)", - [db_name, db_type, "", db_host, db_port, db_user, db_pwd, comment], + "INSERT INTO connect_config(id, db_name, db_type, db_path, db_host, db_port, db_user, db_pwd, comment, user_id)VALUES(nextval('seq_id'),?,?,?,?,?,?,?,?,?)", + [db_name, db_type, "", db_host, db_port, db_user, db_pwd, comment, user_id], ) cursor.commit() self.connect.commit() @@ -120,10 +121,10 @@ class DuckdbConnectConfig: return row_dict return None - def get_db_list(self): + def get_db_list(self, user_id: str = None): if os.path.isfile(duckdb_path): cursor = duckdb.connect(duckdb_path).cursor() - cursor.execute("SELECT * FROM connect_config ") + cursor.execute("SELECT * FROM connect_config where user_id =? or user_id=? or user_id IS NULL", [user_id, '']) fields = [field[0] for field in cursor.description] data = [] diff --git a/pilot/connections/manages/connection_manager.py b/pilot/connections/manages/connection_manager.py index 340c58f0b..1ef212ed9 100644 --- a/pilot/connections/manages/connection_manager.py +++ b/pilot/connections/manages/connection_manager.py @@ -54,7 +54,7 @@ class ConnectManager: def __init__(self, system_app: SystemApp): self.storage = DuckdbConnectConfig() self.db_summary_client = DBSummaryClient(system_app) - # self.__load_config_db() + self.__load_config_db() def __load_config_db(self): if CFG.LOCAL_DB_HOST: @@ -173,8 +173,8 @@ class ConnectManager: print(f"{db_info.db_name} Test connect Failure!{str(e)}") raise ValueError(f"{db_info.db_name} Test connect Failure!{str(e)}") - def get_db_list(self): - return self.storage.get_db_list() + def get_db_list(self, user_id: str = None): + return self.storage.get_db_list(user_id=user_id) def get_db_names(self): return self.storage.get_db_names() @@ -198,7 +198,7 @@ class ConnectManager: # 在这里执行需要异步运行的代码 self.db_summary_client.db_summary_embedding(db_name, db_type) - def add_db(self, db_info: DBConfig): + def add_db(self, db_info: DBConfig, user_id: str = None): print(f"add_db:{db_info.__dict__}") try: db_type = DBType.of_db_type(db_info.db_type) @@ -215,6 +215,7 @@ class ConnectManager: db_info.db_user, db_info.db_pwd, db_info.comment, + user_id, ) # async embedding executor = CFG.SYSTEM_APP.get_component( diff --git a/pilot/memory/chat_history/store_type/duckdb_history.py b/pilot/memory/chat_history/store_type/duckdb_history.py index 28c92a142..e716bcb12 100644 --- a/pilot/memory/chat_history/store_type/duckdb_history.py +++ b/pilot/memory/chat_history/store_type/duckdb_history.py @@ -69,7 +69,7 @@ class DuckdbHistoryMemory(BaseChatHistoryMemory): except Exception as e: print("init create conversation log error!" + str(e)) - def append(self, once_message: OnceConversation) -> None: + def append(self, once_message: OnceConversation, user_id: str = None) -> None: context = self.__get_messages_by_conv_uid(self.chat_seesion_id) conversations: List[OnceConversation] = [] if context: @@ -88,7 +88,7 @@ class DuckdbHistoryMemory(BaseChatHistoryMemory): self.chat_seesion_id, once_message.chat_mode, once_message.get_user_conv().content, - "", + user_id, json.dumps(conversations, ensure_ascii=False), ], ) @@ -149,7 +149,7 @@ class DuckdbHistoryMemory(BaseChatHistoryMemory): return None @staticmethod - def conv_list(cls, user_name: str = None) -> None: + def conv_list(user_name: str = None) -> None: if os.path.isfile(duckdb_path): cursor = duckdb.connect(duckdb_path).cursor() if user_name: diff --git a/pilot/openapi/api_v1/api_v1.py b/pilot/openapi/api_v1/api_v1.py index ea569b3e4..790b9d1c4 100644 --- a/pilot/openapi/api_v1/api_v1.py +++ b/pilot/openapi/api_v1/api_v1.py @@ -4,6 +4,8 @@ import asyncio import os import shutil import logging + +import requests from fastapi import ( APIRouter, Request, @@ -38,7 +40,7 @@ from pilot.server.knowledge.request.request import KnowledgeSpaceRequest from pilot.scene.base_chat import BaseChat from pilot.scene.base import ChatScene from pilot.scene.chat_factory import ChatFactory -from pilot.common.schema import DBType + from pilot.scene.message import OnceConversation from pilot.configs.model_config import LLM_MODEL_CONFIG, KNOWLEDGE_UPLOAD_ROOT_PATH @@ -46,6 +48,7 @@ from pilot.summary.db_summary_client import DBSummaryClient from pilot.memory.chat_history.chat_hisotry_factory import ChatHistory from pilot.model.cluster import BaseModelController, WorkerManager, WorkerManagerFactory from pilot.model.base import FlatSupportedModel +from pilot.user import UserDao, UserRequest, get_user_from_headers router = APIRouter() CFG = Config() @@ -56,6 +59,9 @@ knowledge_service = KnowledgeService() model_semaphore = None global_counter = 0 +GITHUB_CLIENT_ID = "a7353895f4f0821801d9" +GITHUB_CLIENT_SECRET = "14aff2bd2253841faab4c8efab9debeffbbff03e" + def __get_conv_user_message(conversations: dict): messages = conversations["messages"] @@ -67,11 +73,11 @@ def __get_conv_user_message(conversations: dict): def __new_conversation(chat_mode, user_id) -> ConversationVo: unique_id = uuid.uuid1() - return ConversationVo(conv_uid=str(unique_id), chat_mode=chat_mode) + return ConversationVo(conv_uid=str(unique_id), chat_mode=chat_mode, user_id=user_id, user_name=user_id) -def get_db_list(): - dbs = CFG.LOCAL_DB_MANAGE.get_db_list() +def get_db_list(user_id: str = None): + dbs = CFG.LOCAL_DB_MANAGE.get_db_list(user_id=user_id) params: dict = {} for item in dbs: params.update({item["db_name"]: item["db_name"]}) @@ -85,8 +91,8 @@ def plugins_select_info(): return plugins_infos -def get_db_list_info(): - dbs = CFG.LOCAL_DB_MANAGE.get_db_list() +def get_db_list_info(user_id: str = None): + dbs = CFG.LOCAL_DB_MANAGE.get_db_list(user_id=user_id) params: dict = {} for item in dbs: comment = item["comment"] @@ -130,13 +136,13 @@ def get_worker_manager() -> WorkerManager: @router.get("/v1/chat/db/list", response_model=Result[DBConfig]) -async def db_connect_list(): - return Result.succ(CFG.LOCAL_DB_MANAGE.get_db_list()) +async def db_connect_list(user_token: UserRequest = Depends(get_user_from_headers)): + return Result.succ(CFG.LOCAL_DB_MANAGE.get_db_list(user_token.user_id)) @router.post("/v1/chat/db/add", response_model=Result[bool]) -async def db_connect_add(db_config: DBConfig = Body()): - return Result.succ(CFG.LOCAL_DB_MANAGE.add_db(db_config)) +async def db_connect_add(db_config: DBConfig = Body(), user_token: UserRequest = Depends(get_user_from_headers)): + return Result.succ(CFG.LOCAL_DB_MANAGE.add_db(db_config, user_token.user_id)) @router.post("/v1/chat/db/edit", response_model=Result[bool]) @@ -182,10 +188,10 @@ async def db_support_types(): @router.get("/v1/chat/dialogue/list", response_model=Result[ConversationVo]) -async def dialogue_list(user_id: str = None): +async def dialogue_list(user_token: UserRequest = Depends(get_user_from_headers)): dialogues: List = [] chat_history_service = ChatHistory() - datas = chat_history_service.get_store_cls().conv_list(user_id) + datas = chat_history_service.get_store_cls().conv_list(user_name=user_token.user_id) for item in datas: conv_uid = item.get("conv_uid") summary = item.get("summary") @@ -235,9 +241,9 @@ async def dialogue_scenes(): @router.post("/v1/chat/dialogue/new", response_model=Result[ConversationVo]) async def dialogue_new( - chat_mode: str = ChatScene.ChatNormal.value(), user_id: str = None + chat_mode: str = ChatScene.ChatNormal.value(), user_token: UserRequest = Depends(get_user_from_headers) ): - conv_vo = __new_conversation(chat_mode, user_id) + conv_vo = __new_conversation(chat_mode, user_token.user_id) return Result.succ(conv_vo) @@ -259,7 +265,7 @@ async def params_list(chat_mode: str = ChatScene.ChatNormal.value()): @router.post("/v1/chat/mode/params/file/load") async def params_load( - conv_uid: str, chat_mode: str, model_name: str, doc_file: UploadFile = File(...) + conv_uid: str, chat_mode: str, model_name: str, doc_file: UploadFile = File(...), user_token: UserRequest = Depends(get_user_from_headers) ): print(f"params_load: {conv_uid},{chat_mode},{model_name}") try: @@ -285,7 +291,7 @@ async def params_load( select_param=doc_file.filename, model_name=model_name, ) - chat: BaseChat = get_chat_instance(dialogue) + chat: BaseChat = get_chat_instance(dialogue, user_token.user_id) resp = await chat.prepare() ### refresh messages @@ -295,6 +301,20 @@ async def params_load( return Result.faild(code="E000X", msg=f"File Load Error {e}") +@router.post("/v1/user/add") +async def add_user(user_req: UserRequest): + print(f"add user: {user_req}") + user_dao = UserDao() + + results = user_dao.get_by_user_no_and_channel(user_req.user_no, user_req.user_channel) + if len(results) > 0: + return Result.succ(results[0]) + user = user_dao.add_user_if_not_exist(user_req) + if user is not None: + return Result.succ(user) + return Result.faild(msg=f"user(channel={user_req.user_channel}, user_no={user_req.user_no}) is not valid!") + + @router.post("/v1/chat/dialogue/delete") async def dialogue_delete(con_uid: str): history_fac = ChatHistory() @@ -327,7 +347,7 @@ async def dialogue_history_messages(con_uid: str): return Result.succ(get_hist_messages(con_uid)) -def get_chat_instance(dialogue: ConversationVo = Body()) -> BaseChat: +def get_chat_instance(dialogue: ConversationVo = Body(), user_id: str = None) -> BaseChat: logger.info(f"get_chat_instance:{dialogue}") if not dialogue.chat_mode: dialogue.chat_mode = ChatScene.ChatNormal.value() @@ -345,6 +365,7 @@ def get_chat_instance(dialogue: ConversationVo = Body()) -> BaseChat: "current_user_input": dialogue.user_input, "select_param": dialogue.select_param, "model_name": dialogue.model_name, + "user_id": user_id, } chat: BaseChat = CHAT_FACTORY.get_implementation( dialogue.chat_mode, **{"chat_param": chat_param} @@ -353,11 +374,11 @@ def get_chat_instance(dialogue: ConversationVo = Body()) -> BaseChat: @router.post("/v1/chat/prepare") -async def chat_prepare(dialogue: ConversationVo = Body()): +async def chat_prepare(dialogue: ConversationVo = Body(), user_token: UserRequest = Depends(get_user_from_headers)): # dialogue.model_name = CFG.LLM_MODEL logger.info(f"chat_prepare:{dialogue}") ## check conv_uid - chat: BaseChat = get_chat_instance(dialogue) + chat: BaseChat = get_chat_instance(dialogue, user_token.user_id) if len(chat.history_message) > 0: return Result.succ(None) resp = await chat.prepare() @@ -365,11 +386,11 @@ async def chat_prepare(dialogue: ConversationVo = Body()): @router.post("/v1/chat/completions") -async def chat_completions(dialogue: ConversationVo = Body()): +async def chat_completions(dialogue: ConversationVo = Body(), user_token: UserRequest = Depends(get_user_from_headers)): print( f"chat_completions:{dialogue.chat_mode},{dialogue.select_param},{dialogue.model_name}" ) - chat: BaseChat = get_chat_instance(dialogue) + chat: BaseChat = get_chat_instance(dialogue, user_token.user_id) # background_tasks = BackgroundTasks() # background_tasks.add_task(release_model_semaphore) headers = { @@ -419,6 +440,36 @@ async def model_supports(worker_manager: WorkerManager = Depends(get_worker_mana return Result.faild(code="E000X", msg=f"Fetch supportd models error {e}") +@router.get("/v1/github/callback") +async def github_access_token(code: str = None): + try: + logger.info(f"github login callback code={code}") + access_token_url = f"https://github.com/login/oauth/access_token?client_id={GITHUB_CLIENT_ID}&client_secret={GITHUB_CLIENT_SECRET}&code={code}" + logger.info(f"github login access_token url={access_token_url}") + + headers = { + "Accept": "application/json", + } + result = requests.post(access_token_url, headers=headers) + resp_json = result.json() + logger.info(f"github access_token result={resp_json}") + + access_token = resp_json["access_token"] + get_user_header = { + "Authorization": f"Bearer {access_token}" + } + get_user_url = f"https://api.github.com/user" + user_info_resp = requests.get(get_user_url, headers=get_user_header) + logger.info(f"github user_info={user_info_resp.content}") + + user_info = json.loads(user_info_resp.content) + # save to db + return await add_user(UserRequest(user_no=str(user_info["login"]), user_name=user_info["name"], user_channel="github", nick_name=user_info["name"], role="normal", avatar_url=user_info["avatar_url"])) + except Exception as e: + logger.info(f"github login error: {e}") + return Result.faild(f"login error: {e}") + + async def no_stream_generator(chat): msg = await chat.nostream_call() yield f"data: {msg}\n\n" diff --git a/pilot/scene/base_chat.py b/pilot/scene/base_chat.py index 5d685573e..571706b88 100644 --- a/pilot/scene/base_chat.py +++ b/pilot/scene/base_chat.py @@ -80,6 +80,8 @@ class BaseChat(ABC): self.current_message.param_type = self.chat_mode.param_types()[0] self.current_message.param_value = chat_param["select_param"] self.current_tokens_used: int = 0 + if chat_param["user_id"]: + self.user_id = chat_param["user_id"] class Config: """Configuration for this pydantic object.""" @@ -187,7 +189,7 @@ class BaseChat(ABC): f"""ERROR!{str(e)}\n {ai_response_text} """ ) ### store current conversation - self.memory.append(self.current_message) + self.memory.append(self.current_message, self.user_id) async def nostream_call(self): payload = self.__call_base() @@ -233,7 +235,7 @@ class BaseChat(ABC): f"""ERROR!{str(e)}\n {ai_response_text} """ ) ### store dialogue - self.memory.append(self.current_message) + self.memory.append(self.current_message, self.user_id) return self.current_ai_response() def _blocking_stream_call(self): diff --git a/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py b/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py index 9a316cd34..12465998f 100644 --- a/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py +++ b/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py @@ -81,6 +81,7 @@ class ChatExcel(BaseChat): "select_param": self.excel_reader.excel_file_name, "excel_reader": self.excel_reader, "model_name": self.model_name, + "user_id": self.user_id, } learn_chat = ExcelLearning(**chat_param) result = await learn_chat.nostream_call() diff --git a/pilot/scene/chat_data/chat_excel/excel_learning/chat.py b/pilot/scene/chat_data/chat_excel/excel_learning/chat.py index c7eb82584..646704d32 100644 --- a/pilot/scene/chat_data/chat_excel/excel_learning/chat.py +++ b/pilot/scene/chat_data/chat_excel/excel_learning/chat.py @@ -28,6 +28,7 @@ class ExcelLearning(BaseChat): select_param: str = None, excel_reader: Any = None, model_name: str = None, + user_id: str = None, ): chat_mode = ChatScene.ExcelLearning """ """ @@ -39,6 +40,7 @@ class ExcelLearning(BaseChat): "current_user_input": user_input, "select_param": select_param, "model_name": model_name, + "user_id": user_id, } super().__init__(chat_param=chat_param) if parent_mode: diff --git a/pilot/scene/chat_knowledge/v1/chat.py b/pilot/scene/chat_knowledge/v1/chat.py index 9681f13c6..c77e9ce9a 100644 --- a/pilot/scene/chat_knowledge/v1/chat.py +++ b/pilot/scene/chat_knowledge/v1/chat.py @@ -114,4 +114,4 @@ class ChatKnowledge(BaseChat): def get_space_context(self, space_name): service = KnowledgeService() - return service.get_space_context(space_name) + return service.get_space_context(space_name, self.user_id) diff --git a/pilot/server/knowledge/api.py b/pilot/server/knowledge/api.py index 57fadb21e..e055722a2 100644 --- a/pilot/server/knowledge/api.py +++ b/pilot/server/knowledge/api.py @@ -3,7 +3,7 @@ import shutil import tempfile import logging -from fastapi import APIRouter, File, UploadFile, Form +from fastapi import APIRouter, File, UploadFile, Form, Depends from pilot.configs.config import Config from pilot.configs.model_config import ( @@ -27,6 +27,7 @@ from pilot.server.knowledge.request.request import ( ) from pilot.server.knowledge.request.request import KnowledgeSpaceRequest +from pilot.user import UserRequest, get_user_from_headers logger = logging.getLogger(__name__) @@ -38,9 +39,10 @@ knowledge_space_service = KnowledgeService() @router.post("/knowledge/space/add") -def space_add(request: KnowledgeSpaceRequest): +def space_add(request: KnowledgeSpaceRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/space/add params: {request}") try: + request.user_id = user_token.user_id knowledge_space_service.create_knowledge_space(request) return Result.succ([]) except Exception as e: @@ -48,16 +50,17 @@ def space_add(request: KnowledgeSpaceRequest): @router.post("/knowledge/space/list") -def space_list(request: KnowledgeSpaceRequest): +def space_list(request: KnowledgeSpaceRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/space/list params:") try: + request.user_id = user_token.user_id return Result.succ(knowledge_space_service.get_knowledge_space(request)) except Exception as e: return Result.faild(code="E000X", msg=f"space list error {e}") @router.post("/knowledge/space/delete") -def space_delete(request: KnowledgeSpaceRequest): +def space_delete(request: KnowledgeSpaceRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/space/delete params:") try: return Result.succ(knowledge_space_service.delete_space(request.name)) @@ -66,27 +69,27 @@ def space_delete(request: KnowledgeSpaceRequest): @router.post("/knowledge/{space_name}/arguments") -def arguments(space_name: str): +def arguments(space_name: str, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/knowledge/space/arguments params:") try: - return Result.succ(knowledge_space_service.arguments(space_name)) + return Result.succ(knowledge_space_service.arguments(space_name, user_token.user_id)) except Exception as e: return Result.faild(code="E000X", msg=f"space list error {e}") @router.post("/knowledge/{space_name}/argument/save") -def arguments_save(space_name: str, argument_request: SpaceArgumentRequest): +def arguments_save(space_name: str, argument_request: SpaceArgumentRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/knowledge/space/argument/save params:") try: return Result.succ( - knowledge_space_service.argument_save(space_name, argument_request) + knowledge_space_service.argument_save(space_name, argument_request, user_token.user_id) ) except Exception as e: return Result.faild(code="E000X", msg=f"space list error {e}") @router.post("/knowledge/{space_name}/document/add") -def document_add(space_name: str, request: KnowledgeDocumentRequest): +def document_add(space_name: str, request: KnowledgeDocumentRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/document/add params: {space_name}, {request}") try: return Result.succ( @@ -100,7 +103,7 @@ def document_add(space_name: str, request: KnowledgeDocumentRequest): @router.post("/knowledge/{space_name}/document/list") -def document_list(space_name: str, query_request: DocumentQueryRequest): +def document_list(space_name: str, query_request: DocumentQueryRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/document/list params: {space_name}, {query_request}") try: return Result.succ( @@ -111,7 +114,7 @@ def document_list(space_name: str, query_request: DocumentQueryRequest): @router.post("/knowledge/{space_name}/document/delete") -def document_delete(space_name: str, query_request: DocumentQueryRequest): +def document_delete(space_name: str, query_request: DocumentQueryRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/document/list params: {space_name}, {query_request}") try: return Result.succ( @@ -161,11 +164,11 @@ async def document_upload( @router.post("/knowledge/{space_name}/document/sync") -def document_sync(space_name: str, request: DocumentSyncRequest): +def document_sync(space_name: str, request: DocumentSyncRequest, user_token: UserRequest = Depends(get_user_from_headers)): logger.info(f"Received params: {space_name}, {request}") try: knowledge_space_service.sync_knowledge_document( - space_name=space_name, sync_request=request + space_name=space_name, sync_request=request, user_id=user_token.user_id, ) return Result.succ([]) except Exception as e: @@ -173,7 +176,7 @@ def document_sync(space_name: str, request: DocumentSyncRequest): @router.post("/knowledge/{space_name}/chunk/list") -def document_list(space_name: str, query_request: ChunkQueryRequest): +def document_list(space_name: str, query_request: ChunkQueryRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"/document/list params: {space_name}, {query_request}") try: return Result.succ(knowledge_space_service.get_document_chunks(query_request)) @@ -182,7 +185,7 @@ def document_list(space_name: str, query_request: ChunkQueryRequest): @router.post("/knowledge/{vector_name}/query") -def similar_query(space_name: str, query_request: KnowledgeQueryRequest): +def similar_query(space_name: str, query_request: KnowledgeQueryRequest, user_token: UserRequest = Depends(get_user_from_headers)): print(f"Received params: {space_name}, {query_request}") embedding_factory = CFG.SYSTEM_APP.get_component( "embedding_factory", EmbeddingFactory diff --git a/pilot/server/knowledge/request/request.py b/pilot/server/knowledge/request/request.py index b83165c19..5c3320609 100644 --- a/pilot/server/knowledge/request/request.py +++ b/pilot/server/knowledge/request/request.py @@ -22,6 +22,8 @@ class KnowledgeSpaceRequest(BaseModel): desc: str = None """owner: owner""" owner: str = None + """user_id: user_id""" + user_id: str = None class KnowledgeDocumentRequest(BaseModel): diff --git a/pilot/server/knowledge/service.py b/pilot/server/knowledge/service.py index 701246195..f4b8305ca 100644 --- a/pilot/server/knowledge/service.py +++ b/pilot/server/knowledge/service.py @@ -73,7 +73,7 @@ class KnowledgeService: - request: KnowledgeSpaceRequest """ query = KnowledgeSpaceEntity( - name=request.name, + name=request.name, user_id=request.user_id ) spaces = knowledge_space_dao.get_knowledge_space(query) if len(spaces) > 0: @@ -108,7 +108,7 @@ class KnowledgeService: - request: KnowledgeSpaceRequest """ query = KnowledgeSpaceEntity( - name=request.name, vector_type=request.vector_type, owner=request.owner + name=request.name, vector_type=request.vector_type, owner=request.owner, user_id=request.user_id ) responses = [] spaces = knowledge_space_dao.get_knowledge_space(query) @@ -128,12 +128,12 @@ class KnowledgeService: responses.append(res) return responses - def arguments(self, space_name): + def arguments(self, space_name: str, user_id: str): """show knowledge space arguments Args: - space_name: Knowledge Space Name """ - query = KnowledgeSpaceEntity(name=space_name) + query = KnowledgeSpaceEntity(name=space_name, user_id=user_id) spaces = knowledge_space_dao.get_knowledge_space(query) if len(spaces) != 1: raise Exception(f"there are no or more than one space called {space_name}") @@ -144,13 +144,13 @@ class KnowledgeService: context = space.context return json.loads(context) - def argument_save(self, space_name, argument_request: SpaceArgumentRequest): + def argument_save(self, space_name, argument_request: SpaceArgumentRequest, user_id: str): """save argument Args: - space_name: Knowledge Space Name - argument_request: SpaceArgumentRequest """ - query = KnowledgeSpaceEntity(name=space_name) + query = KnowledgeSpaceEntity(name=space_name, user_id=user_id) spaces = knowledge_space_dao.get_knowledge_space(query) if len(spaces) != 1: raise Exception(f"there are no or more than one space called {space_name}") @@ -178,7 +178,7 @@ class KnowledgeService: res.page = request.page return res - def sync_knowledge_document(self, space_name, sync_request: DocumentSyncRequest): + def sync_knowledge_document(self, space_name, sync_request: DocumentSyncRequest, user_id: str = None): """sync knowledge document chunk into vector store Args: - space: Knowledge Space Name @@ -209,7 +209,7 @@ class KnowledgeService: f" doc:{doc.doc_name} status is {doc.status}, can not sync" ) - space_context = self.get_space_context(space_name) + space_context = self.get_space_context(space_name, user_id) chunk_size = ( CFG.KNOWLEDGE_CHUNK_SIZE if space_context is None @@ -431,13 +431,14 @@ class KnowledgeService: context_template_string = json.dumps(context_template, indent=4) return context_template_string - def get_space_context(self, space_name): + def get_space_context(self, space_name, user_id: str = None): """get space contect Args: - space_name: space name """ request = KnowledgeSpaceRequest() request.name = space_name + request.user_id = user_id spaces = self.get_knowledge_space(request) if len(spaces) != 1: raise Exception( diff --git a/pilot/server/knowledge/space_db.py b/pilot/server/knowledge/space_db.py index b833906da..3f8ad957d 100644 --- a/pilot/server/knowledge/space_db.py +++ b/pilot/server/knowledge/space_db.py @@ -24,9 +24,10 @@ class KnowledgeSpaceEntity(Base): context = Column(Text) gmt_created = Column(DateTime) gmt_modified = Column(DateTime) + user_id = Column(String(100)) def __repr__(self): - return f"KnowledgeSpaceEntity(id={self.id}, name='{self.name}', vector_type='{self.vector_type}', desc='{self.desc}', owner='{self.owner}' context='{self.context}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" + return f"KnowledgeSpaceEntity(id={self.id}, name='{self.name}', vector_type='{self.vector_type}', desc='{self.desc}', owner='{self.owner}' context='{self.context}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}', user_id='{self.user_id}')" class KnowledgeSpaceDao(BaseDao): @@ -44,6 +45,7 @@ class KnowledgeSpaceDao(BaseDao): owner=space.owner, gmt_created=datetime.now(), gmt_modified=datetime.now(), + user_id=space.user_id, ) session.add(knowledge_space) session.commit() @@ -52,6 +54,10 @@ class KnowledgeSpaceDao(BaseDao): def get_knowledge_space(self, query: KnowledgeSpaceEntity): session = self.get_session() knowledge_spaces = session.query(KnowledgeSpaceEntity) + if query.user_id is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.user_id == query.user_id + ) if query.id is not None: knowledge_spaces = knowledge_spaces.filter( KnowledgeSpaceEntity.id == query.id diff --git a/pilot/user/__init__.py b/pilot/user/__init__.py new file mode 100644 index 000000000..79cef7567 --- /dev/null +++ b/pilot/user/__init__.py @@ -0,0 +1,13 @@ +from pilot.user.user_db import ( + UserEntity, + UserDao +) + +from pilot.user.user_request import UserRequest, get_user_from_headers + +__all__ = [ + "UserEntity", + "UserDao", + "UserRequest", + "get_user_from_headers", +] diff --git a/pilot/user/meta_data/dbgpt.db b/pilot/user/meta_data/dbgpt.db new file mode 100644 index 000000000..e69de29bb diff --git a/pilot/user/user_db.py b/pilot/user/user_db.py new file mode 100644 index 000000000..9e969aa85 --- /dev/null +++ b/pilot/user/user_db.py @@ -0,0 +1,146 @@ +from pilot.configs.config import Config +from sqlalchemy.ext.declarative import declarative_base +from datetime import datetime + +from sqlalchemy import Column, Integer, String, DateTime +from pilot.connections.rdbms.base_dao import BaseDao +from pilot.user.user_request import UserRequest +import uuid + +CFG = Config() +Base = declarative_base() + +DB_AUTH = "auth" +TABLE_USER = "user" + + +class UserEntity(Base): + __tablename__ = TABLE_USER + user_id = Column(String(36), primary_key=True) + nick_name = Column(String(127)) + user_no = Column(String(127)) + user_channel = Column(String(127)) + role = Column(String(100)) + email = Column(String(100)) + gmt_created = Column(DateTime) + gmt_modified = Column(DateTime) + avatar_url = Column(String(256)) + + def __repr__(self): + return f"UserEntity(user_id='{self.user_id}', nick_name='{self.nick_name}', user_no='{self.user_no}', user_channel='{self.user_channel}', role='{self.role}', email='{self.email}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" + + +class UserDao(BaseDao): + def __init__(self): + """ + create database `auth` firstly. + """ + super().__init__( + database=DB_AUTH, orm_base=Base, create_not_exist_table=False + ) + + def get_by_user_id(self, user_id: str): + """ + get user info by user_id + """ + session = self.Session() + user_query = session.query(UserEntity).filter(UserEntity.user_id == user_id) + result = user_query.all() + session.close() + + if len(result) > 0: + return result[0] + raise f"user({user_id}) is not existed!" + + def get_by_user_no_and_channel(self, user_no: str, user_channel: str): + """ + query user by user_no + """ + session = self.Session() + user_query = session.query(UserEntity).filter( + UserEntity.user_no == user_no + ).filter( + UserEntity.user_channel == user_channel + ) + + result = user_query.all() + session.close() + return result + + def add_user_if_not_exist(self, user_req: UserRequest): + """ + check before do addition, add new user. + """ + if user_req is None: + raise f"user info is empty!" + + # new user, by user_channel and user_no: + if user_req.user_no is not None and user_req.user_channel is not None: + result = self.get_by_user_no_and_channel(user_req.user_no, user_req.user_channel) + if len(result) == 0: + session = self.Session() + user = UserEntity( + user_id=uuid.uuid4().hex, + nick_name=user_req.nick_name, + user_no=user_req.user_no, + user_channel=user_req.user_channel, + role=user_req.role, + email=user_req.email, + gmt_created=datetime.now(), + gmt_modified=datetime.now(), + avatar_url=user_req.avatar_url + ) + session.add(user) + session.commit() + session.close() + return user + return None + + def get_all_users(self): + """ + query user by user_no + """ + session = self.Session() + user_query = session.query(UserEntity).order_by( + UserEntity.gmt_created.desc() + ) + result = user_query.all() + session.close() + return result + + def update_user(self, user: UserEntity): + session = self.Session() + session.merge(user) + session.commit() + session.close() + return True + + def delete_user(self, user: UserEntity): + session = self.Session() + if user: + session.delete(user) + session.commit() + session.close() + + +if __name__ == "__main__": + user_dao = UserDao() + + result = user_dao.get_all_users() + print(result) + + user_dao.add_user_if_not_exist(UserRequest(user_no="101", nick_name="shine", user_channel="GITHUB")) + result2 = user_dao.get_by_user_no_and_channel("101", "GITHUB") + print(result2) + + result3 = user_dao.get_all_users() + print(result3) + + user_dao.delete_user(result2[0]) + result4 = user_dao.get_all_users() + print(result4) + + + user: UserEntity = result3[0] + ret5 = user_dao.get_by_user_id(user.user_id) + print(ret5) diff --git a/pilot/user/user_request.py b/pilot/user/user_request.py new file mode 100644 index 000000000..870fcfe55 --- /dev/null +++ b/pilot/user/user_request.py @@ -0,0 +1,22 @@ +import json + +from fastapi import Header +from pydantic import BaseModel + + +class UserRequest(BaseModel): + user_id: str = None + user_no: str = None + user_name: str = None # 同user_id + user_channel: str = None + role: str = "normal" + nick_name: str = None + email: str = None + avatar_url: str = None + + +def get_user_from_headers(user_id: str = Header(...)): + try: + return UserRequest(user_id=user_id) + except Exception as e: + raise f"Authentication failed."