Source code for langchain_community.llms.octoai_endpoint

from typing import Any, Dict

from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env

from langchain_community.llms.openai import BaseOpenAI
from langchain_community.utils.openai import is_openai_v1

DEFAULT_BASE_URL = "https://text.octoai.run/v1/"
DEFAULT_MODEL = "codellama-7b-instruct"


class OctoAIEndpoint(BaseOpenAI):
    """OctoAI LLM终端点 - 兼容OpenAI。

    OctoAIEndpoint是一个用于与OctoAI Compute Service大型语言模型终端点交互的类。

    要使用，您应该设置环境变量``OCTOAI_API_TOKEN``为您的API令牌，或将其作为命名参数传递给构造函数。

    示例:
        .. code-block:: python

            from langchain_community.llms.octoai_endpoint  import OctoAIEndpoint

            llm = OctoAIEndpoint(
                model="llama-2-13b-chat-fp16",
                max_tokens=200,
                presence_penalty=0,
                temperature=0.1,
                top_p=0.9,
            )"""

    """传递给模型的关键字参数。"""
    octoai_api_base: str = Field(default=DEFAULT_BASE_URL)
    octoai_api_token: SecretStr = Field(default=None)
    model_name: str = Field(default=DEFAULT_MODEL)

    @classmethod
    def is_lc_serializable(cls) -> bool:
        return False

    @property
    def _invocation_params(self) -> Dict[str, Any]:
        """获取用于调用模型的参数。"""

        params: Dict[str, Any] = {
            "model": self.model_name,
            **self._default_params,
        }
        if not is_openai_v1():
            params.update(
                {
                    "api_key": self.octoai_api_token.get_secret_value(),
                    "api_base": self.octoai_api_base,
                }
            )

        return {**params, **super()._invocation_params}

    @property
    def _llm_type(self) -> str:
        """llm的返回类型。"""
        return "octoai_endpoint"

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥和Python包。"""
        values["octoai_api_base"] = get_from_dict_or_env(
            values,
            "octoai_api_base",
            "OCTOAI_API_BASE",
            default=DEFAULT_BASE_URL,
        )
        values["octoai_api_token"] = convert_to_secret_str(
            get_from_dict_or_env(values, "octoai_api_token", "OCTOAI_API_TOKEN")
        )
        values["model_name"] = get_from_dict_or_env(
            values,
            "model_name",
            "MODEL_NAME",
            default=DEFAULT_MODEL,
        )

        try:
            import openai

            if is_openai_v1():
                client_params = {
                    "api_key": values["octoai_api_token"].get_secret_value(),
                    "base_url": values["octoai_api_base"],
                }
                if not values.get("client"):
                    values["client"] = openai.OpenAI(**client_params).completions
                if not values.get("async_client"):
                    values["async_client"] = openai.AsyncOpenAI(
                        **client_params
                    ).completions
            else:
                values["openai_api_base"] = values["octoai_api_base"]
                values["openai_api_key"] = values["octoai_api_token"].get_secret_value()
                values["client"] = openai.Completion
        except ImportError:
            raise ImportError(
                "Could not import openai python package. "
                "Please install it with `pip install openai`."
            )

        if "endpoint_url" in values["model_kwargs"]:
            raise ValueError(
                "`endpoint_url` was deprecated, please use `octoai_api_base`."
            )

        return values