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

Author: lu-ohai

ads/llm/langchain/plugins/llm_md.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, List, Optional
 
 import requests
+from oci.auth import signers
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 
 from ads.llm.langchain.plugins.base import BaseLLM
@@ -134,11 +135,22 @@ def send_request(
             header.pop("content_type", DEFAULT_CONTENT_TYPE_JSON)
             or DEFAULT_CONTENT_TYPE_JSON
         )
+        timeout = kwargs.pop("timeout", DEFAULT_TIME_OUT)
         request_kwargs = {"json": data}
         request_kwargs["headers"] = header
-        request_kwargs["auth"] = self.auth.get("signer")
-        timeout = kwargs.pop("timeout", DEFAULT_TIME_OUT)
-        response = requests.post(endpoint, timeout=timeout, **request_kwargs, **kwargs)
+        signer = self.auth.get("signer")
+
+        attempts = 0
+        while attempts < 2:
+            request_kwargs["auth"] = signer
+            response = requests.post(
+                endpoint, timeout=timeout, **request_kwargs, **kwargs
+            )
+            if response.status_code == 401 and self.is_principal_signer(signer):
+                signer.refresh_security_token()
+                attempts += 1
+                continue
+            break
 
         try:
             response.raise_for_status()
@@ -155,6 +167,21 @@ def send_request(
 
         return response_json
 
+    @staticmethod
+    def is_principal_signer(signer):
+        """Checks if the signer is instance principal or resource principal signer."""
+        if (
+            isinstance(signer, signers.InstancePrincipalsSecurityTokenSigner)
+            or isinstance(signer, signers.ResourcePrincipalsFederationSigner)
+            or isinstance(signer, signers.EphemeralResourcePrincipalSigner)
+            or isinstance(signer, signers.EphemeralResourcePrincipalV21Signer)
+            or isinstance(signer, signers.NestedResourcePrincipals)
+            or isinstance(signer, signers.OkeWorkloadIdentityResourcePrincipalSigner)
+        ):
+            return True
+        else:
+            return False
+
 
 class ModelDeploymentTGI(ModelDeploymentLLM):
     """OCI Data Science Model Deployment TGI Endpoint.

ads/llm/patch.py

@@ -1,5 +1,6 @@
 from langchain.schema.runnable import RunnableParallel
-from langchain.load import dumpd, load
+from langchain.load.dump import dumpd
+from langchain.load.load import load
 
 
 class RunnableParallelSerializer:

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 Reviver, load as __lc_load
+from langchain.load.load import Reviver
 from langchain.load.serializable import Serializable
 
 from ads.common.auth import default_signer

ads/llm/templates/score_chain.jinja2

@@ -132,22 +132,6 @@ def pre_inference(data, input_schema_path):
     """
     return deserialize(data, input_schema_path)
 
-
-def post_inference(yhat):
-    """
-    Post-process the model results
-
-    Parameters
-    ----------
-    yhat: Data format after calling model.predict.
-
-    Returns
-    -------
-    yhat: Data format after any processing.
-
-    """
-    return yhat
-
 def predict(data, model=load_model(), input_schema_path=os.path.join(os.path.dirname(os.path.realpath(__file__)), "input_schema.json")):
     """
     Returns prediction given the model and data to predict
@@ -165,5 +149,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.invoke(features))
-    return {'prediction': yhat}
+    output = model.invoke(features)
+    return {'output': output}