merge conflicts

Author: Ziqun Ye

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

@@ -3,28 +3,20 @@
 
 # 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 enum import Enum
 from typing import Any, Dict, List, Optional
 
 from langchain.llms.base import LLM
-from langchain.pydantic_v1 import BaseModel, root_validator, Field
+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 StrEnum(str, Enum):
-    """Enum with string members
-    https://docs.python.org/3.11/library/enum.html#enum.StrEnum
-    """
-
-    # Pydantic uses Python's standard enum classes to define choices.
-    # https://docs.pydantic.dev/latest/api/standard_library_types/#enum
-
-
-class BaseLLM(LLM):
+class BaseLLM(LLM, Serializable):
     """Base OCI LLM class. Contains common attributes."""
 
-    auth: dict = Field(default_factory=default_signer)
+    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."""
@@ -54,8 +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
@@ -65,15 +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."""
 
-    auth: dict = Field(default_factory=default_signer)
-    """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."""
-
-    compartment_id: str
+    compartment_id: str = None
     """Compartment ID of the caller."""
 
     endpoint_kwargs: Dict[str, Any] = {}
@@ -104,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/langchain/plugins/contant.py

@@ -3,8 +3,17 @@
 
 # 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 enum import Enum
+
+
+class StrEnum(str, Enum):
+    """Enum with string members
+    https://docs.python.org/3.11/library/enum.html#enum.StrEnum
+    """
+
+    # Pydantic uses Python's standard enum classes to define choices.
+    # https://docs.pydantic.dev/latest/api/standard_library_types/#enum
 
-from ads.llm.langchain.plugins.base import StrEnum
 
 DEFAULT_TIME_OUT = 300
 DEFAULT_CONTENT_TYPE_JSON = "application/json"

ads/llm/langchain/plugins/llm_gen_ai.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict, List, Optional
 
 from langchain.callbacks.manager import CallbackManagerForLLMRun
-from oci.util import to_dict
+
 from ads.llm.langchain.plugins.base import BaseLLM, GenerativeAiClientModel
 from ads.llm.langchain.plugins.contant import *
 
@@ -31,10 +31,10 @@ class GenerativeAI(GenerativeAiClientModel, BaseLLM):
 
     """
 
-    task: Task = Task.TEXT_GENERATION
+    task: str = "text_generation"
     """Indicates the task."""
 
-    model: Optional[str] = OCIGenerativeAIModel.COHERE_COMMAND
+    model: Optional[str] = "cohere.command"
     """Model name to use."""
 
     frequency_penalty: float = None
@@ -46,13 +46,13 @@ class GenerativeAI(GenerativeAiClientModel, BaseLLM):
     truncate: Optional[str] = None
     """Specify how the client handles inputs longer than the maximum token."""
 
-    length: str = LengthParam.AUTO
+    length: str = "AUTO"
     """Indicates the approximate length of the summary. """
 
-    format: str = FormatParam.PARAGRAPH
+    format: str = "PARAGRAPH"
     """Indicates the style in which the summary will be delivered - in a free form paragraph or in bullet points."""
 
-    extractiveness: str = ExtractivenessParam.AUTO
+    extractiveness: str = "AUTO"
     """Controls how close to the original text the summary is. High extractiveness summaries will lean towards reusing sentences verbatim, while low extractiveness summaries will tend to paraphrase more."""
 
     additional_command: str = ""
@@ -181,8 +181,8 @@ def _process_response(self, response: Any, num_generations: int = 1) -> str:
     def completion_with_retry(self, **kwargs: Any) -> Any:
         from oci.generative_ai.models import (
             GenerateTextDetails,
-            SummarizeTextDetails,
             OnDemandServingMode,
+            SummarizeTextDetails,
         )
 
         # TODO: Add retry logic for OCI

ads/llm/langchain/plugins/llm_md.py

@@ -9,6 +9,7 @@
 
 import requests
 from langchain.callbacks.manager import CallbackManagerForLLMRun
+
 from ads.llm.langchain.plugins.base import BaseLLM
 from ads.llm.langchain.plugins.contant import (
     DEFAULT_CONTENT_TYPE_JSON,
@@ -25,8 +26,8 @@ class ModelDeploymentLLM(BaseLLM):
     """The uri of the endpoint from the deployed Model Deployment model."""
 
     best_of: int = 1
-    """Generates best_of completions server-side and returns the "best" 
-    (the one with the highest log probability per token). 
+    """Generates best_of completions server-side and returns the "best"
+    (the one with the highest log probability per token).
     """
 
     @property
@@ -230,7 +231,7 @@ class ModelDeploymentVLLM(ModelDeploymentLLM):
     """Whether to use beam search instead of sampling."""
 
     ignore_eos: bool = False
-    """Whether to ignore the EOS token and continue generating tokens after 
+    """Whether to ignore the EOS token and continue generating tokens after
     the EOS token is generated."""
 
     logprobs: Optional[int] = None

ads/llm/serialize.py

@@ -23,7 +23,8 @@
 from opensearchpy.client import OpenSearch
 
 from ads.common.auth import default_signer
-from ads.llm import GenerativeAI, ModelDeploymentTGI, ModelDeploymentVLLM
+from ads.common.object_storage_details import ObjectStorageDetails
+from ads.llm import GenerativeAI, ModelDeploymentVLLM, ModelDeploymentTGI
 from ads.llm.chain import GuardrailSequence
 from ads.llm.guardrails.base import CustomGuardrailBase
 from ads.llm.patch import RunnableParallel, RunnableParallelSerializer
@@ -241,12 +242,11 @@ def ignore_unknown(self, node):
         None, _SafeLoaderIgnoreUnknown.ignore_unknown
     )
 
-    if uri.startswith("oci://"):
-        storage_options = default_signer()
-    else:
-        storage_options = {}
+    storage_options = default_signer() if ObjectStorageDetails.is_oci_path(uri) else {}
+
     with fsspec.open(uri, **storage_options) as f:
         config = yaml.load(f, Loader=_SafeLoaderIgnoreUnknown)
+
     return load(
         config, secrets_map=secrets_map, valid_namespaces=valid_namespaces, **kwargs
     )
@@ -275,6 +275,22 @@ def default(obj: Any) -> Any:
     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.
 
@@ -293,14 +309,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