Support OpenSearchVectorSearch and FAISS serializer for vector store (#461)

z7ye · web-flow · commit 0315d8b2826b · 2023-12-01T15:46:53.000-08:00
diff --git a/ads/llm/serialize.py b/ads/llm/serialize.py
@@ -4,27 +4,32 @@
 # Copyright (c) 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import base64
 import json
 import os
 import tempfile
+from copy import deepcopy
 from typing import Any, Dict, List, Optional
 
 import fsspec
 import yaml
 from langchain import llms
-from langchain.llms import loading
+from langchain.chains import RetrievalQA
 from langchain.chains.loading import load_chain_from_config
+from langchain.llms import loading
+from langchain.load import dumpd
 from langchain.load.load import Reviver
 from langchain.load.serializable import Serializable
+from langchain.vectorstores import FAISS, OpenSearchVectorSearch
+from opensearchpy.client import OpenSearch
 
 from ads.common.auth import default_signer
 from ads.common.object_storage_details import ObjectStorageDetails
-from ads.llm import GenerativeAI, ModelDeploymentVLLM, ModelDeploymentTGI
+from ads.llm import GenerativeAI, ModelDeploymentTGI, ModelDeploymentVLLM
 from ads.llm.chain import GuardrailSequence
 from ads.llm.guardrails.base import CustomGuardrailBase
 from ads.llm.patch import RunnableParallel, RunnableParallelSerializer
 
-
 # This is a temp solution for supporting custom LLM in legacy load_chain
 __lc_llm_dict = llms.get_type_to_cls_dict()
 __lc_llm_dict[GenerativeAI.__name__] = lambda: GenerativeAI
@@ -39,11 +44,129 @@ def __new_type_to_cls_dict():
 llms.get_type_to_cls_dict = __new_type_to_cls_dict
 loading.get_type_to_cls_dict = __new_type_to_cls_dict
 
+
+class OpenSearchVectorDBSerializer:
+    """
+    Serializer for OpenSearchVectorSearch class
+    """
+    @staticmethod
+    def type():
+        return OpenSearchVectorSearch.__name__
+
+    @staticmethod
+    def load(config: dict, **kwargs):
+        config["kwargs"]["embedding_function"] = load(
+            config["kwargs"]["embedding_function"], **kwargs
+        )
+        return OpenSearchVectorSearch(
+            **config["kwargs"],
+            http_auth=(
+                os.environ.get("OCI_OPENSEARCH_USERNAME", None),
+                os.environ.get("OCI_OPENSEARCH_PASSWORD", None),
+            ),
+            verify_certs=True if os.environ.get("OCI_OPENSEARCH_VERIFY_CERTS", None).lower() == "true" else False,
+            ca_certs=os.environ.get("OCI_OPENSEARCH_CA_CERTS", None),
+        )
+
+    @staticmethod
+    def save(obj):
+        serialized = dumpd(obj)
+        serialized["type"] = "constructor"
+        serialized["_type"] = OpenSearchVectorDBSerializer.type()
+        kwargs = {}
+        for key, val in obj.__dict__.items():
+            if key == "client":
+                if isinstance(val, OpenSearch):
+                    client_info = val.transport.hosts[0]
+                    opensearch_url = (
+                        f"https://{client_info['host']}:{client_info['port']}"
+                    )
+                    kwargs.update({"opensearch_url": opensearch_url})
+                else:
+                    raise NotImplementedError("Only support OpenSearch client.")
+                continue
+            kwargs[key] = dump(val)
+        serialized["kwargs"] = kwargs
+        return serialized
+
+
+class FaissSerializer:
+    """
+    Serializer for OpenSearchVectorSearch class
+    """
+    @staticmethod
+    def type():
+        return FAISS.__name__
+
+    @staticmethod
+    def load(config: dict, **kwargs):
+        embedding_function = load(config["embedding_function"], **kwargs)
+        decoded_pkl = base64.b64decode(json.loads(config["vectordb"]))
+        return FAISS.deserialize_from_bytes(
+            embeddings=embedding_function, serialized=decoded_pkl
+        )  # Load the index
+
+    @staticmethod
+    def save(obj):
+        serialized = {}
+        serialized["_type"] = FaissSerializer.type()
+        pkl = obj.serialize_to_bytes()
+        # Encoding bytes to a base64 string
+        encoded_pkl = base64.b64encode(pkl).decode('utf-8')
+        # Serializing the base64 string
+        serialized["vectordb"] = json.dumps(encoded_pkl)
+        serialized["embedding_function"] = dump(obj.__dict__["embedding_function"])
+        return serialized
+
+# Mapping class to vector store serialization functions
+vectordb_serialization = {"OpenSearchVectorSearch": OpenSearchVectorDBSerializer, "FAISS": FaissSerializer}
+
+
+class RetrievalQASerializer:
+    """
+    Serializer for RetrieverQA class
+    """
+    @staticmethod
+    def type():
+        return "retrieval_qa"
+
+    @staticmethod
+    def load(config: dict, **kwargs):
+        config_param = deepcopy(config)
+        retriever_kwargs = config_param.pop("retriever_kwargs")
+        vectordb_serializer = vectordb_serialization[config_param["vectordb"]["class"]]
+        vectordb = vectordb_serializer.load(config_param.pop("vectordb"), **kwargs)
+        retriever = vectordb.as_retriever(**retriever_kwargs)
+        return load_chain_from_config(config=config_param, retriever=retriever)
+
+    @staticmethod
+    def save(obj):
+        serialized = obj.dict()
+        retriever_kwargs = {}
+        for key, val in obj.retriever.__dict__.items():
+            if key not in ["tags", "metadata", "vectorstore"]:
+                retriever_kwargs[key] = val
+        serialized["retriever_kwargs"] = retriever_kwargs
+        serialized["vectordb"] = {"class": obj.retriever.vectorstore.__class__.__name__}
+
+        vectordb_serializer = vectordb_serialization[serialized["vectordb"]["class"]]
+        serialized["vectordb"].update(
+            vectordb_serializer.save(obj.retriever.vectorstore)
+        )
+
+        if serialized["vectordb"]["class"] not in vectordb_serialization:
+            raise NotImplementedError(
+                f"VectorDBSerializer for {serialized['vectordb']['class']} is not implemented."
+            )
+        return serialized
+
+
 # Mapping class to custom serialization functions
 custom_serialization = {
     GuardrailSequence: GuardrailSequence.save,
     CustomGuardrailBase: CustomGuardrailBase.save,
     RunnableParallel: RunnableParallelSerializer.save,
+    RetrievalQA: RetrievalQASerializer.save,
 }
 
 # Mapping _type to custom deserialization functions
@@ -52,6 +175,7 @@ def __new_type_to_cls_dict():
     GuardrailSequence.type(): GuardrailSequence.load,
     CustomGuardrailBase.type(): CustomGuardrailBase.load,
     RunnableParallelSerializer.type(): RunnableParallelSerializer.load,
+    RetrievalQASerializer.type(): RetrievalQASerializer.load,
 }
 
 
