Toggle Menu
Prev Up Next

# Source code for langchain_community.chat_models.pai_eas_endpoint

    import json
    import logging
    from typing import Any, AsyncIterator, Dict, List, Optional, cast

    import requests
    from langchain_core.callbacks import (
        AsyncCallbackManagerForLLMRun,
        CallbackManagerForLLMRun,
    )
    from langchain_core.language_models.chat_models import BaseChatModel
    from langchain_core.messages import (
        AIMessage,
        AIMessageChunk,
        BaseMessage,
        ChatMessage,
        HumanMessage,
        SystemMessage,
    )
    from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
    from langchain_core.pydantic_v1 import root_validator
    from langchain_core.utils import get_from_dict_or_env

    from langchain_community.llms.utils import enforce_stop_tokens

    logger = logging.getLogger(__name__)

    class PaiEasChatEndpoint(BaseChatModel):
        """Alibaba Cloud PAI-EAS LLM Service chat model API.

            To use, must have a deployed eas chat llm service on AliCloud. One can set the
        environment variable ``eas_service_url`` and ``eas_service_token`` set with your eas
        service url and service token.

        Example:
            .. code-block:: python

                from langchain_community.chat_models import PaiEasChatEndpoint
                eas_chat_endpoint = PaiEasChatEndpoint(
                    eas_service_url="your_service_url",
                    eas_service_token="your_service_token"
                )
        """

        """PAI-EAS Service URL"""
        eas_service_url: str

        """PAI-EAS Service TOKEN"""
        eas_service_token: str

        """PAI-EAS Service Infer Params"""
        max_new_tokens: Optional[int] = 512
        temperature: Optional[float] = 0.8
        top_p: Optional[float] = 0.1
        top_k: Optional[int] = 10
        do_sample: Optional[bool] = False
        use_cache: Optional[bool] = True
        stop_sequences: Optional[List[str]] = None

        """Enable stream chat mode."""
        streaming: bool = False

        """Key/value arguments to pass to the model. Reserved for future use"""
        model_kwargs: Optional[dict] = None

        version: Optional[str] = "2.0"

        timeout: Optional[int] = 5000

        @root_validator()
        def validate_environment(cls, values: Dict) -> Dict:
            """Validate that api key and python package exists in environment."""
            values["eas_service_url"] = get_from_dict_or_env(
                values, "eas_service_url", "EAS_SERVICE_URL"
            )
            values["eas_service_token"] = get_from_dict_or_env(
                values, "eas_service_token", "EAS_SERVICE_TOKEN"
            )

            return values

        @property
        def _identifying_params(self) -> Dict[str, Any]:
            """Get the identifying parameters."""
            _model_kwargs = self.model_kwargs or {}
            return {
                "eas_service_url": self.eas_service_url,
                "eas_service_token": self.eas_service_token,
                **{"model_kwargs": _model_kwargs},
            }

        @property
        def _llm_type(self) -> str:
            """Return type of llm."""
            return "pai_eas_chat_endpoint"

        @property
        def _default_params(self) -> Dict[str, Any]:
            """Get the default parameters for calling Cohere API."""
            return {
                "max_new_tokens": self.max_new_tokens,
                "temperature": self.temperature,
                "top_k": self.top_k,
                "top_p": self.top_p,
                "stop_sequences": [],
                "do_sample": self.do_sample,
                "use_cache": self.use_cache,
            }

        def _invocation_params(
            self, stop_sequences: Optional[List[str]], **kwargs: Any
        ) -> dict:
            params = self._default_params
            if self.model_kwargs:
                params.update(self.model_kwargs)
            if self.stop_sequences is not None and stop_sequences is not None:
                raise ValueError("`stop` found in both the input and default params.")
            elif self.stop_sequences is not None:
                params["stop"] = self.stop_sequences
            else:
                params["stop"] = stop_sequences
            return {**params, **kwargs}

        def format_request_payload(
            self, messages: List[BaseMessage], **model_kwargs: Any
        ) -> dict:
            prompt: Dict[str, Any] = {}
            user_content: List[str] = []
            assistant_content: List[str] = []

            for message in messages:
                """Converts message to a dict according to role"""
                content = cast(str, message.content)
                if isinstance(message, HumanMessage):
                    user_content = user_content + [content]
                elif isinstance(message, AIMessage):
                    assistant_content = assistant_content + [content]
                elif isinstance(message, SystemMessage):
                    prompt["system_prompt"] = content
                elif isinstance(message, ChatMessage) and message.role in [
                    "user",
                    "assistant",
                    "system",
                ]:
                    if message.role == "system":
                        prompt["system_prompt"] = content
                    elif message.role == "user":
                        user_content = user_content + [content]
                    elif message.role == "assistant":
                        assistant_content = assistant_content + [content]
                else:
                    supported = ",".join([role for role in ["user", "assistant", "system"]])
                    raise ValueError(
                        f"""Received unsupported role.
                        Supported roles for the LLaMa Foundation Model: {supported}"""
                    )
            prompt["prompt"] = user_content[len(user_content) - 1]
            history = [
                history_item
                for _, history_item in enumerate(zip(user_content[:-1], assistant_content))
            ]

            prompt["history"] = history

            return {**prompt, **model_kwargs}

        def _format_response_payload(
            self, output: bytes, stop_sequences: Optional[List[str]]
        ) -> str:
            """Formats response"""
            try:
                text = json.loads(output)["response"]
                if stop_sequences:
                    text = enforce_stop_tokens(text, stop_sequences)
                return text
            except Exception as e:
                if isinstance(e, json.decoder.JSONDecodeError):
                    return output.decode("utf-8")
                raise e

        def _generate(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> ChatResult:
            output_str = self._call(messages, stop=stop, run_manager=run_manager, **kwargs)
            message = AIMessage(content=output_str)
            generation = ChatGeneration(message=message)
            return ChatResult(generations=[generation])

        def _call(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> str:
            params = self._invocation_params(stop, **kwargs)

            request_payload = self.format_request_payload(messages, **params)
            response_payload = self._call_eas(request_payload)
            generated_text = self._format_response_payload(response_payload, params["stop"])

            if run_manager:
                run_manager.on_llm_new_token(generated_text)

            return generated_text

        def _call_eas(self, query_body: dict) -> Any:
            """Generate text from the eas service."""
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"{self.eas_service_token}",
            }

            # make request
            response = requests.post(
                self.eas_service_url, headers=headers, json=query_body, timeout=self.timeout
            )

            if response.status_code != 200:
                raise Exception(
                    f"Request failed with status code {response.status_code}"
                    f" and message {response.text}"
                )

            return response.text

        def _call_eas_stream(self, query_body: dict) -> Any:
            """Generate text from the eas service."""
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"{self.eas_service_token}",
            }

            # make request
            response = requests.post(
                self.eas_service_url, headers=headers, json=query_body, timeout=self.timeout
            )

            if response.status_code != 200:
                raise Exception(
                    f"Request failed with status code {response.status_code}"
                    f" and message {response.text}"
                )

            return response

        def _convert_chunk_to_message_message(
            self,
            chunk: str,
        ) -> AIMessageChunk:
            data = json.loads(chunk.encode("utf-8"))
            return AIMessageChunk(content=data.get("response", ""))

        async def _astream(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> AsyncIterator[ChatGenerationChunk]:
            params = self._invocation_params(stop, **kwargs)

            request_payload = self.format_request_payload(messages, **params)
            request_payload["use_stream_chat"] = True

            response = self._call_eas_stream(request_payload)
            for chunk in response.iter_lines(
                chunk_size=8192, decode_unicode=False, delimiter=b"\0"
            ):
                if chunk:
                    content = self._convert_chunk_to_message_message(chunk)

                    # identify stop sequence in generated text, if any
                    stop_seq_found: Optional[str] = None
                    for stop_seq in params["stop"]:
                        if stop_seq in content.content:
                            stop_seq_found = stop_seq

                    # identify text to yield
                    text: Optional[str] = None
                    if stop_seq_found:
                        content.content = content.content[
                            : content.content.index(stop_seq_found)
                        ]

                    # yield text, if any
                    if text:
                        cg_chunk = ChatGenerationChunk(message=content)
                        if run_manager:
                            await run_manager.on_llm_new_token(
                                cast(str, content.content), chunk=cg_chunk
                            )
                        yield cg_chunk

                    # break if stop sequence found
                    if stop_seq_found:
                        break