codekingpro/portable-devtools

1"""Wrapper around Moonshot chat models."""2 3from typing import Dict4 5from langchain_core.utils import (6    convert_to_secret_str,7    get_from_dict_or_env,8    pre_init,9)10 11from langchain_community.chat_models import ChatOpenAI12from langchain_community.llms.moonshot import MOONSHOT_SERVICE_URL_BASE, MoonshotCommon13 14 15class MoonshotChat(MoonshotCommon, ChatOpenAI):  # type: ignore[misc]16    """Moonshot chat model integration.17 18    Setup:19        Install ``openai`` and set environment variables ``MOONSHOT_API_KEY``.20 21        .. code-block:: bash22 23            pip install openai24            export MOONSHOT_API_KEY="your-api-key"25 26    Key init args — completion params:27        model: str28            Name of Moonshot model to use.29        temperature: float30            Sampling temperature.31        max_tokens: Optional[int]32            Max number of tokens to generate.33 34    Key init args — client params:35        api_key: Optional[str]36            Moonshot API KEY. If not passed in will be read from env var MOONSHOT_API_KEY.37        api_base: Optional[str]38            Base URL for API requests.39 40    See full list of supported init args and their descriptions in the params section.41 42    Instantiate:43        .. code-block:: python44 45            from langchain_community.chat_models import MoonshotChat46 47            chat = MoonshotChat(48                temperature=0.5,49                api_key="your-api-key",50                model="moonshot-v1-8k",51                # api_base="...",52                # other params...53            )54 55    Invoke:56        .. code-block:: python57 58            messages = [59                ("system", "你是一名专业的翻译家，可以将用户的中文翻译为英文。"),60                ("human", "我喜欢编程。"),61            ]62            chat.invoke(messages)63 64        .. code-block:: python65 66            AIMessage(67                content='I like programming.',68                additional_kwargs={},69                response_metadata={70                    'token_usage': {71                        'completion_tokens': 5,72                        'prompt_tokens': 27,73                        'total_tokens': 3274                    },75                    'model_name': 'moonshot-v1-8k',76                    'system_fingerprint': None,77                    'finish_reason': 'stop',78                    'logprobs': None79                },80                id='run-71c03f4e-6628-41d5-beb6-d2559ae68266-0'81            )82 83    Stream:84        .. code-block:: python85 86            for chunk in chat.stream(messages):87                print(chunk)88 89        .. code-block:: python90 91            content='' additional_kwargs={} response_metadata={} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'92            content='I' additional_kwargs={} response_metadata={} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'93            content=' like' additional_kwargs={} response_metadata={} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'94            content=' programming' additional_kwargs={} response_metadata={} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'95            content='.' additional_kwargs={} response_metadata={} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'96            content='' additional_kwargs={} response_metadata={'finish_reason': 'stop'} id='run-80d77096-8b83-4c39-a84d-71d9c746da92'97 98        .. code-block:: python99 100            stream = chat.stream(messages)101            full = next(stream)102            for chunk in stream:103                full += chunk104            full105 106        .. code-block:: python107 108            AIMessageChunk(109                content='I like programming.',110                additional_kwargs={},111                response_metadata={'finish_reason': 'stop'},112                id='run-10c80976-7aa5-4ff7-ba3e-1251665557ef'113            )114 115    Async:116        .. code-block:: python117 118            await chat.ainvoke(messages)119 120            # stream:121            # async for chunk in chat.astream(messages):122            #    print(chunk)123 124            # batch:125            # await chat.abatch([messages])126 127        .. code-block:: python128 129            [AIMessage(content='I like programming.', additional_kwargs={}, response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 27, 'total_tokens': 32}, 'model_name': 'moonshot-v1-8k', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-2938b005-9204-4b9f-b273-1c3272fce9e5-0')]130 131    Response metadata132        .. code-block:: python133 134            ai_msg = chat.invoke(messages)135            ai_msg.response_metadata136 137        .. code-block:: python138 139            {140                'token_usage': {141                    'completion_tokens': 5,142                    'prompt_tokens': 27,143                    'total_tokens': 32144                },145                'model_name': 'moonshot-v1-8k',146                'system_fingerprint': None,147                'finish_reason': 'stop',148                'logprobs': None149            }150 151    """  # noqa: E501152 153    @pre_init154    def validate_environment(cls, values: Dict) -> Dict:155        """Validate that the environment is set up correctly."""156        values["moonshot_api_key"] = convert_to_secret_str(157            get_from_dict_or_env(158                values,159                ["moonshot_api_key", "api_key", "openai_api_key"],160                "MOONSHOT_API_KEY",161            )162        )163 164        try:165            import openai166 167        except ImportError:168            raise ImportError(169                "Could not import openai python package. "170                "Please install it with `pip install openai`."171            )172 173        client_params = {174            "api_key": values["moonshot_api_key"].get_secret_value(),175            "base_url": values["base_url"]176            if "base_url" in values177            else MOONSHOT_SERVICE_URL_BASE,178        }179 180        if not values.get("client"):181            values["client"] = openai.OpenAI(**client_params).chat.completions182        if not values.get("async_client"):183            values["async_client"] = openai.AsyncOpenAI(184                **client_params185            ).chat.completions186 187        return values188