Mirrowel commited on
Commit
5538104
·
1 Parent(s): 72525d5

feat: Implement NvidiaProvider for fetching models from NVIDIA API

Browse files
src/rotator_library/providers/nvidia_provider.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import httpx
2
+ import logging
3
+ from typing import List
4
+ from .provider_interface import ProviderInterface
5
+
6
+ lib_logger = logging.getLogger('rotator_library')
7
+ lib_logger.propagate = False # Ensure this logger doesn't propagate to root
8
+ if not lib_logger.handlers:
9
+ lib_logger.addHandler(logging.NullHandler())
10
+
11
+ class NvidiaProvider(ProviderInterface):
12
+ """
13
+ Provider implementation for the NVIDIA API.
14
+ """
15
+ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]:
16
+ """
17
+ Fetches the list of available models from the NVIDIA API.
18
+ """
19
+ try:
20
+ response = await client.get(
21
+ "https://integrate.api.nvidia.com/v1/models",
22
+ headers={"Authorization": f"Bearer {api_key}"}
23
+ )
24
+ response.raise_for_status()
25
+ return [f"nvidia_nim/{model['id']}" for model in response.json().get("data", [])]
26
+ except httpx.RequestError as e:
27
+ lib_logger.error(f"Failed to fetch NVIDIA models: {e}")
28
+ return []