feat Milvus integration

application/vectorstore/milvus.py

@@ -1,5 +1,6 @@
 from typing import List, Optional
-from langchain_community.vectorstores.milvus import Milvus
+from uuid import uuid4
+
 
 from application.core.settings import settings
 from application.vectorstore.base import BaseVectorStore
@@ -8,28 +9,26 @@ 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,
-            }
+        from langchain_milvus import Milvus
+
+        connection_args = {
+            "uri": settings.MILVUS_URI,
+            "token": settings.MILVUS_TOKEN,
+        }
         self._docsearch = Milvus(
             embedding_function=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key),
-            collection_name=settings.COLLECTION_NAME,
+            collection_name=settings.MILVUS_COLLECTION_NAME,
             connection_args=connection_args,
-            drop_old=True,
         )
+        self._path = path
 
     def search(self, question, k=2, *args, **kwargs):
-        return self._docsearch.similarity_search(query=question, k=k, *args, **kwargs)
+        return self._docsearch.similarity_search(query=question, k=k, filter={"path": self._path} *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)
+        ids = [str(uuid4()) for _ in range(len(texts))]
+
+        return self._docsearch.add_texts(texts=texts, metadatas=metadatas, ids=ids, *args, **kwargs)
 
     def save_local(self, *args, **kwargs):
         pass