feat: add partner chat provider names

Author: Pouyanpi

nemoguardrails/llm/providers/providers.py

@@ -25,7 +25,7 @@
 import importlib
 import logging
 import warnings
-from typing import Dict, List, Type
+from typing import Dict, List, Set, Type
 
 from langchain.chat_models.base import BaseChatModel
 from langchain_community import llms
@@ -82,6 +82,16 @@ def _discover_langchain_community_llm_providers():
     return type_to_cls_dict
 
 
+# this is needed as we perform the mapping in langchain_initializer.py
+_CUSTOM_CHAT_PROVIDERS = {"nim"}
+
+
+def _discover_langchain_partner_chat_providers() -> Set[str]:
+    from langchain.chat_models.base import _SUPPORTED_PROVIDERS
+
+    return _SUPPORTED_PROVIDERS | _CUSTOM_CHAT_PROVIDERS
+
+
 def _discover_langchain_community_chat_providers():
     """Creates a mapping from provider name to chat model class.
     The provider name is defined as the last segment of the module path.
@@ -159,6 +169,17 @@ def get_community_chat_provider_names() -> List[str]:
     return list(sorted(list(_chat_providers.keys())))
 
 
+def _get_all_chat_provider_names() -> List[str]:
+    """Consolidates all chat provider names."""
+
+    return list(_chat_providers.keys() | _discover_langchain_partner_chat_providers())
+
+
+def get_chat_provider_names() -> List[str]:
+    """Returns the list of supported chat providers."""
+    return list(sorted(_get_all_chat_provider_names()))
+
+
 def _get_text_completion_provider(provider_name: str) -> Type[BaseLLM]:
     if provider_name not in _llm_providers:
         raise RuntimeError(f"Could not find LLM provider '{provider_name}'")

tests/llm_providers/test_version_compatibility.py

@@ -24,8 +24,9 @@
     _chat_providers,
     _discover_langchain_community_chat_providers,
     _discover_langchain_community_llm_providers,
+    _discover_langchain_partner_chat_providers,
     _llm_providers,
-    _parse_version,
+    get_chat_provider_names,
     get_community_chat_provider_names,
     get_llm_provider_names,
 )
@@ -133,7 +134,7 @@
     "yi",
     "you",
 ]
-_CHAT_PROVIDERS_NAMES = [
+_COMMUNITY_CHAT_PROVIDERS_NAMES = [
     "azure_openai",
     "bedrock",
     "anthropic",
@@ -195,6 +196,26 @@
     "llamacpp",
     "yi",
 ]
+
+_PARTNER_CHAT_PROVIDERS_NAMES = {
+    "anthropic",
+    "azure_openai",
+    "bedrock",
+    "bedrock_converse",
+    "cohere",
+    "deepseek",
+    "fireworks",
+    "google_anthropic_vertex",
+    "google_genai",
+    "google_vertexai",
+    "groq",
+    "huggingface",
+    "mistralai",
+    "nim",
+    "ollama",
+    "openai",
+    "together",
+}
 # at some point we might care about certain providers
 CRITICAL_LLM_PROVIDERS = [
     "openai",
@@ -316,18 +337,40 @@ def test_provider_imports():
 
 
 def test_discover_langchain_community_chat_providers():
-    """Test that the function correctly discovers LangChain chat providers."""
+    """Test that the function correctly discovers LangChain community chat providers."""
+
     providers = _discover_langchain_community_chat_providers()
     chat_provider_names = get_community_chat_provider_names()
     assert set(chat_provider_names) == set(
         providers.keys()
     ), "it seems that we are registering a provider that is not in the LC community chat provider"
-    assert _CHAT_PROVIDERS_NAMES == list(providers.keys()), (
+    assert _COMMUNITY_CHAT_PROVIDERS_NAMES == list(providers.keys()), (
         "LangChain chat community providers may have changed. "
         "please investigate and update the test if necessary."
     )
 
 
+def test_dicsover_partner_chat_providers():
+    """Test that the function correctly discovers LangChain partner chat providers."""
+
+    partner_chat_providers = _discover_langchain_partner_chat_providers()
+    assert _PARTNER_CHAT_PROVIDERS_NAMES.issubset(partner_chat_providers), (
+        "LangChain partner chat providers may have changed. Update "
+        "_PARTNER_CHAT_PROVIDERS_NAMES to include all expected providers."
+    )
+    chat_providers = get_chat_provider_names()
+
+    assert partner_chat_providers.issubset(
+        chat_providers
+    ), "partner chat providers are not a subset of the list of chat providers"
+
+    if not partner_chat_providers == _PARTNER_CHAT_PROVIDERS_NAMES:
+        warnings.warn(
+            "LangChain partner chat providers may have changed. Update "
+            "_PARTNER_CHAT_PROVIDERS_NAMES to include all expected providers."
+        )
+
+
 def test_discover_langchain_community_llm_providers():
     providers = _discover_langchain_community_llm_providers()
     llm_provider_names = get_llm_provider_names()