resolve conflict

Author: Ziqun Ye

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

@@ -19,6 +19,7 @@
 from langchain.llms import loading
 from langchain.load import dumpd
 from langchain.load.load import load as lc_load
+from langchain.load.load import Reviver
 from langchain.load.serializable import Serializable
 from langchain.vectorstores import FAISS, OpenSearchVectorSearch
 from opensearchpy.client import OpenSearch
@@ -200,14 +201,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
@@ -268,6 +287,9 @@ 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.")

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,8 @@ 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))
-    return {'prediction': yhat}
+    output = model.invoke(features)
+    # Return the output as is if the output is a dictionary
+    if isinstance(output, dict):
+        return output
+    return {'output': output}