Merge branch 'feature/guardrails' of https://github.com/oracle/accelerated-data-science into add_langchain_deployment

Author: lu-ohai

ads/llm/__init__.py

@@ -4,8 +4,15 @@
 # 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/
 
-
-from ads.llm.langchain.plugins.llm_gen_ai import GenerativeAI
-from ads.llm.langchain.plugins.llm_md import ModelDeploymentTGI
-from ads.llm.langchain.plugins.llm_md import ModelDeploymentVLLM
-from ads.llm.langchain.plugins.embeddings import GenerativeAIEmbeddings
+try:
+    import langchain
+    from ads.llm.langchain.plugins.llm_gen_ai import GenerativeAI
+    from ads.llm.langchain.plugins.llm_md import ModelDeploymentTGI
+    from ads.llm.langchain.plugins.llm_md import ModelDeploymentVLLM
+    from ads.llm.langchain.plugins.embeddings import GenerativeAIEmbeddings
+except ImportError as ex:
+    if ex.name == "langchain":
+        raise ImportError(
+            f"{ex.msg}\nPlease install/update langchain with `pip install langchain -U`"
+        ) from ex
+    raise ex

ads/llm/chain.py

@@ -192,15 +192,15 @@ def _save_to_file(self, chain_dict, filename, overwrite=False):
             )
 
         file_ext = pathlib.Path(expanded_path).suffix.lower()
+        if file_ext not in [".yaml", ".json"]:
+            raise ValueError(
+                f"{self.__class__.__name__} can only be saved as yaml or json format."
+            )
         with open(expanded_path, "w", encoding="utf-8") as f:
             if file_ext == ".yaml":
                 yaml.safe_dump(chain_dict, f, default_flow_style=False)
             elif file_ext == ".json":
                 json.dump(chain_dict, f)
-            else:
-                raise ValueError(
-                    f"{self.__class__.__name__} can only be saved as yaml or json format."
-                )
 
     def save(self, filename: str = None, overwrite: bool = False) -> dict:
         """Serialize the sequence to a dictionary.

ads/llm/langchain/plugins/base.py

@@ -6,14 +6,13 @@
 from typing import Any, Dict, List, Optional
 
 from langchain.llms.base import LLM
-from langchain.load.serializable import Serializable
 from langchain.pydantic_v1 import BaseModel, Field, root_validator
 
 from ads.common.auth import default_signer
 from ads.config import COMPARTMENT_OCID
 
 
-class BaseLLM(LLM, Serializable):
+class BaseLLM(LLM):
     """Base OCI LLM class. Contains common attributes."""
 
     auth: dict = Field(default_factory=default_signer, exclude=True)
@@ -58,6 +57,14 @@ def is_lc_serializable(cls) -> bool:
 class GenerativeAiClientModel(BaseModel):
     """Base model for generative AI embedding model and LLM."""
 
+    # This auth is the same as the auth in BaseLLM class.
+    # However, this is needed for the Gen AI embedding model.
+    # Do not remove this attribute
+    auth: dict = Field(default_factory=default_signer, exclude=True)
+    """ADS auth dictionary for OCI authentication.
+    This can be generated by calling `ads.common.auth.api_keys()` or `ads.common.auth.resource_principal()`.
+    If this is not provided then the `ads.common.default_signer()` will be used."""
+
     client: Any  #: :meta private:
     """OCI GenerativeAiClient."""
 

ads/llm/langchain/plugins/llm_md.py

@@ -234,7 +234,7 @@ def _construct_json_body(self, prompt, params):
         }
 
     def _process_response(self, response_json: dict):
-        return str(response_json.get("generated_text", response_json)) + "\n"
+        return str(response_json.get("generated_text", response_json))
 
 
 class ModelDeploymentVLLM(ModelDeploymentLLM):

ads/llm/parsers.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,
 }
 
 

ads/llm/serialize.py

@@ -150,4 +150,7 @@ def predict(data, model=load_model(), input_schema_path=os.path.join(os.path.dir
     """
     features = pre_inference(data, input_schema_path)
     output = model.invoke(features)
+    # Return the output as is if the output is a dictionary
+    if isinstance(output, dict):
+        return output
     return {'output': output}

ads/llm/templates/score_chain.jinja2

@@ -1,5 +1,5 @@
 -r test-requirements.txt
--e ".[bds,data,geo,huggingface,notebook,onnx,opctl,optuna,pii,spark,tensorflow,text,torch,viz]"
+-e ".[bds,data,geo,huggingface,llm,notebook,onnx,opctl,optuna,pii,spark,tensorflow,text,torch,viz]"
 arff
 category_encoders
 dask

dev-requirements.txt

@@ -181,6 +181,9 @@ pii = [
   "spacy-transformers==1.2.5",
   "spacy==3.6.1",
 ]
+llm = [
+  "langchain>=0.0.295"
+]
 
 [project.urls]
 "Github" = "https://github.com/oracle/accelerated-data-science"

pyproject.toml

@@ -4,7 +4,9 @@
 # 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 json
 import os
+import tempfile
 from typing import Any, List, Dict, Mapping, Optional
 from unittest import TestCase
 from langchain.callbacks.manager import CallbackManagerForLLMRun
@@ -146,8 +148,8 @@ def test_guardrail_sequence_with_template_and_toxicity(self):
         def test_fn(chain: GuardrailSequence):
             output = chain.run("cats", num_generations=5)
             self.assertIsInstance(output, GuardrailIO)
-            self.assertIsInstance(output.data, list)
-            self.assertEqual(len(output.data), 1)
+            self.assertIsInstance(output.data, str)
+            self.assertEqual(output.data, "Tell me a joke about cats")
             self.assertIsInstance(output.info, list)
             self.assertEqual(len(output.info), len(chain.steps))
 
@@ -166,10 +168,37 @@ def test_guardrail_sequence_with_filtering(self):
         def test_fn(chain: GuardrailSequence):
             output = chain.run(self.TOXIC_CONTENT)
             self.assertIsInstance(output, GuardrailIO)
-            self.assertIsInstance(output.data, list)
-            self.assertEqual(len(output.data), 1)
-            self.assertEqual(output.data[0], message)
+            self.assertIsInstance(output.data, str)
+            self.assertEqual(output.data, message)
             self.assertIsInstance(output.info, list)
             self.assertEqual(len(output.info), len(chain.steps))
 
         self.assert_before_and_after_serialization(test_fn, chain)
+
+    def test_empty_sequence(self):
+        """Tests empty sequence."""
+        seq = GuardrailSequence()
+        self.assertEqual(seq.steps, [])
+
+    def test_save_to_file(self):
+        """Tests saving to file."""
+        message = "Let's talk something else."
+        toxicity = HuggingFaceEvaluation(
+            path="toxicity",
+            load_args=self.LOAD_ARGS,
+            threshold=0.5,
+            custom_msg=message,
+        )
+        chain = GuardrailSequence.from_sequence(self.FAKE_LLM | toxicity)
+        try:
+            temp = tempfile.NamedTemporaryFile(suffix=".json", delete=False)
+            temp.close()
+            with self.assertRaises(FileExistsError):
+                serialized = chain.save(temp.name)
+            with self.assertRaises(ValueError):
+                chain.save("abc.html")
+            serialized = chain.save(temp.name, overwrite=True)
+            with open(temp.name, "r", encoding="utf-8") as f:
+                self.assertEqual(json.load(f), serialized)
+        finally:
+            os.unlink(temp.name)