Merge pull request #1045 from Jacksonxhx/Jackson

Integrated Milvus Vector DB into main

application/core/settings.py

@@ -18,7 +18,7 @@ class Settings(BaseSettings):
     DEFAULT_MAX_HISTORY: int = 150
     MODEL_TOKEN_LIMITS: dict = {"gpt-3.5-turbo": 4096, "claude-2": 1e5}
     UPLOAD_FOLDER: str = "inputs"
-    VECTOR_STORE: str = "faiss"  # "faiss" or "elasticsearch" or "qdrant"
+    VECTOR_STORE: str = "faiss"  # "faiss" or "elasticsearch" or "qdrant" or "milvus"
     RETRIEVERS_ENABLED: list = ["classic_rag", "duckduck_search"] # also brave_search
 
     API_URL: str = "http://localhost:7091"  # backend url for celery worker
@@ -62,6 +62,11 @@ class Settings(BaseSettings):
     QDRANT_PATH: Optional[str] = None
     QDRANT_DISTANCE_FUNC: str = "Cosine"
 
+    # Milvus vectorstore config
+    MILVUS_COLLECTION_NAME: Optional[str] = "docsgpt"
+    MILVUS_URI: Optional[str] = "./milvus_local.db"   # milvus lite version as default
+    MILVUS_TOKEN: Optional[str] = ""
+
     BRAVE_SEARCH_API_KEY: Optional[str] = None
 
     FLASK_DEBUG_MODE: bool = False

application/requirements.txt

@@ -10,12 +10,14 @@ elasticsearch==8.14.0
 escodegen==1.0.11
 esprima==4.0.1
 Flask==3.0.1
-faiss-cpu==1.8.0
+faiss-cpu==1.8.0.post1
 gunicorn==23.0.0
 html2text==2020.1.16
 javalang==0.13.0
-langchain==0.1.4
-langchain-openai==0.0.5
+langchain==0.2.16
+langchain-community==0.2.16
+langchain-core==0.2.38
+langchain-openai==0.1.23
 openapi3_parser==1.1.16
 pandas==2.2.2
 pydantic_settings==2.4.0
@@ -27,7 +29,7 @@ redis==5.0.1
 Requests==2.32.0
 retry==0.9.2
 sentence-transformers
-tiktoken
+tiktoken==0.7.0
 torch
 tqdm==4.66.3
 transformers==4.44.0

application/vectorstore/base.py

@@ -1,13 +1,30 @@
 from abc import ABC, abstractmethod
 import os
-from langchain_community.embeddings import (
-    HuggingFaceEmbeddings,
-    CohereEmbeddings,
-    HuggingFaceInstructEmbeddings,
-)
+from sentence_transformers import SentenceTransformer
 from langchain_openai import OpenAIEmbeddings
 from application.core.settings import settings
 
+class EmbeddingsWrapper:
+    def __init__(self, model_name, *args, **kwargs):
+        self.model = SentenceTransformer(model_name, config_kwargs={'allow_dangerous_deserialization': True}, *args, **kwargs)
+        self.dimension = self.model.get_sentence_embedding_dimension()
+
+    def embed_query(self, query: str):
+        return self.model.encode(query).tolist()
+
+    def embed_documents(self, documents: list):
+        return self.model.encode(documents).tolist()
+
+    def __call__(self, text):
+        if isinstance(text, str):
+            return self.embed_query(text)
+        elif isinstance(text, list):
+            return self.embed_documents(text)
+        else:
+            raise ValueError("Input must be a string or a list of strings")
+
+
+
 class EmbeddingsSingleton:
     _instances = {}
 
@@ -23,16 +40,15 @@ def get_instance(embeddings_name, *args, **kwargs):
     def _create_instance(embeddings_name, *args, **kwargs):
         embeddings_factory = {
             "openai_text-embedding-ada-002": OpenAIEmbeddings,
-            "huggingface_sentence-transformers/all-mpnet-base-v2": HuggingFaceEmbeddings,
-            "huggingface_sentence-transformers-all-mpnet-base-v2": HuggingFaceEmbeddings,
-            "huggingface_hkunlp/instructor-large": HuggingFaceInstructEmbeddings,
-            "cohere_medium": CohereEmbeddings
+            "huggingface_sentence-transformers/all-mpnet-base-v2": lambda: EmbeddingsWrapper("sentence-transformers/all-mpnet-base-v2"),
+            "huggingface_sentence-transformers-all-mpnet-base-v2": lambda: EmbeddingsWrapper("sentence-transformers/all-mpnet-base-v2"),
+            "huggingface_hkunlp/instructor-large": lambda: EmbeddingsWrapper("hkunlp/instructor-large"),
         }
 
