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
139
140 | class NVIDIA(OpenAILike):
"""NVIDIA's API Catalog Connector."""
_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:
"""
Initialize an instance of the NVIDIA class.
This class provides an interface to the NVIDIA NIM. By default, it connects to a hosted NIM,
but you can switch to an on-premises NIM by providing a `base_url`.
Args:
model (str, optional): The model to use for the NIM.
nvidia_api_key (str, optional): The API key for the NVIDIA NIM. Defaults to None.
api_key (str, optional): An alternative parameter for providing the API key. Defaults to None.
base_url (str, optional): The base URL for the NIM. Use this to switch to an on-premises NIM.
max_tokens (int, optional): The maximum number of tokens to generate. Defaults to 1024.
**kwargs: Additional keyword arguments.
API Keys:
- The recommended way to provide the API key is through the `NVIDIA_API_KEY` environment variable.
Raises:
DeprecationWarning: If an API key is not provided for a hosted NIM, a warning is issued. This will become an error in version 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":
"""
Deprecated: use NVIDIA(base_url="...") instead.
"""
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
|