adding unit test

Author: Ziqun Ye

ads/llm/serialize.py

@@ -4,6 +4,7 @@
 # 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
@@ -19,18 +20,15 @@
 from langchain.load import dumpd
 from langchain.load.load import load as lc_load
 from langchain.load.serializable import Serializable
-from langchain.vectorstores import OpenSearchVectorSearch, FAISS
+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
-import json
-import base64
-
 
 # This is a temp solution for supporting custom LLM in legacy load_chain
 __lc_llm_dict = llms.get_type_to_cls_dict()
@@ -58,15 +56,15 @@ def type():
     @staticmethod
     def load(config: dict, **kwargs):
         config["kwargs"]["embedding_function"] = load(
-            config["kwargs"]["embedding_function"]
+            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=os.environ.get("oci_opensearch_verify_certs", False),
+            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),
         )
 
@@ -102,7 +100,7 @@ def type():
 
     @staticmethod
     def load(config: dict, **kwargs):
-        embedding_function = load(config["embedding_function"])
+        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
@@ -124,7 +122,7 @@ def save(obj):
 vectordb_serialization = {"OpenSearchVectorSearch": OpenSearchVectorDBSerializer, "FAISS": FaissSerializer}
 
 
-class RetrieverQASerializer:
+class RetrievalQASerializer:
     """
     Serializer for RetrieverQA class
     """
@@ -137,7 +135,7 @@ 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"))
+        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)
 
@@ -168,7 +166,7 @@ def save(obj):
     GuardrailSequence: GuardrailSequence.save,
     CustomGuardrailBase: CustomGuardrailBase.save,
     RunnableParallel: RunnableParallelSerializer.save,
-    RetrievalQA: RetrieverQASerializer.save,
+    RetrievalQA: RetrievalQASerializer.save,
 }
 
 # Mapping _type to custom deserialization functions
@@ -177,7 +175,7 @@ def save(obj):
     GuardrailSequence.type(): GuardrailSequence.load,
     CustomGuardrailBase.type(): CustomGuardrailBase.load,
     RunnableParallelSerializer.type(): RunnableParallelSerializer.load,
-    RetrieverQASerializer.type(): RetrieverQASerializer.load,
+    RetrievalQASerializer.type(): RetrievalQASerializer.load,
 }
 
 

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()