Integrating the Yi family of models. (#24491)

Thank you for contributing to LangChain!

- [x] **PR title**: "community:add Yi LLM", "docs:add Yi Documentation"
                          
- [x] **PR message**: ***Delete this entire checklist*** and replace
with
- **Description:** This PR adds support for the Yi model to LangChain.
- **Dependencies:**
[langchain_core,requests,contextlib,typing,logging,json,langchain_community]
    - **Twitter handle:** 01.AI


- [x] **Add tests and docs**: I've added the corresponding documentation
to the relevant paths

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
Co-authored-by: isaac hershenson <ihershenson@hmc.edu>
This commit is contained in:
Haijian Wang
2024-07-27 01:57:33 +08:00
committed by GitHub
parent ad7581751f
commit cda3025ee1
9 changed files with 856 additions and 13 deletions

View File

@@ -640,12 +640,6 @@ def _import_yuan2() -> Type[BaseLLM]:
return Yuan2
def _import_you() -> Type[BaseLLM]:
from langchain_community.llms.you import You
return You
def _import_volcengine_maas() -> Type[BaseLLM]:
from langchain_community.llms.volcengine_maas import VolcEngineMaasLLM
@@ -658,6 +652,18 @@ def _import_sparkllm() -> Type[BaseLLM]:
return SparkLLM
def _import_you() -> Type[BaseLLM]:
from langchain_community.llms.you import You
return You
def _import_yi() -> Type[BaseLLM]:
from langchain_community.llms.yi import YiLLM
return YiLLM
def __getattr__(name: str) -> Any:
if name == "AI21":
return _import_ai21()
@@ -853,18 +859,20 @@ def __getattr__(name: str) -> Any:
return _import_yandex_gpt()
elif name == "Yuan2":
return _import_yuan2()
elif name == "You":
return _import_you()
elif name == "VolcEngineMaasLLM":
return _import_volcengine_maas()
elif name == "SparkLLM":
return _import_sparkllm()
elif name == "YiLLM":
return _import_yi()
elif name == "You":
return _import_you()
elif name == "type_to_cls_dict":
# for backwards compatibility
type_to_cls_dict: Dict[str, Type[BaseLLM]] = {
k: v() for k, v in get_type_to_cls_dict().items()
}
return type_to_cls_dict
elif name == "SparkLLM":
return _import_sparkllm()
else:
raise AttributeError(f"Could not find: {name}")
@@ -967,8 +975,9 @@ __all__ = [
"Writer",
"Xinference",
"YandexGPT",
"You",
"Yuan2",
"YiLLM",
"You",
]
@@ -1065,7 +1074,8 @@ def get_type_to_cls_dict() -> Dict[str, Callable[[], Type[BaseLLM]]]:
"qianfan_endpoint": _import_baidu_qianfan_endpoint,
"yandex_gpt": _import_yandex_gpt,
"yuan2": _import_yuan2,
"you": _import_you,
"VolcEngineMaasLLM": _import_volcengine_maas,
"SparkLLM": _import_sparkllm,
"yi": _import_yi,
"you": _import_you,
}

View File

@@ -0,0 +1,104 @@
from __future__ import annotations
import json
import logging
from typing import Any, Dict, List, Literal, Optional
import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Field, SecretStr
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from langchain_community.llms.utils import enforce_stop_tokens
logger = logging.getLogger(__name__)
class YiLLM(LLM):
"""Yi large language models."""
model: str = "yi-large"
temperature: float = 0.3
top_p: float = 0.95
timeout: int = 60
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
yi_api_key: Optional[SecretStr] = None
region: Literal["auto", "domestic", "international"] = "auto"
yi_api_url_domestic: str = "https://api.lingyiwanwu.com/v1/chat/completions"
yi_api_url_international: str = "https://api.01.ai/v1/chat/completions"
def __init__(self, **kwargs: Any):
kwargs["yi_api_key"] = convert_to_secret_str(
get_from_dict_or_env(kwargs, "yi_api_key", "YI_API_KEY")
)
super().__init__(**kwargs)
@property
def _default_params(self) -> Dict[str, Any]:
return {
"model": self.model,
"temperature": self.temperature,
"top_p": self.top_p,
**self.model_kwargs,
}
def _post(self, request: Any) -> Any:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.yi_api_key.get_secret_value()}", # type: ignore
}
urls = []
if self.region == "domestic":
urls = [self.yi_api_url_domestic]
elif self.region == "international":
urls = [self.yi_api_url_international]
else: # auto
urls = [self.yi_api_url_domestic, self.yi_api_url_international]
for url in urls:
try:
response = requests.post(
url,
headers=headers,
json=request,
timeout=self.timeout,
)
if response.status_code == 200:
parsed_json = json.loads(response.text)
return parsed_json["choices"][0]["message"]["content"]
elif (
response.status_code != 403
): # If not a permission error, raise immediately
response.raise_for_status()
except requests.RequestException as e:
if url == urls[-1]: # If this is the last URL to try
raise ValueError(f"An error has occurred: {e}")
else:
logger.warning(f"Failed to connect to {url}, trying next URL")
continue
raise ValueError("Failed to connect to all available URLs")
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
request = self._default_params
request["messages"] = [{"role": "user", "content": prompt}]
request.update(kwargs)
text = self._post(request)
if stop is not None:
text = enforce_stop_tokens(text, stop)
return text
@property
def _llm_type(self) -> str:
"""Return type of chat_model."""
return "yi-llm"