31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138 | class NVIDIA(OpenAILike):
"""NVIDIA的API目录连接器。"""
_is_hosted: bool = PrivateAttr(True)
_mode: str = PrivateAttr("nvidia")
def __init__(
self,
model: str = DEFAULT_MODEL,
nvidia_api_key: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = BASE_URL,
max_tokens: Optional[int] = 1024,
**kwargs: Any,
) -> None:
"""初始化一个NVIDIA类的实例。
该类提供了与NVIDIA NIM的接口。默认情况下，它连接到托管的NIM，但您可以通过提供`base_url`来切换到本地NIM。
Args:
model (str, 可选): 用于NIM的模型。
nvidia_api_key (str, 可选): NVIDIA NIM的API密钥。默认为None。
api_key (str, 可选): 提供API密钥的另一种参数。默认为None。
base_url (str, 可选): NIM的基本URL。使用此参数切换到本地NIM。
max_tokens (int, 可选): 生成的最大标记数。默认为1024。
**kwargs: 附加关键字参数。
API密钥:
- 通过`NVIDIA_API_KEY`环境变量提供API密钥是推荐的方式。
Raises:
DeprecationWarning: 如果未为托管的NIM提供API密钥，则会发出警告。这将在版本0.2.0中变为错误。
"""
api_key = get_from_param_or_env(
"api_key",
nvidia_api_key or api_key,
"NVIDIA_API_KEY",
"NO_API_KEY_PROVIDED",
)
self._is_hosted = base_url in KNOWN_URLS
if self._is_hosted and api_key == "NO_API_KEY_PROVIDED":
warnings.warn(
"An API key is required for the hosted NIM. This will become an error in 0.2.0.",
)
super().__init__(
model=model,
api_key=api_key,
api_base=base_url,
max_tokens=max_tokens,
is_chat_model=True,
default_headers={"User-Agent": "llama-index-llms-nvidia"},
**kwargs,
)
@property
def available_models(self) -> List[Model]:
models = self._get_client().models.list().data
# only exclude models in hosted mode. in non-hosted mode, the administrator has control
# over the model name and may deploy an excluded name that will work.
if self._is_hosted:
exclude = {
"mistralai/mixtral-8x22b-v0.1", # not a /chat/completion endpoint
}
models = [model for model in models if model.id not in exclude]
return models
@classmethod
def class_name(cls) -> str:
return "NVIDIA"
@deprecated(
version="0.1.3",
reason="Will be removed in 0.2. Construct with `base_url` instead.",
)
def mode(
self,
mode: Optional[Literal["nvidia", "nim"]] = "nvidia",
*,
base_url: Optional[str] = None,
model: Optional[str] = None,
api_key: Optional[str] = None,
) -> "NVIDIA":
"""
已弃用：请改用NVIDIA(base_url="...")。
"""
if mode == "nim":
if not base_url:
raise ValueError("base_url is required for nim mode")
if mode == "nvidia":
api_key = get_from_param_or_env(
"api_key",
api_key,
"NVIDIA_API_KEY",
)
base_url = base_url or BASE_URL
self._mode = mode
if base_url:
self.api_base = base_url
if model:
self.model = model
if api_key:
self.api_key = api_key
return self
|