Mirrowel commited on
Commit
97ae153
·
1 Parent(s): b079814

chore(providers): remap nvidia provider name for litellm alignment

Browse files

This change ensures the 'nvidia' provider is correctly recognized by Litellm, which expects the name 'nvidia_nim'. This remapping is essential for proper integration and functionality when using Litellm with NVIDIA models.

src/rotator_library/providers/__init__.py CHANGED
@@ -28,6 +28,9 @@ def _register_providers():
28
  if isinstance(attribute, type) and issubclass(attribute, ProviderInterface) and attribute is not ProviderInterface:
29
  # The provider name is derived from the module name (e.g., 'openai_provider' -> 'openai')
30
  provider_name = module_name.replace("_provider", "")
 
 
 
31
  PROVIDER_PLUGINS[provider_name] = attribute
32
  #print(f"Registered provider: {provider_name}")
33
 
 
28
  if isinstance(attribute, type) and issubclass(attribute, ProviderInterface) and attribute is not ProviderInterface:
29
  # The provider name is derived from the module name (e.g., 'openai_provider' -> 'openai')
30
  provider_name = module_name.replace("_provider", "")
31
+ # Remap 'nvidia' to 'nvidia_nim' to align with litellm's provider name
32
+ if provider_name == "nvidia":
33
+ provider_name = "nvidia_nim"
34
  PROVIDER_PLUGINS[provider_name] = attribute
35
  #print(f"Registered provider: {provider_name}")
36