Source code for langchain_community.llms.forefrontai

from typing import Any, Dict, List, Mapping, Optional

import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Extra, 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


class ForefrontAI(LLM):
    """ForefrontAI大型语言模型。

    要使用，您应该设置环境变量``FOREFRONTAI_API_KEY``
    并使用您的API密钥。

    示例:
        .. code-block:: python

            from langchain_community.llms import ForefrontAI
            forefrontai = ForefrontAI(endpoint_url="")"""

    endpoint_url: str = ""
    """要使用的模型名称。"""

    temperature: float = 0.7
    """使用哪种采样温度。"""

    length: int = 256
    """生成完成的最大令牌数。"""

    top_p: float = 1.0
    """每一步需要考虑的标记的总概率质量。"""

    top_k: int = 40
    """保留用于top-k过滤的最高概率词汇标记的数量。"""

    repetition_penalty: int = 1
    """根据频率惩罚重复的标记。"""

    forefrontai_api_key: SecretStr

    base_url: Optional[str] = None
    """基础URL的使用，如果为None，则根据模型名称决定。"""

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥。"""
        values["forefrontai_api_key"] = convert_to_secret_str(
            get_from_dict_or_env(values, "forefrontai_api_key", "FOREFRONTAI_API_KEY")
        )
        return values

    @property
    def _default_params(self) -> Mapping[str, Any]:
        """获取调用ForefrontAI API的默认参数。"""
        return {
            "temperature": self.temperature,
            "length": self.length,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "repetition_penalty": self.repetition_penalty,
        }

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """获取识别参数。"""
        return {**{"endpoint_url": self.endpoint_url}, **self._default_params}

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

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """调用ForefrontAI的完整端点。

参数：
    prompt: 传递给模型的提示。
    stop: 生成时可选的停止词列表。

返回：
    模型生成的字符串。

示例：
    .. code-block:: python

        response = ForefrontAI("Tell me a joke.")
"""
        auth_value = f"Bearer {self.forefrontai_api_key.get_secret_value()}"
        response = requests.post(
            url=self.endpoint_url,
            headers={
                "Authorization": auth_value,
                "Content-Type": "application/json",
            },
            json={"text": prompt, **self._default_params, **kwargs},
        )
        response_json = response.json()
        text = response_json["result"][0]["completion"]
        if stop is not None:
            # I believe this is required since the stop tokens
            # are not enforced by the model parameters
            text = enforce_stop_tokens(text, stop)
        return text