diff --git a/app/model/KnowledgeDoc.py b/app/model/KnowledgeDoc.py index 176408f..435a26f 100644 --- a/app/model/KnowledgeDoc.py +++ b/app/model/KnowledgeDoc.py @@ -2,6 +2,7 @@ import asyncio from enum import Enum from typing import Optional +from llama_index.core import Document from pydantic import computed_field from sqlalchemy.orm import selectinload from sqlmodel import Field, Relationship, select @@ -9,6 +10,7 @@ from sqlmodel import Field, Relationship, select from app.model.BasicModel import BasicModel from app.model.UserModel import UserServiceModel from app.utils.create_module_service import create_model_service +from app.utils.db_utils import AsyncSessionDep from app.utils.milvus_utils import milvus_service @@ -58,10 +60,12 @@ class KnowledgeDocModel(BasicModel, table=True): def creator(self, value: Optional[UserServiceModel]) -> None: pass + # 删除文档字后,要删除文档在Milvus中对应的文档块 async def handle_after_delete(delete_cls, row_dict, session): await milvus_service.async_delete(delete_cls.id) + # 批量删除文档字后,要删除文档在Milvus中对应的文档块 async def handle_after_batch_delete(delete_cls_list, row_dict_list, session): task_list = [ @@ -70,6 +74,7 @@ async def handle_after_batch_delete(delete_cls_list, row_dict_list, session): ] await asyncio.gather(*task_list) + # /*---------------------------------------KnowledgeDocService-------------------------------------------*/ # 专门给KnowledgeService使用的,用来给知识库上传文档的时候操作Mysql数据库中的doc表 @@ -79,10 +84,50 @@ KnowledgeDocService = create_model_service( after_batch_delete=handle_after_batch_delete ) + # /*---------------------------------------KnowledgeDocServiceWithCreator-------------------------------------------*/ + +# 新建文档之后,要给对应的文档创建向量索引 +async def handle_after_insert(insert_cls: KnowledgeDocModel, row_dict: dict, session: AsyncSessionDep): + await milvus_service.async_create_index_from_documents([Document( + doc_id=insert_cls.id, + text=insert_cls.content, + metadata=insert_cls.model_dump() + )]) + + +# 批量新建之后,给所有对应的文档创建向量索引 +async def handle_after_batch_insert(insert_cls_list: list[KnowledgeDocModel], row_dict_list: list[dict], session: AsyncSessionDep): + document_list = [ + Document( + doc_id=insert_cls.id, + text=insert_cls.content, + metadata=insert_cls.model_dump() + ) + for insert_cls in insert_cls_list + ] + await milvus_service.async_create_index_from_documents(document_list) + + +# 更新文档之后,先删除文档索引再重新创建 +async def handle_after_update(update_cls: KnowledgeDocModel, row_dict: dict, session: AsyncSessionDep): + await milvus_service.async_delete(update_cls.id) + await handle_after_insert(update_cls, row_dict, session) + + +# 批量更新文档之后,先删除文档索引再重新创建 +async def handle_after_batch_update(update_cls_list: list[KnowledgeDocModel], row_dict_list: list[dict], session: AsyncSessionDep): + await handle_after_batch_delete(update_cls_list, row_dict_list, session) + await handle_after_batch_insert(update_cls_list, row_dict_list, session) + + # 专门给在线文档模块使用的 KnowledgeDocServiceWithCreator = create_model_service( KnowledgeDocModel, + after_insert=handle_after_insert, + after_batch_insert=handle_after_batch_insert, + after_update=handle_after_update, + after_batch_update=handle_after_batch_update, after_delete=handle_after_delete, after_batch_delete=handle_after_batch_delete, custom_query=lambda: select(KnowledgeDocModel)