Source code for langchain_community.llms.cerebriumai
import logging
from typing import Any, Dict, List, Mapping, Optional, cast
import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Extra, Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from langchain_community.llms.utils import enforce_stop_tokens
logger = logging.getLogger(__name__)
[docs]class CerebriumAI(LLM):
    """CerebriumAI large language models.
    To use, you should have the ``cerebrium`` python package installed.
    You should also have the environment variable ``CEREBRIUMAI_API_KEY``
    set with your API key or pass it as a named argument in the constructor.
    Any parameters that are valid to be passed to the call can be passed
    in, even if not explicitly saved on this class.
    Example:
        .. code-block:: python
            from langchain_community.llms import CerebriumAI
            cerebrium = CerebriumAI(endpoint_url="", cerebriumai_api_key="my-api-key")
    """
    endpoint_url: str = ""
    """model endpoint to use"""
    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
    """Holds any model parameters valid for `create` call not
    explicitly specified."""
    cerebriumai_api_key: Optional[SecretStr] = None
    class Config:
        """Configuration for this pydantic config."""
        extra = Extra.forbid
    @root_validator(pre=True)
    def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]: