feat: KnowledgeDocServiceWithCreator, 新建或者更新完在线文档之后，更新文档在Milvus中的向量记录

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)