codekingpro/portable-devtools

1from typing import Any, Dict, List, Optional2 3import requests4from langchain_core.callbacks import CallbackManagerForLLMRun5from langchain_core.language_models import LLM6from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env, pre_init7from pydantic import (8    BaseModel,9    ConfigDict,10    Field,11    SecretStr,12    model_validator,13)14 15from langchain_community.llms.utils import enforce_stop_tokens16 17MOONSHOT_SERVICE_URL_BASE = "https://api.moonshot.cn/v1"18 19 20class _MoonshotClient(BaseModel):21    """An API client that talks to the Moonshot server."""22 23    api_key: SecretStr24    """The API key to use for authentication."""25    base_url: str = MOONSHOT_SERVICE_URL_BASE26 27    def completion(self, request: Any) -> Any:28        headers = {"Authorization": f"Bearer {self.api_key.get_secret_value()}"}29        response = requests.post(30            f"{self.base_url}/chat/completions",31            headers=headers,32            json=request,33        )34        if not response.ok:35            raise ValueError(f"HTTP {response.status_code} error: {response.text}")36        return response.json()["choices"][0]["message"]["content"]37 38 39class MoonshotCommon(BaseModel):40    """Common parameters for Moonshot LLMs."""41 42    client: Any43    base_url: str = MOONSHOT_SERVICE_URL_BASE44    moonshot_api_key: Optional[SecretStr] = Field(default=None, alias="api_key")45    """Moonshot API key. Get it here: https://platform.moonshot.cn/console/api-keys"""46    model_name: str = Field(default="moonshot-v1-8k", alias="model")47    """Model name. Available models listed here: https://platform.moonshot.cn/pricing"""48    max_tokens: int = 102449    """Maximum number of tokens to generate."""50    temperature: float = 0.351    """Temperature parameter (higher values make the model more creative)."""52 53    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())54 55    @property56    def lc_secrets(self) -> dict:57        """A map of constructor argument names to secret ids.58 59        For example,60            {"moonshot_api_key": "MOONSHOT_API_KEY"}61        """62        return {"moonshot_api_key": "MOONSHOT_API_KEY"}63 64    @property65    def _default_params(self) -> Dict[str, Any]:66        """Get the default parameters for calling OpenAI API."""67        return {68            "model": self.model_name,69            "max_tokens": self.max_tokens,70            "temperature": self.temperature,71        }72 73    @property74    def _invocation_params(self) -> Dict[str, Any]:75        return {**{"model": self.model_name}, **self._default_params}76 77    @model_validator(mode="before")78    @classmethod79    def build_extra(cls, values: Dict[str, Any]) -> Any:80        """Build extra parameters.81        Override the superclass method, prevent the model parameter from being82        overridden.83        """84        return values85 86    @pre_init87    def validate_environment(cls, values: Dict) -> Dict:88        """Validate that api key and python package exists in environment."""89        values["moonshot_api_key"] = convert_to_secret_str(90            get_from_dict_or_env(values, "moonshot_api_key", "MOONSHOT_API_KEY")91        )92 93        values["client"] = _MoonshotClient(94            api_key=values["moonshot_api_key"],95            base_url=values["base_url"]96            if "base_url" in values97            else MOONSHOT_SERVICE_URL_BASE,98        )99        return values100 101    @property102    def _llm_type(self) -> str:103        """Return type of llm."""104        return "moonshot"105 106 107class Moonshot(MoonshotCommon, LLM):108    """Moonshot large language models.109 110    To use, you should have the environment variable ``MOONSHOT_API_KEY`` set with your111    API key. Referenced from https://platform.moonshot.cn/docs112 113    Example:114        .. code-block:: python115 116            from langchain_community.llms.moonshot import Moonshot117 118            moonshot = Moonshot(model="moonshot-v1-8k")119    """120 121    model_config = ConfigDict(122        populate_by_name=True,123    )124 125    def _call(126        self,127        prompt: str,128        stop: Optional[List[str]] = None,129        run_manager: Optional[CallbackManagerForLLMRun] = None,130        **kwargs: Any,131    ) -> str:132        request = self._invocation_params133        request["messages"] = [{"role": "user", "content": prompt}]134        request.update(kwargs)135        text = self.client.completion(request)136        if stop is not None:137            # This is required since the stop tokens138            # are not enforced by the model parameters139            text = enforce_stop_tokens(text, stop)140 141        return text142