codekingpro/portable-devtools

1"""Wrapper around Anyscale Endpoint"""2 3from typing import (4    Any,5    Dict,6    List,7    Mapping,8    Optional,9    Set,10)11 12from langchain_core.callbacks import (13    AsyncCallbackManagerForLLMRun,14    CallbackManagerForLLMRun,15)16from langchain_core.outputs import Generation, GenerationChunk, LLMResult17from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env, pre_init18from pydantic import Field, SecretStr19 20from langchain_community.llms.openai import (21    BaseOpenAI,22    acompletion_with_retry,23    completion_with_retry,24)25from langchain_community.utils.openai import is_openai_v126 27DEFAULT_BASE_URL = "https://api.endpoints.anyscale.com/v1"28DEFAULT_MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"29 30 31def update_token_usage(32    keys: Set[str], response: Dict[str, Any], token_usage: Dict[str, Any]33) -> None:34    """Update token usage."""35    _keys_to_use = keys.intersection(response["usage"])36    for _key in _keys_to_use:37        if _key not in token_usage:38            token_usage[_key] = response["usage"][_key]39        else:40            token_usage[_key] += response["usage"][_key]41 42 43def create_llm_result(44    choices: Any, prompts: List[str], token_usage: Dict[str, int], model_name: str45) -> LLMResult:46    """Create the LLMResult from the choices and prompts."""47    generations = []48    for i, _ in enumerate(prompts):49        choice = choices[i]50        generations.append(51            [52                Generation(53                    text=choice["message"]["content"],54                    generation_info=dict(55                        finish_reason=choice.get("finish_reason"),56                        logprobs=choice.get("logprobs"),57                    ),58                )59            ]60        )61    llm_output = {"token_usage": token_usage, "model_name": model_name}62    return LLMResult(generations=generations, llm_output=llm_output)63 64 65class Anyscale(BaseOpenAI):66    """Anyscale large language models.67 68    To use, you should have the environment variable ``ANYSCALE_API_KEY``set with your69    Anyscale Endpoint, or pass it as a named parameter to the constructor.70    To use with Anyscale Private Endpoint, please also set ``ANYSCALE_BASE_URL``.71 72    Example:73        .. code-block:: python74            from langchain_classic.llms import Anyscale75            anyscalellm = Anyscale(anyscale_api_key="ANYSCALE_API_KEY")76            # To leverage Ray for parallel processing77            @ray.remote(num_cpus=1)78            def send_query(llm, text):79                resp = llm.invoke(text)80                return resp81            futures = [send_query.remote(anyscalellm, text) for text in texts]82            results = ray.get(futures)83    """84 85    """Key word arguments to pass to the model."""86    anyscale_api_base: str = Field(default=DEFAULT_BASE_URL)87    anyscale_api_key: SecretStr = Field(default=SecretStr(""))88    model_name: str = Field(default=DEFAULT_MODEL)89 90    prefix_messages: List = Field(default_factory=list)91 92    @classmethod93    def is_lc_serializable(cls) -> bool:94        return False95 96    @pre_init97    def validate_environment(cls, values: Dict) -> Dict:98        """Validate that api key and python package exists in environment."""99        values["anyscale_api_base"] = get_from_dict_or_env(100            values,101            "anyscale_api_base",102            "ANYSCALE_API_BASE",103            default=DEFAULT_BASE_URL,104        )105        values["anyscale_api_key"] = convert_to_secret_str(106            get_from_dict_or_env(values, "anyscale_api_key", "ANYSCALE_API_KEY")107        )108        values["model_name"] = get_from_dict_or_env(109            values,110            "model_name",111            "MODEL_NAME",112            default=DEFAULT_MODEL,113        )114 115        try:116            import openai117 118            if is_openai_v1():119                client_params = {120                    "api_key": values["anyscale_api_key"].get_secret_value(),121                    "base_url": values["anyscale_api_base"],122                    # To do: future support123                    # "organization": values["openai_organization"],124                    # "timeout": values["request_timeout"],125                    # "max_retries": values["max_retries"],126                    # "default_headers": values["default_headers"],127                    # "default_query": values["default_query"],128                    # "http_client": values["http_client"],129                }130                if not values.get("client"):131                    values["client"] = openai.OpenAI(**client_params).completions132                if not values.get("async_client"):133                    values["async_client"] = openai.AsyncOpenAI(134                        **client_params135                    ).completions136            else:137                values["openai_api_base"] = values["anyscale_api_base"]138                values["openai_api_key"] = values["anyscale_api_key"].get_secret_value()139                values["client"] = openai.Completion140        except ImportError:141            raise ImportError(142                "Could not import openai python package. "143                "Please install it with `pip install openai`."144            )145        if values["streaming"] and values["n"] > 1:146            raise ValueError("Cannot stream results when n > 1.")147        if values["streaming"] and values["best_of"] > 1:148            raise ValueError("Cannot stream results when best_of > 1.")149 150        return values151 152    @property153    def _identifying_params(self) -> Mapping[str, Any]:154        """Get the identifying parameters."""155        return {156            **{"model_name": self.model_name},157            **super()._identifying_params,158        }159 160    @property161    def _invocation_params(self) -> Dict[str, Any]:162        """Get the parameters used to invoke the model."""163        openai_creds: Dict[str, Any] = {164            "model": self.model_name,165        }166        if not is_openai_v1():167            openai_creds.update(168                {169                    "api_key": self.anyscale_api_key.get_secret_value(),170                    "api_base": self.anyscale_api_base,171                }172            )173        return {**openai_creds, **super()._invocation_params}174 175    @property176    def _llm_type(self) -> str:177        """Return type of llm."""178        return "Anyscale LLM"179 180    def _generate(181        self,182        prompts: List[str],183        stop: Optional[List[str]] = None,184        run_manager: Optional[CallbackManagerForLLMRun] = None,185        **kwargs: Any,186    ) -> LLMResult:187        """Call out to OpenAI's endpoint with k unique prompts.188 189        Args:190            prompts: The prompts to pass into the model.191            stop: Optional list of stop words to use when generating.192 193        Returns:194            The full LLM output.195 196        Example:197            .. code-block:: python198 199                response = openai.generate(["Tell me a joke."])200        """201        # TODO: write a unit test for this202        params = self._invocation_params203        params = {**params, **kwargs}204        sub_prompts = self.get_sub_prompts(params, prompts, stop)205        choices = []206        token_usage: Dict[str, int] = {}207        # Get the token usage from the response.208        # Includes prompt, completion, and total tokens used.209        _keys = {"completion_tokens", "prompt_tokens", "total_tokens"}210        system_fingerprint: Optional[str] = None211        for _prompts in sub_prompts:212            if self.streaming:213                if len(_prompts) > 1:214                    raise ValueError("Cannot stream results with multiple prompts.")215 216                generation: Optional[GenerationChunk] = None217                for chunk in self._stream(_prompts[0], stop, run_manager, **kwargs):218                    if generation is None:219                        generation = chunk220                    else:221                        generation += chunk222                assert generation is not None223                choices.append(224                    {225                        "text": generation.text,226                        "finish_reason": generation.generation_info.get("finish_reason")227                        if generation.generation_info228                        else None,229                        "logprobs": generation.generation_info.get("logprobs")230                        if generation.generation_info231                        else None,232                    }233                )234            else:235                response = completion_with_retry(236                    ## THis is the ONLY change from BaseOpenAI()._generate()237                    self,238                    prompt=_prompts[0],239                    run_manager=run_manager,240                    **params,241                )242                if not isinstance(response, dict):243                    # V1 client returns the response in an PyDantic object instead of244                    # dict. For the transition period, we deep convert it to dict.245                    response = response.dict()246 247                choices.extend(response["choices"])248                update_token_usage(_keys, response, token_usage)249                if not system_fingerprint:250                    system_fingerprint = response.get("system_fingerprint")251        return self.create_llm_result(252            choices,253            prompts,254            params,255            token_usage,256            system_fingerprint=system_fingerprint,257        )258 259    async def _agenerate(260        self,261        prompts: List[str],262        stop: Optional[List[str]] = None,263        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,264        **kwargs: Any,265    ) -> LLMResult:266        """Call out to OpenAI's endpoint async with k unique prompts."""267        params = self._invocation_params268        params = {**params, **kwargs}269        sub_prompts = self.get_sub_prompts(params, prompts, stop)270        choices = []271        token_usage: Dict[str, int] = {}272        # Get the token usage from the response.273        # Includes prompt, completion, and total tokens used.274        _keys = {"completion_tokens", "prompt_tokens", "total_tokens"}275        system_fingerprint: Optional[str] = None276        for _prompts in sub_prompts:277            if self.streaming:278                if len(_prompts) > 1:279                    raise ValueError("Cannot stream results with multiple prompts.")280 281                generation: Optional[GenerationChunk] = None282                async for chunk in self._astream(283                    _prompts[0], stop, run_manager, **kwargs284                ):285                    if generation is None:286                        generation = chunk287                    else:288                        generation += chunk289                assert generation is not None290                choices.append(291                    {292                        "text": generation.text,293                        "finish_reason": generation.generation_info.get("finish_reason")294                        if generation.generation_info295                        else None,296                        "logprobs": generation.generation_info.get("logprobs")297                        if generation.generation_info298                        else None,299                    }300                )301            else:302                response = await acompletion_with_retry(303                    ## THis is the ONLY change from BaseOpenAI()._agenerate()304                    self,305                    prompt=_prompts[0],306                    run_manager=run_manager,307                    **params,308                )309                if not isinstance(response, dict):310                    response = response.dict()311                choices.extend(response["choices"])312                update_token_usage(_keys, response, token_usage)313        return self.create_llm_result(314            choices,315            prompts,316            params,317            token_usage,318            system_fingerprint=system_fingerprint,319        )320