codekingpro/portable-devtools

1from __future__ import annotations2 3from typing import Any, Dict, Iterator, List, Mapping, Optional, cast4 5from langchain_core.callbacks import CallbackManagerForLLMRun6from langchain_core.language_models.chat_models import BaseChatModel7from langchain_core.messages import (8    AIMessage,9    AIMessageChunk,10    BaseMessage,11    FunctionMessage,12    HumanMessage,13    SystemMessage,14)15from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult16 17from langchain_community.llms.volcengine_maas import VolcEngineMaasBase18 19 20def _convert_message_to_dict(message: BaseMessage) -> dict:21    if isinstance(message, SystemMessage):22        message_dict = {"role": "system", "content": message.content}23    elif isinstance(message, HumanMessage):24        message_dict = {"role": "user", "content": message.content}25    elif isinstance(message, AIMessage):26        message_dict = {"role": "assistant", "content": message.content}27    elif isinstance(message, FunctionMessage):28        message_dict = {"role": "function", "content": message.content}29    else:30        raise ValueError(f"Got unknown type {message}")31    return message_dict32 33 34def convert_dict_to_message(_dict: Mapping[str, Any]) -> AIMessage:35    """Convert a dict to a message."""36 37    content = _dict.get("choice", {}).get("message", {}).get("content", "")38    return AIMessage(content=content)39 40 41class VolcEngineMaasChat(BaseChatModel, VolcEngineMaasBase):42    """Volc Engine Maas hosts a plethora of models.43 44    You can utilize these models through this class.45 46    To use, you should have the ``volcengine`` python package installed.47    and set access key and secret key by environment variable or direct pass those48    to this class.49    access key, secret key are required parameters which you could get help50    https://www.volcengine.com/docs/6291/6556851 52    In order to use them, it is necessary to install the 'volcengine' Python package.53    The access key and secret key must be set either via environment variables or54    passed directly to this class.55    access key and secret key are mandatory parameters for which assistance can be56    sought at https://www.volcengine.com/docs/6291/65568.57 58    The two methods are as follows:59    * Environment Variable60    Set the environment variables 'VOLC_ACCESSKEY' and 'VOLC_SECRETKEY' with your61    access key and secret key.62 63    * Pass Directly to Class64    Example:65        .. code-block:: python66 67            from langchain_community.llms import VolcEngineMaasLLM68            model = VolcEngineMaasChat(model="skylark-lite-public",69                                          volc_engine_maas_ak="your_ak",70                                          volc_engine_maas_sk="your_sk")71    """72 73    @property74    def _llm_type(self) -> str:75        """Return type of chat model."""76        return "volc-engine-maas-chat"77 78    @classmethod79    def is_lc_serializable(cls) -> bool:80        """Return whether this model can be serialized by Langchain."""81        return False82 83    @property84    def _identifying_params(self) -> Dict[str, Any]:85        return {86            **{"endpoint": self.endpoint, "model": self.model},87            **super()._identifying_params,88        }89 90    def _convert_prompt_msg_params(91        self,92        messages: List[BaseMessage],93        **kwargs: Any,94    ) -> Dict[str, Any]:95        model_req = {96            "model": {97                "name": self.model,98            }99        }100        if self.model_version is not None:101            model_req["model"]["version"] = self.model_version102        return {103            **model_req,104            "messages": [_convert_message_to_dict(message) for message in messages],105            "parameters": {**self._default_params, **kwargs},106        }107 108    def _stream(109        self,110        messages: List[BaseMessage],111        stop: Optional[List[str]] = None,112        run_manager: Optional[CallbackManagerForLLMRun] = None,113        **kwargs: Any,114    ) -> Iterator[ChatGenerationChunk]:115        if stop is not None:116            kwargs["stop"] = stop117        params = self._convert_prompt_msg_params(messages, **kwargs)118        for res in self.client.stream_chat(params):119            if res:120                msg = convert_dict_to_message(res)121                chunk = ChatGenerationChunk(message=AIMessageChunk(content=msg.content))122                if run_manager:123                    run_manager.on_llm_new_token(cast(str, msg.content), chunk=chunk)124                yield chunk125 126    def _generate(127        self,128        messages: List[BaseMessage],129        stop: Optional[List[str]] = None,130        run_manager: Optional[CallbackManagerForLLMRun] = None,131        **kwargs: Any,132    ) -> ChatResult:133        completion = ""134        if self.streaming:135            for chunk in self._stream(messages, stop, run_manager, **kwargs):136                completion += chunk.text137        else:138            if stop is not None:139                kwargs["stop"] = stop140            params = self._convert_prompt_msg_params(messages, **kwargs)141            res = self.client.chat(params)142            msg = convert_dict_to_message(res)143            completion = cast(str, msg.content)144 145        message = AIMessage(content=completion)146        return ChatResult(generations=[ChatGeneration(message=message)])147