Fix/serialization of llm plugins (#445)

Author: qiuosier

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
@@ -68,11 +59,6 @@ class GenerativeAiClientModel(BaseModel):
     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 of the caller."""
 

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

@@ -18,6 +18,7 @@
 from langchain.load.serializable import Serializable
 
 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.chain import GuardrailSequence
 from ads.llm.guardrails.base import CustomGuardrailBase
@@ -115,12 +116,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
     )