-        if embeddings_name not in embeddings_factory:
-            raise ValueError(f"Invalid embeddings_name: {embeddings_name}")
-
-        return embeddings_factory[embeddings_name](*args, **kwargs)
+        if embeddings_name in embeddings_factory:
+            return embeddings_factory[embeddings_name](*args, **kwargs)
+        else:
+            return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
 
 class BaseVectorStore(ABC):
     def __init__(self):
@@ -58,22 +74,14 @@ def _get_embeddings(self, embeddings_name, embeddings_key=None):
                     embeddings_name,
                     openai_api_key=embeddings_key
                 )
-        elif embeddings_name == "cohere_medium":
-            embedding_instance = EmbeddingsSingleton.get_instance(
-                embeddings_name,
-                cohere_api_key=embeddings_key
-            )
         elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
             if os.path.exists("./model/all-mpnet-base-v2"):
                 embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name,
-                    model_name="./model/all-mpnet-base-v2",
-                    model_kwargs={"device": "cpu"}
+                    embeddings_name="./model/all-mpnet-base-v2",
                 )
             else:
                 embedding_instance = EmbeddingsSingleton.get_instance(
                     embeddings_name,
-                    model_kwargs={"device": "cpu"}
                 )
         else:
             embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)

application/vectorstore/faiss.py

@@ -14,7 +14,8 @@ def __init__(self, path, embeddings_key, docs_init=None):
             )
         else:
             self.docsearch = FAISS.load_local(
-                self.path, embeddings
+                self.path, embeddings, 
+                allow_dangerous_deserialization=True
             )
         self.assert_embedding_dimensions(embeddings)
 
@@ -37,10 +38,10 @@ def assert_embedding_dimensions(self, embeddings):
         """
         if settings.EMBEDDINGS_NAME == "huggingface_sentence-transformers/all-mpnet-base-v2":
             try:
-                word_embedding_dimension = embeddings.client[1].word_embedding_dimension
+                word_embedding_dimension = embeddings.dimension
             except AttributeError as e:
-                raise AttributeError("word_embedding_dimension not found in embeddings.client[1]") from e
+                raise AttributeError("'dimension' attribute not found in embeddings instance. Make sure the embeddings object is properly initialized.") from e
             docsearch_index_dimension = self.docsearch.index.d
             if word_embedding_dimension != docsearch_index_dimension:
-                raise ValueError(f"word_embedding_dimension ({word_embedding_dimension}) " +
-                                 f"!= docsearch_index_word_embedding_dimension ({docsearch_index_dimension})")
+                raise ValueError(f"Embedding dimension mismatch: embeddings.dimension ({word_embedding_dimension}) " +
+                                 f"!= docsearch index dimension ({docsearch_index_dimension})")

application/vectorstore/milvus.py

@@ -0,0 +1,37 @@
+from typing import List, Optional
+from uuid import uuid4
+
+
+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__()
+        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.MILVUS_COLLECTION_NAME,
+            connection_args=connection_args,
+        )
+        self._path = path
+
+    def search(self, question, k=2, *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):
+        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
+
+    def delete_index(self, *args, **kwargs):
+        pass

application/vectorstore/vector_creator.py

@@ -1,5 +1,6 @@
 from application.vectorstore.faiss import FaissStore
 from application.vectorstore.elasticsearch import ElasticsearchStore
+from application.vectorstore.milvus import MilvusStore
 from application.vectorstore.mongodb import MongoDBVectorStore
 from application.vectorstore.qdrant import QdrantStore
 
@@ -10,6 +11,7 @@ class VectorCreator:
         "elasticsearch": ElasticsearchStore,
         "mongodb": MongoDBVectorStore,
         "qdrant": QdrantStore,
+        "milvus": MilvusStore,
     }
 
     @classmethod

docker-compose-azure.yaml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   frontend:
     build: ./frontend

docker-compose-dev.yaml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
 
   redis:

docker-compose-local.yaml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   frontend:
     build: ./frontend

docker-compose-mock.yaml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   frontend:
     build: ./frontend

docker-compose.yaml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   frontend:
     build: ./frontend