mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-02 03:26:17 +00:00
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:
@@ -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,
|
||||
}
|
||||
|
104
libs/community/langchain_community/llms/yi.py
Normal file
104
libs/community/langchain_community/llms/yi.py
Normal 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"
|
Reference in New Issue
Block a user