integrated milvus db

application/vectorstore/milvus.py

@@ -0,0 +1,38 @@
+from typing import List, Optional
+from langchain_community.vectorstores.milvus import Milvus
+
+from application.core.settings import settings
+from application.vectorstore.base import BaseVectorStore
+
+
+class MilvusStore(BaseVectorStore):
+    def __init__(self, path: str = "", embeddings_key: str = "embeddings"):
+        super().__init__()
+        if path:
+            connection_args ={
+                "uri": path,
+                "tpken": settings.MILVUS_TOKEN,
+            }
+        else:
+            connection_args = {
+                "uri": settings.MILVUS_URL,
+                'token': settings.MILVUS_TOKEN,
+            }
+        self._docsearch = Milvus(
+            embedding_function=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key),
+            collection_name=settings.COLLECTION_NAME,
+            connection_args=connection_args,
+            drop_old=True,
+        )
+
+    def search(self, question, k=2, *args, **kwargs):
+        return self._docsearch.similarity_search(query=question, k=k, *args, **kwargs)
+
+    def add_texts(self, texts: List[str], metadatas: Optional[List[dict]], *args, **kwargs):
+        return self._docsearch.add_texts(texts=texts, metadatas=metadatas, *args, **kwargs)
+
+    def save_local(self, *args, **kwargs):
+        pass
+
+    def delete_index(self, *args, **kwargs):
+        pass