feat: milvus基本的增删改查

app/utils/milvus_utils.py

@@ -0,0 +1,91 @@
+from typing import List, Optional
+
+from llama_index.core import Document, VectorStoreIndex
+from llama_index.vector_stores.milvus import MilvusVectorStore
+
+from app.config.env import env
+from app.utils.llm_utils import create_embeddings, create_llama_index_llm
+
+
+class MilvusService:
+  def __init__(self):
+    milvus_vector_store: Optional[MilvusVectorStore] = None
+    self._milvus_vector_store = milvus_vector_store
+
+    self.embeddings = create_embeddings()
+
+  # 获取一个MilvusVectorStore实例，如果不存在就异步创建
+  async def get_vector_store(self):
+    if not self._milvus_vector_store:
+      print("create new_vector_store...")
+      self._milvus_vector_store = MilvusVectorStore(
+        uri=env.milvus_uri,
+        dim=env.milvus_dimension,
+        collection_name=env.milvus_collection,
+        db_name=env.milvus_database,
+        embedding_field="embedding",
+        id_field="id",
+        similarity_metric="COSINE",
+        consistency_level="Strong",
+        overwrite=False,
+      )
+    return self._milvus_vector_store
+
+  # 检查Milvus连接是否正常
+  async def check_milvus_connection(self):
+    retrieve_result = await self.async_search("hello")
+    print("✅ Milvus connection successful：", f"{env.milvus_uri}:{env.milvus_database}/{env.milvus_collection}", f"[{retrieve_result}]")
+
+  async def async_create_index_from_documents(self, documents: List[Document]) -> VectorStoreIndex:
+    vector_store = await self.get_vector_store()
+    """异步创建向量索引并插入文档"""
+    vector_index: VectorStoreIndex = VectorStoreIndex.from_vector_store(
+      vector_store=vector_store,
+      embed_model=self.embeddings,
+    )
+    from llama_index.core.node_parser import SimpleNodeParser
+    parser = SimpleNodeParser()
+    nodes = parser.get_nodes_from_documents(documents)
+    await vector_index.ainsert_nodes(nodes)
+    return vector_index
+
+  # 检索milvus文档
+  async def async_search(self, query: str, top_k: int = 5) -> List[dict]:
+    """异步向量搜索"""
+    # 创建查询引擎，设置top_k参数
+    vector_index = VectorStoreIndex.from_vector_store(
+      vector_store=await self.get_vector_store(),
+      embed_model=self.embeddings,
+    )
+    query_engine = vector_index.as_query_engine(
+      streaming=False,
+      similarity_top_k=top_k,  # 这里正确使用top_k参数
+      llm=create_llama_index_llm()
+    )
+
+    # 异步执行查询
+    response = await query_engine.aquery(query)
+
+    return {
+      "answer": str(response),
+      "sources": [
+        {
+          "text": node.node.text,
+          "metadata": node.node.metadata,
+          "score": node.score
+        }
+        for node in response.source_nodes
+      ]
+    }
+
+  async def async_delete(self, doc_id: str):
+    vector_index = VectorStoreIndex.from_vector_store(
+      vector_store=await self.get_vector_store(),
+      embed_model=self.embeddings,
+    )
+    await vector_index.adelete_ref_doc(doc_id)
+    return None
+
+
+# 全局实例
+milvus_service = MilvusService()