codekingpro/portable-devtools

1"""Wrapper around Konko AI's Completion API."""2 3import logging4import warnings5from typing import Any, Dict, List, Optional6 7from langchain_core.callbacks import (8    AsyncCallbackManagerForLLMRun,9    CallbackManagerForLLMRun,10)11from langchain_core.language_models.llms import LLM12from pydantic import ConfigDict, SecretStr, model_validator13 14from langchain_community.utils.openai import is_openai_v115 16logger = logging.getLogger(__name__)17 18 19class Konko(LLM):20    """Konko AI models.21 22    To use, you'll need an API key. This can be passed in as init param23    ``konko_api_key`` or set as environment variable ``KONKO_API_KEY``.24 25    Konko AI API reference: https://docs.konko.ai/reference/26    """27 28    base_url: str = "https://api.konko.ai/v1/completions"29    """Base inference API URL."""30    konko_api_key: SecretStr31    """Konko AI API key."""32    model: str33    """Model name. Available models listed here: 34       https://docs.konko.ai/reference/get_models35    """36    temperature: Optional[float] = None37    """Model temperature."""38    top_p: Optional[float] = None39    """Used to dynamically adjust the number of choices for each predicted token based 40        on the cumulative probabilities. A value of 1 will always yield the same 41        output. A temperature less than 1 favors more correctness and is appropriate 42        for question answering or summarization. A value greater than 1 introduces more 43        randomness in the output.44    """45    top_k: Optional[int] = None46    """Used to limit the number of choices for the next predicted word or token. It 47        specifies the maximum number of tokens to consider at each step, based on their 48        probability of occurrence. This technique helps to speed up the generation 49        process and can improve the quality of the generated text by focusing on the 50        most likely options.51    """52    max_tokens: Optional[int] = None53    """The maximum number of tokens to generate."""54    repetition_penalty: Optional[float] = None55    """A number that controls the diversity of generated text by reducing the 56        likelihood of repeated sequences. Higher values decrease repetition.57    """58    logprobs: Optional[int] = None59    """An integer that specifies how many top token log probabilities are included in 60        the response for each token generation step.61    """62 63    model_config = ConfigDict(64        extra="forbid",65    )66 67    @model_validator(mode="before")68    @classmethod69    def validate_environment(cls, values: Dict[str, Any]) -> Any:70        """Validate that python package exists in environment."""71        try:72            import konko73 74        except ImportError:75            raise ImportError(76                "Could not import konko python package. "77                "Please install it with `pip install konko`."78            )79        if not hasattr(konko, "_is_legacy_openai"):80            warnings.warn(81                "You are using an older version of the 'konko' package. "82                "Please consider upgrading to access new features"83                "including the completion endpoint."84            )85        return values86 87    def construct_payload(88        self,89        prompt: str,90        stop: Optional[List[str]] = None,91        **kwargs: Any,92    ) -> Dict[str, Any]:93        stop_to_use = stop[0] if stop and len(stop) == 1 else stop94        payload: Dict[str, Any] = {95            **self.default_params,96            "prompt": prompt,97            "stop": stop_to_use,98            **kwargs,99        }100        return {k: v for k, v in payload.items() if v is not None}101 102    @property103    def _llm_type(self) -> str:104        """Return type of model."""105        return "konko"106 107    @staticmethod108    def get_user_agent() -> str:109        from langchain_community import __version__110 111        return f"langchain/{__version__}"112 113    @property114    def default_params(self) -> Dict[str, Any]:115        return {116            "model": self.model,117            "temperature": self.temperature,118            "top_p": self.top_p,119            "top_k": self.top_k,120            "max_tokens": self.max_tokens,121            "repetition_penalty": self.repetition_penalty,122        }123 124    def _call(125        self,126        prompt: str,127        stop: Optional[List[str]] = None,128        run_manager: Optional[CallbackManagerForLLMRun] = None,129        **kwargs: Any,130    ) -> str:131        """Call out to Konko's text generation endpoint.132 133        Args:134            prompt: The prompt to pass into the model.135 136        Returns:137            The string generated by the model..138        """139        import konko140 141        payload = self.construct_payload(prompt, stop, **kwargs)142 143        try:144            if is_openai_v1():145                response = konko.completions.create(**payload)146            else:147                response = konko.Completion.create(**payload)148 149        except AttributeError:150            raise ValueError(151                "`konko` has no `Completion` attribute, this is likely "152                "due to an old version of the konko package. Try upgrading it "153                "with `pip install --upgrade konko`."154            )155 156        if is_openai_v1():157            output = response.choices[0].text158        else:159            output = response["choices"][0]["text"]160 161        return output162 163    async def _acall(164        self,165        prompt: str,166        stop: Optional[List[str]] = None,167        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,168        **kwargs: Any,169    ) -> str:170        """Asynchronously call out to Konko's text generation endpoint.171 172        Args:173            prompt: The prompt to pass into the model.174 175        Returns:176            The string generated by the model.177        """178        import konko179 180        payload = self.construct_payload(prompt, stop, **kwargs)181 182        try:183            if is_openai_v1():184                client = konko.AsyncKonko()185                response = await client.completions.create(**payload)186            else:187                response = await konko.Completion.acreate(**payload)188 189        except AttributeError:190            raise ValueError(191                "`konko` has no `Completion` attribute, this is likely "192                "due to an old version of the konko package. Try upgrading it "193                "with `pip install --upgrade konko`."194            )195 196        if is_openai_v1():197            output = response.choices[0].text198        else:199            output = response["choices"][0]["text"]200 201        return output202