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

Author: lu-ohai

ads/llm/chain.py

@@ -178,6 +178,9 @@ def run(self, input: Any, num_generations: int = 1, **kwargs) -> GuardrailIO:
         if self.log_info or os.environ.get(LOG_ADS_GUARDRAIL_INFO) == "1":
             # LOG_ADS_GUARDRAIL_INFO is set to "1" in score.py by default.
             print(obj.dict())
+        # If the output is a singleton list, take it out of the list.
+        if isinstance(obj.data, list) and len(obj.data) == 1:
+            obj.data = obj.data[0]
         return obj
 
     def _save_to_file(self, chain_dict, filename, overwrite=False):

ads/llm/guardrails/base.py

@@ -78,7 +78,7 @@ def __repr__(self) -> str:
             steps.append(f"{run_info.name} - {run_info.metrics}")
         if run_info:
             steps.append(str(run_info.output))
-        return "\n".join(steps)
+        return "\n".join(steps) + "\n\n" + str(self)
 
 
 class BlockedByGuardrail(ToolException):

ads/llm/langchain/plugins/base.py

@@ -46,7 +46,7 @@ def _print_response(self, completion, response):
 
     @classmethod
     def get_lc_namespace(cls) -> List[str]:
-        """Get the namespace of the langchain object."""
+        """Get the namespace of the LangChain object."""
         return ["ads", "llm"]
 
     @classmethod
@@ -56,10 +56,12 @@ def is_lc_serializable(cls) -> bool:
 
 
 class GenerativeAiClientModel(BaseModel):
+    """Base model for generative AI embedding model and LLM."""
+
     client: Any  #: :meta private:
     """OCI GenerativeAiClient."""
 
-    compartment_id: str
+    compartment_id: str = None
     """Compartment ID of the caller."""
 
     endpoint_kwargs: Dict[str, Any] = {}
@@ -90,7 +92,9 @@ def validate_environment(  # pylint: disable=no-self-argument
             client_kwargs.update(values["client_kwargs"])
             values["client"] = GenerativeAiClient(**auth, **client_kwargs)
         # Set default compartment ID
-        if "compartment_id" not in values and COMPARTMENT_OCID:
-            values["compartment_id"] = COMPARTMENT_OCID
-
+        if not values.get("compartment_id"):
+            if COMPARTMENT_OCID:
+                values["compartment_id"] = COMPARTMENT_OCID
+            else:
+                raise ValueError("Please specify compartment_id.")
         return values

ads/llm/serialize.py

@@ -14,7 +14,7 @@
 from langchain import llms
 from langchain.llms import loading
 from langchain.chains.loading import load_chain_from_config
-from langchain.load.load import load as __lc_load
+from langchain.load.load import Reviver, load as __lc_load
 from langchain.load.serializable import Serializable
 
 from ads.common.auth import default_signer
@@ -76,14 +76,32 @@ def load(
     Returns:
         Revived LangChain objects.
     """
+    # Add ADS as valid namespace
     if not valid_namespaces:
         valid_namespaces = []
     if "ads" not in valid_namespaces:
         valid_namespaces.append("ads")
 
+    reviver = Reviver(secrets_map, valid_namespaces)
+
+    def _load(obj: Any) -> Any:
+        if isinstance(obj, dict):
+            if "_type" in obj and obj["_type"] in custom_deserialization:
+                if valid_namespaces:
+                    kwargs["valid_namespaces"] = valid_namespaces
+                if secrets_map:
+                    kwargs["secret_map"] = secrets_map
+                return custom_deserialization[obj["_type"]](obj, **kwargs)
+            # Need to revive leaf nodes before reviving this node
+            loaded_obj = {k: _load(v) for k, v in obj.items()}
+            return reviver(loaded_obj)
+        if isinstance(obj, list):
+            return [_load(o) for o in obj]
+        return obj
+
     if isinstance(obj, dict) and "_type" in obj:
         obj_type = obj["_type"]
-        # Check if the object requires a custom function to load.
+        # Check if the object has custom load function.
         if obj_type in custom_deserialization:
             if valid_namespaces:
                 kwargs["valid_namespaces"] = valid_namespaces
@@ -93,7 +111,7 @@ def load(
         # Legacy chain
         return load_chain_from_config(obj, **kwargs)
 
-    return __lc_load(obj, secrets_map=secrets_map, valid_namespaces=valid_namespaces)
+    return _load(obj)
 
 
 def load_from_yaml(
@@ -144,11 +162,30 @@ def default(obj: Any) -> Any:
     TypeError
         If the object is not LangChain serializable.
     """
+    for super_class, save_fn in custom_serialization.items():
+        if isinstance(obj, super_class):
+            return save_fn(obj)
     if isinstance(obj, Serializable) and obj.is_lc_serializable():
         return obj.to_json()
     raise TypeError(f"Serialization of {type(obj)} is not supported.")
 
 
+def __save(obj):
+    """Calls the legacy save method to save the object to temp json
+    then load it into a dictionary.
+    """
+    try:
+        temp_file = tempfile.NamedTemporaryFile(
+            mode="w", encoding="utf-8", suffix=".json", delete=False
+        )
+        temp_file.close()
+        obj.save(temp_file.name)
+        with open(temp_file.name, "r", encoding="utf-8") as f:
+            return json.load(f)
+    finally:
+        os.unlink(temp_file.name)
+
+
 def dump(obj: Any) -> Dict[str, Any]:
     """Return a json dict representation of an object.
 
@@ -167,14 +204,14 @@ def dump(obj: Any) -> Dict[str, Any]:
     ):
         # The object is not is_lc_serializable.
         # However, it supports the legacy save() method.
-        try:
-            temp_file = tempfile.NamedTemporaryFile(
-                mode="w", encoding="utf-8", suffix=".json", delete=False
-            )
-            temp_file.close()
-            obj.save(temp_file.name)
-            with open(temp_file.name, "r", encoding="utf-8") as f:
-                return json.load(f)
-        finally:
-            os.unlink(temp_file.name)
-    return json.loads(json.dumps(obj, default=default))
+        return __save(obj)
+    # The object is is_lc_serializable.
+    # However, some properties may not be serializable
+    # Here we try to dump the object and fallback to the save() method
+    # if there is an error.
+    try:
+        return json.loads(json.dumps(obj, default=default))
+    except TypeError as ex:
+        if isinstance(obj, Serializable) and hasattr(obj, "save"):
+            return __save(obj)
+        raise ex

ads/llm/templates/score_chain.jinja2

@@ -165,5 +165,5 @@ def predict(data, model=load_model(), input_schema_path=os.path.join(os.path.dir
 
     """
     features = pre_inference(data, input_schema_path)
-    yhat = post_inference(model.run(features))
+    yhat = post_inference(model.invoke(features))
     return {'prediction': yhat}