Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 65 additions & 1 deletion nemoguardrails/actions/llm/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,70 @@ def __init__(self, inner_exception: Any):
self.inner_exception = inner_exception


def _infer_provider_from_module(llm: BaseLanguageModel) -> Optional[str]:
"""Infer provider name from the LLM's module path.

This function extracts the provider name from LangChain package naming conventions:
- langchain_openai -> openai
- langchain_anthropic -> anthropic
- langchain_google_genai -> google_genai
- langchain_nvidia_ai_endpoints -> nvidia_ai_endpoints
- langchain_community.chat_models.ollama -> ollama

For patched/wrapped classes, checks base classes as well.

Args:
llm: The LLM instance

Returns:
The inferred provider name, or None if it cannot be determined
"""
module = type(llm).__module__

if module.startswith("langchain_"):
package = module.split(".")[0]
provider = package.replace("langchain_", "")

if provider == "community":
parts = module.split(".")
if len(parts) >= 3:
provider = parts[-1]
return provider
Comment thread
Pouyanpi marked this conversation as resolved.
else:
return provider
Comment thread
Pouyanpi marked this conversation as resolved.

for base_class in type(llm).__mro__[1:]:
base_module = base_class.__module__
if base_module.startswith("langchain_"):
package = base_module.split(".")[0]
provider = package.replace("langchain_", "")

if provider == "community":
parts = base_module.split(".")
if len(parts) >= 3:
provider = parts[-1]
return provider
else:
return provider

return None


def get_llm_provider(llm: BaseLanguageModel) -> Optional[str]:
"""Get the provider name for an LLM instance by inferring from module path.

This function extracts the provider name from LangChain package naming conventions.
See _infer_provider_from_module for details on the inference logic.

Args:
llm: The LLM instance

Returns:
The provider name if it can be inferred, None otherwise
"""
return _infer_provider_from_module(llm)


def _infer_model_name(llm: BaseLanguageModel):
"""Helper to infer the model name based from an LLM instance.

Expand Down Expand Up @@ -126,7 +190,7 @@ def _setup_llm_call_info(
llm_call_info_var.set(llm_call_info)

llm_call_info.llm_model_name = model_name or _infer_model_name(llm)
llm_call_info.llm_provider_name = model_provider
llm_call_info.llm_provider_name = model_provider or _infer_provider_from_module(llm)


def _prepare_callbacks(
Expand Down
125 changes: 125 additions & 0 deletions tests/test_actions_llm_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from nemoguardrails.actions.llm.utils import _infer_provider_from_module


class MockOpenAILLM:
__module__ = "langchain_openai.chat_models"


class MockAnthropicLLM:
__module__ = "langchain_anthropic.chat_models"


class MockNVIDIALLM:
__module__ = "langchain_nvidia_ai_endpoints.chat_models"


class MockCommunityOllama:
__module__ = "langchain_community.chat_models.ollama"


class MockUnknownLLM:
__module__ = "some_custom_package.models"


class MockNVIDIAOriginal:
__module__ = "langchain_nvidia_ai_endpoints.chat_models"


class MockPatchedNVIDIA(MockNVIDIAOriginal):
__module__ = "nemoguardrails.llm.providers._langchain_nvidia_ai_endpoints_patch"


def test_infer_provider_openai():
llm = MockOpenAILLM()
provider = _infer_provider_from_module(llm)
assert provider == "openai"


def test_infer_provider_anthropic():
llm = MockAnthropicLLM()
provider = _infer_provider_from_module(llm)
assert provider == "anthropic"


def test_infer_provider_nvidia_ai_endpoints():
llm = MockNVIDIALLM()
provider = _infer_provider_from_module(llm)
assert provider == "nvidia_ai_endpoints"


def test_infer_provider_community_ollama():
llm = MockCommunityOllama()
provider = _infer_provider_from_module(llm)
assert provider == "ollama"


def test_infer_provider_unknown():
llm = MockUnknownLLM()
provider = _infer_provider_from_module(llm)
assert provider is None


def test_infer_provider_from_patched_class():
llm = MockPatchedNVIDIA()
provider = _infer_provider_from_module(llm)
assert provider == "nvidia_ai_endpoints"


def test_infer_provider_checks_base_classes():
class BaseOpenAI:
__module__ = "langchain_openai.chat_models"

class CustomWrapper(BaseOpenAI):
__module__ = "my_custom_wrapper.llms"

llm = CustomWrapper()
provider = _infer_provider_from_module(llm)
assert provider == "openai"


def test_infer_provider_multiple_inheritance():
class BaseNVIDIA:
__module__ = "langchain_nvidia_ai_endpoints.chat_models"

class Mixin:
__module__ = "some_mixin.utils"

class MultipleInheritance(Mixin, BaseNVIDIA):
__module__ = "custom_package.models"

llm = MultipleInheritance()
provider = _infer_provider_from_module(llm)
assert provider == "nvidia_ai_endpoints"


def test_infer_provider_deeply_nested_inheritance():
class Original:
__module__ = "langchain_anthropic.chat_models"

class Wrapper1(Original):
__module__ = "wrapper1.models"

class Wrapper2(Wrapper1):
__module__ = "wrapper2.models"

class Wrapper3(Wrapper2):
__module__ = "wrapper3.models"

llm = Wrapper3()
provider = _infer_provider_from_module(llm)
assert provider == "anthropic"