How to create a ChatModel with secrets

[felipebcs]

@dosu-bot, how can I create a ChatModel that extends ChatOpenAI, but adds custom secret fields and that still can be serialized?

[dosubot[bot]]

To create a ChatModel that extends ChatOpenAI, adds custom secret fields, and can still be serialized, you can follow this example:

from typing import Optional, Dict, Any, List
from pydantic import Field, SecretStr
from langchain_openai.chat_models.base import BaseChatOpenAI
from langchain_core.utils import secret_from_env

class CustomChatModel(BaseChatOpenAI):
    """Custom Chat Model extending ChatOpenAI with additional secret fields."""

    model_name: str = Field(alias="model")
    custom_api_key: Optional[SecretStr] = Field(
        alias="api_key",
        default_factory=secret_from_env("CUSTOM_API_KEY", default=None),
    )
    """Custom API key, automatically read from env variable `CUSTOM_API_KEY` if not provided."""

    custom_api_base: str = Field(default="https://api.custom.com/v1/")
    """Base URL path for API requests."""

    model_config = ConfigDict(
        populate_by_name=True,
    )

    @property
    def lc_secrets(self) -> Dict[str, str]:
        """A map of constructor argument names to secret ids."""
        return {"custom_api_key": "CUSTOM_API_KEY"}

    @classmethod
    def get_lc_namespace(cls) -> List[str]:
        """Get the namespace of the langchain object."""
        return ["custom_namespace", "chat_models"]

    @property
    def lc_attributes(self) -> Dict[str, Any]:
        """List of attribute names that should be included in the serialized kwargs."""
        attributes: Dict[str, Any] = {}

        if self.custom_api_base:
            attributes["custom_api_base"] = self.custom_api_base

        return attributes

    @classmethod
    def is_lc_serializable(cls) -> bool:
        """Return whether this model can be serialized by Langchain."""
        return True

    @property
    def _llm_type(self) -> str:
        """Return type of chat model."""
        return "custom-chat"

    # Additional methods and properties can be added here as needed

This class extends BaseChatOpenAI and includes custom secret fields (custom_api_key) and other attributes (custom_api_base). It implements the necessary properties and methods to ensure the model can be serialized. You can use the dumpd, dumps, load, and loads functions from the langchain_core.load module for serialization and deserialization ^[1][2].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

The custom_api_key is not shown in the serialized object because secrets, such as API keys, are intentionally withheld from serialized representations. This is controlled by the .lc_secrets attribute of the LangChain object. During deserialization, you can load the API key back onto the object using the secrets_map parameter ^[1].

To continue talking to Dosu, mention @dosu.

[felipebcs]

@dosu thanks, I'll test that. And how can I use those new variables (say custom_api_key) during initialization of the CustomChatModel so I can override values from the inherited class?

[dosubot[bot]]

To use new variables like custom_api_key during the initialization of the CustomChatModel and override values from the inherited class, you can pass them as arguments when creating an instance of the model. Here's how you can do it:

# Assuming you have the necessary imports and the CustomChatModel class defined

# Initialize the CustomChatModel with custom values
custom_chat_model = CustomChatModel(
    model_name="your_model_name",
    custom_api_key=SecretStr("your_custom_api_key"),
    custom_api_base="https://api.custom.com/v2/"
)

# Now, custom_chat_model is initialized with the specified values

In this example, custom_api_key and custom_api_base are provided during initialization, allowing you to override any default values or inherited values from the BaseChatOpenAI class. The SecretStr ensures that the custom_api_key is handled securely. This approach allows you to customize the behavior of your CustomChatModel instance as needed.

To continue talking to Dosu, mention @dosu.

[felipebcs]

@dosu but how would I get the value of custom_api_key and override a property on BaseChatOpenAI called openai_api_key?

[dosubot[bot]]

To override the openai_api_key property in the BaseChatOpenAI class with the value of custom_api_key, you can use a method similar to the validate_environment_override method from the ChatEverlyAI class. Here's how you can implement this:

from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from langchain_openai.chat_models.base import BaseChatOpenAI
from typing import Optional, Any

class CustomChatOpenAI(BaseChatOpenAI):
    custom_api_key: Optional[str] = None

    @classmethod
    def validate_environment_override(cls, values: dict) -> Any:
        """Override the openai_api_key with custom_api_key."""
        values["openai_api_key"] = convert_to_secret_str(
            get_from_dict_or_env(
                values,
                "custom_api_key",
                "CUSTOM_API_KEY_ENV_VAR",  # Replace with your environment variable name
            )
        )
        return values

This code snippet demonstrates how to retrieve the custom_api_key from either the provided values or an environment variable and use it to override the openai_api_key. Make sure to replace "CUSTOM_API_KEY_ENV_VAR" with the actual name of your environment variable if needed ^[1].

To continue talking to Dosu, mention @dosu.