codekingpro/portable-devtools

1from typing import Any, Dict, List, Mapping, Optional2 3import requests4from langchain_core.callbacks import CallbackManagerForLLMRun5from langchain_core.language_models.llms import LLM6from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env7from pydantic import ConfigDict, SecretStr, model_validator8 9from langchain_community.llms.utils import enforce_stop_tokens10 11 12class ForefrontAI(LLM):13    """ForefrontAI large language models.14 15    To use, you should have the environment variable ``FOREFRONTAI_API_KEY``16    set with your API key.17 18    Example:19        .. code-block:: python20 21            from langchain_community.llms import ForefrontAI22            forefrontai = ForefrontAI(endpoint_url="")23    """24 25    endpoint_url: str = ""26    """Model name to use."""27 28    temperature: float = 0.729    """What sampling temperature to use."""30 31    length: int = 25632    """The maximum number of tokens to generate in the completion."""33 34    top_p: float = 1.035    """Total probability mass of tokens to consider at each step."""36 37    top_k: int = 4038    """The number of highest probability vocabulary tokens to39    keep for top-k-filtering."""40 41    repetition_penalty: int = 142    """Penalizes repeated tokens according to frequency."""43 44    forefrontai_api_key: SecretStr45 46    base_url: Optional[str] = None47    """Base url to use, if None decides based on model name."""48 49    model_config = ConfigDict(50        extra="forbid",51    )52 53    @model_validator(mode="before")54    @classmethod55    def validate_environment(cls, values: Dict) -> Any:56        """Validate that api key exists in environment."""57        values["forefrontai_api_key"] = convert_to_secret_str(58            get_from_dict_or_env(values, "forefrontai_api_key", "FOREFRONTAI_API_KEY")59        )60        return values61 62    @property63    def _default_params(self) -> Mapping[str, Any]:64        """Get the default parameters for calling ForefrontAI API."""65        return {66            "temperature": self.temperature,67            "length": self.length,68            "top_p": self.top_p,69            "top_k": self.top_k,70            "repetition_penalty": self.repetition_penalty,71        }72 73    @property74    def _identifying_params(self) -> Mapping[str, Any]:75        """Get the identifying parameters."""76        return {**{"endpoint_url": self.endpoint_url}, **self._default_params}77 78    @property79    def _llm_type(self) -> str:80        """Return type of llm."""81        return "forefrontai"82 83    def _call(84        self,85        prompt: str,86        stop: Optional[List[str]] = None,87        run_manager: Optional[CallbackManagerForLLMRun] = None,88        **kwargs: Any,89    ) -> str:90        """Call out to ForefrontAI's complete endpoint.91 92        Args:93            prompt: The prompt to pass into the model.94            stop: Optional list of stop words to use when generating.95 96        Returns:97            The string generated by the model.98 99        Example:100            .. code-block:: python101 102                response = ForefrontAI("Tell me a joke.")103        """104        auth_value = f"Bearer {self.forefrontai_api_key.get_secret_value()}"105        response = requests.post(106            url=self.endpoint_url,107            headers={108                "Authorization": auth_value,109                "Content-Type": "application/json",110            },111            json={"text": prompt, **self._default_params, **kwargs},112        )113        response_json = response.json()114        text = response_json["result"][0]["completion"]115        if stop is not None:116            # I believe this is required since the stop tokens117            # are not enforced by the model parameters118            text = enforce_stop_tokens(text, stop)119        return text120