diff --git a/tests/unitary/with_extras/langchain/test_serializers.py b/tests/unitary/with_extras/langchain/test_serializers.py
@@ -0,0 +1,161 @@
+import unittest
+from langchain.load.serializable import Serializable
+from langchain.schema.embeddings import Embeddings
+
+from langchain.vectorstores import OpenSearchVectorSearch, FAISS
+
+
+import unittest
+from ads.llm.serialize import OpenSearchVectorDBSerializer, FaissSerializer, RetrievalQASerializer
+from tests.unitary.with_extras.langchain.test_guardrails import FakeLLM
+import os
+from unittest import mock
+from typing import Any, Dict, List, Mapping, Optional
+from langchain.chains import RetrievalQA
+from langchain import llms
+from langchain.llms import loading
+
+
+
+
+class FakeEmbeddings(Serializable, Embeddings):
+    """Fake LLM for testing purpose."""
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom_embeddings"
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """This class can be serialized with default LangChain serialization."""
+        return True
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        return [[1] * 1024 for text in texts]
+
+    def embed_query(self, text: str) -> List[float]:
+        return [1] * 1024
+    
+    
+class TestOpensearchSearchVectorSerializers(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.env_patcher = mock.patch.dict(os.environ, {"OCI_OPENSEARCH_USERNAME": "username",
+                                                    "OCI_OPENSEARCH_PASSWORD": "password",
+                                                    "OCI_OPENSEARCH_VERIFY_CERTS": "True",
+                                                    "OCI_OPENSEARCH_CA_CERTS": "/path/to/cert.pem"})
+        cls.env_patcher.start()
+        cls.index_name = "test_index"
+        cls.embeddings = FakeEmbeddings()
+        cls.opensearch = OpenSearchVectorSearch(
+            "https://localhost:8888",
+            embedding_function=cls.embeddings,
+            index_name=cls.index_name,
+            engine="lucene",
+            http_auth=(os.environ["OCI_OPENSEARCH_USERNAME"], os.environ["OCI_OPENSEARCH_PASSWORD"]),
+            verify_certs=os.environ["OCI_OPENSEARCH_VERIFY_CERTS"],
+            ca_certs=os.environ["OCI_OPENSEARCH_CA_CERTS"],
+        )
+        cls.serializer = OpenSearchVectorDBSerializer()
+        super().setUpClass()
+
+    def test_type(self):
+        # Test type()
+        self.assertEqual(self.serializer.type(), "OpenSearchVectorSearch")
+
+    def test_save(self):
+        serialized = self.serializer.save(self.opensearch)
+        assert serialized["id"] == ['langchain', 'vectorstores', 'opensearch_vector_search', 'OpenSearchVectorSearch']
+        assert serialized["kwargs"]["opensearch_url"] == "https://localhost:8888"
+        assert serialized["kwargs"]["engine"] == "lucene"
+        assert serialized["_type"] == "OpenSearchVectorSearch"
+
+    def test_load(self):
+        serialized = self.serializer.save(self.opensearch)
+        new_opensearch = self.serializer.load(serialized, valid_namespaces=["tests"])
+        assert isinstance(new_opensearch, OpenSearchVectorSearch)
+
+
+class TestFAISSSerializers(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.embeddings = FakeEmbeddings()
+        text_embedding_pair = [("test", [1] * 1024)]
+        cls.db = FAISS.from_embeddings(text_embedding_pair, cls.embeddings)
+        cls.serializer = FaissSerializer()
+        super().setUpClass()
+
+    def test_type(self):
+        self.assertEqual(self.serializer.type(), "FAISS")
+
+    def test_save(self):
+        serialized = self.serializer.save(self.db)
+        assert serialized["embedding_function"]["id"] == ["tests", "unitary", "with_extras", "langchain", "test_serializers", "FakeEmbeddings"]
+        assert isinstance(serialized["vectordb"], str)
+
+    def test_load(self):
+        serialized = self.serializer.save(self.db)
+        new_db = self.serializer.load(serialized, valid_namespaces=["tests"])
+        assert isinstance(new_db, FAISS)
+
+
+class TestRetrievalQASerializer(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # Create a sample RetrieverQA object for testing
+        cls.llm = FakeLLM()
+        cls.embeddings = FakeEmbeddings()
+        text_embedding_pair = [("test", [1] * 1024)]
+        cls.db = FAISS.from_embeddings(text_embedding_pair, cls.embeddings)
+        cls.serializer = FaissSerializer()
+        cls.retriever = cls.db.as_retriever()
+        cls.qa = RetrievalQA.from_chain_type(llm=cls.llm,
+                                            chain_type="stuff",
+                                            retriever=cls.retriever)
+        cls.serializer = RetrievalQASerializer()
+        from copy import deepcopy
+        cls.original_type_to_cls_dict = deepcopy(llms.get_type_to_cls_dict())
+        __lc_llm_dict = llms.get_type_to_cls_dict()
+        __lc_llm_dict["custom_embedding"] = lambda: FakeEmbeddings
+        __lc_llm_dict["custom"] = lambda: FakeLLM
+
+        def __new_type_to_cls_dict():
+            return __lc_llm_dict
+
+        llms.get_type_to_cls_dict = __new_type_to_cls_dict
+        loading.get_type_to_cls_dict = __new_type_to_cls_dict
+
+    def test_type(self):
+        self.assertEqual(self.serializer.type(), "retrieval_qa")
+
+    def test_save(self):
+        # Serialize the RetrieverQA object
+        serialized = self.serializer.save(self.qa)
+
+        # Ensure that the serialized object is a dictionary
+        self.assertIsInstance(serialized, dict)
+
+        # Ensure that the serialized object contains the necessary keys
+        self.assertIn("combine_documents_chain", serialized)
+        self.assertIn("retriever_kwargs", serialized)
+        serialized["vectordb"]["class"] == "FAISS"
+
+    def test_load(self):
+        # Create a sample config dictionary
+        serialized = self.serializer.save(self.qa)
+
+        # Deserialize the serialized object
+        deserialized = self.serializer.load(serialized, valid_namespaces=["tests"])
+
+        # Ensure that the deserialized object is an instance of RetrieverQA
+        self.assertIsInstance(deserialized, RetrievalQA)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        llms.get_type_to_cls_dict = cls.original_type_to_cls_dict
+        loading.get_type_to_cls_dict = cls.original_type_to_cls_dict
+        return super().tearDownClass()
+
+
+if __name__ == "__main__":
+    unittest.main()
