codekingpro/portable-devtools

1from typing import Any, Dict, List, Optional2 3import requests4from langchain_core.callbacks import CallbackManagerForLLMRun5from langchain_core.language_models import LLM6from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env, pre_init7from pydantic import (8    BaseModel,9    ConfigDict,10    Field,11    SecretStr,12    model_validator,13)14 15from langchain_community.llms.utils import enforce_stop_tokens16 17SOLAR_SERVICE_URL_BASE = "https://api.upstage.ai/v1/solar"18SOLAR_SERVICE = "https://api.upstage.ai"19 20 21class _SolarClient(BaseModel):22    """An API client that talks to the Solar server."""23 24    api_key: SecretStr25    """The API key to use for authentication."""26    base_url: str = SOLAR_SERVICE_URL_BASE27 28    def completion(self, request: Any) -> Any:29        headers = {"Authorization": f"Bearer {self.api_key.get_secret_value()}"}30        response = requests.post(31            f"{self.base_url}/chat/completions",32            headers=headers,33            json=request,34        )35        if not response.ok:36            raise ValueError(f"HTTP {response.status_code} error: {response.text}")37        return response.json()["choices"][0]["message"]["content"]38 39 40class SolarCommon(BaseModel):41    """Common configuration for Solar LLMs."""42 43    _client: _SolarClient44    base_url: str = SOLAR_SERVICE_URL_BASE45    solar_api_key: Optional[SecretStr] = Field(default=None, alias="api_key")46    """Solar API key. Get it here: https://console.upstage.ai/services/solar"""47    model_name: str = Field(default="solar-mini", alias="model")48    """Model name. Available models listed here: https://console.upstage.ai/services/solar"""49    max_tokens: int = Field(default=1024)50    temperature: float = 0.351 52    model_config = ConfigDict(53        populate_by_name=True,54        arbitrary_types_allowed=True,55        extra="ignore",56        protected_namespaces=(),57    )58 59    @property60    def lc_secrets(self) -> dict:61        return {"solar_api_key": "SOLAR_API_KEY"}62 63    @property64    def _default_params(self) -> Dict[str, Any]:65        return {66            "model": self.model_name,67            "max_tokens": self.max_tokens,68            "temperature": self.temperature,69        }70 71    @property72    def _invocation_params(self) -> Dict[str, Any]:73        return {**{"model": self.model_name}, **self._default_params}74 75    @model_validator(mode="before")76    @classmethod77    def build_extra(cls, values: Dict[str, Any]) -> Any:78        return values79 80    @pre_init81    def validate_environment(cls, values: Dict) -> Dict:82        api_key = get_from_dict_or_env(values, "solar_api_key", "SOLAR_API_KEY")83        if api_key is None or len(api_key) == 0:84            raise ValueError("SOLAR_API_KEY must be configured")85 86        values["solar_api_key"] = convert_to_secret_str(api_key)87 88        if "base_url" not in values:89            values["base_url"] = SOLAR_SERVICE_URL_BASE90 91        if "base_url" in values and not values["base_url"].startswith(SOLAR_SERVICE):92            raise ValueError("base_url must match with: " + SOLAR_SERVICE)93 94        values["_client"] = _SolarClient(95            api_key=values["solar_api_key"], base_url=values["base_url"]96        )97        return values98 99    @property100    def _llm_type(self) -> str:101        return "solar"102 103 104class Solar(SolarCommon, LLM):105    """Solar large language models.106 107    To use, you should have the environment variable108    ``SOLAR_API_KEY`` set with your API key.109    Referenced from https://console.upstage.ai/services/solar110    """111 112    model_config = ConfigDict(113        populate_by_name=True,114    )115 116    def _call(117        self,118        prompt: str,119        stop: Optional[List[str]] = None,120        run_manager: Optional[CallbackManagerForLLMRun] = None,121        **kwargs: Any,122    ) -> str:123        request = self._invocation_params124        request["messages"] = [{"role": "user", "content": prompt}]125        request.update(kwargs)126        text = self._client.completion(request)127        if stop is not None:128            # This is required since the stop tokens129            # are not enforced by the model parameters130            text = enforce_stop_tokens(text, stop)131 132        return text133