codekingpro/portable-devtools

1from typing import Any, Dict2 3from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env, pre_init4from pydantic import Field, SecretStr5 6from langchain_community.llms.openai import BaseOpenAI7from langchain_community.utils.openai import is_openai_v18 9DEFAULT_BASE_URL = "https://text.octoai.run/v1/"10DEFAULT_MODEL = "codellama-7b-instruct"11 12 13class OctoAIEndpoint(BaseOpenAI):14    """OctoAI LLM Endpoints - OpenAI compatible.15 16    OctoAIEndpoint is a class to interact with OctoAI Compute Service large17    language model endpoints.18 19    To use, you should have the environment variable ``OCTOAI_API_TOKEN`` set20    with your API token, or pass it as a named parameter to the constructor.21 22    Example:23        .. code-block:: python24 25            from langchain_community.llms.octoai_endpoint  import OctoAIEndpoint26 27            llm = OctoAIEndpoint(28                model="llama-2-13b-chat-fp16",29                max_tokens=200,30                presence_penalty=0,31                temperature=0.1,32                top_p=0.9,33            )34 35    """36 37    """Key word arguments to pass to the model."""38    octoai_api_base: str = Field(default=DEFAULT_BASE_URL)39    octoai_api_token: SecretStr = Field(default=SecretStr(""))40    model_name: str = Field(default=DEFAULT_MODEL)41 42    @classmethod43    def is_lc_serializable(cls) -> bool:44        return False45 46    @property47    def _invocation_params(self) -> Dict[str, Any]:48        """Get the parameters used to invoke the model."""49 50        params: Dict[str, Any] = {51            "model": self.model_name,52            **self._default_params,53        }54        if not is_openai_v1():55            params.update(56                {57                    "api_key": self.octoai_api_token.get_secret_value(),58                    "api_base": self.octoai_api_base,59                }60            )61 62        return {**params, **super()._invocation_params}63 64    @property65    def _llm_type(self) -> str:66        """Return type of llm."""67        return "octoai_endpoint"68 69    @pre_init70    def validate_environment(cls, values: Dict) -> Dict:71        """Validate that api key and python package exists in environment."""72        values["octoai_api_base"] = get_from_dict_or_env(73            values,74            "octoai_api_base",75            "OCTOAI_API_BASE",76            default=DEFAULT_BASE_URL,77        )78        values["octoai_api_token"] = convert_to_secret_str(79            get_from_dict_or_env(values, "octoai_api_token", "OCTOAI_API_TOKEN")80        )81        values["model_name"] = get_from_dict_or_env(82            values,83            "model_name",84            "MODEL_NAME",85            default=DEFAULT_MODEL,86        )87 88        try:89            import openai90 91            if is_openai_v1():92                client_params = {93                    "api_key": values["octoai_api_token"].get_secret_value(),94                    "base_url": values["octoai_api_base"],95                }96                if not values.get("client"):97                    values["client"] = openai.OpenAI(**client_params).completions98                if not values.get("async_client"):99                    values["async_client"] = openai.AsyncOpenAI(100                        **client_params101                    ).completions102            else:103                values["openai_api_base"] = values["octoai_api_base"]104                values["openai_api_key"] = values["octoai_api_token"].get_secret_value()105                values["client"] = openai.Completion106        except ImportError:107            raise ImportError(108                "Could not import openai python package. "109                "Please install it with `pip install openai`."110            )111 112        if "endpoint_url" in values["model_kwargs"]:113            raise ValueError(114                "`endpoint_url` was deprecated, please use `octoai_api_base`."115            )116 117        return values118