From ca20e527fc95183316e68b238e38960df2e33c33 Mon Sep 17 00:00:00 2001 From: Evan Lohn Date: Thu, 22 May 2025 18:13:18 -0700 Subject: [PATCH] fix tool calling for bedrock claude models (#4761) * fix tool calling for bedrock claude models * unit test * fix unit test --- backend/onyx/tools/utils.py | 5 +++ .../tests/unit/onyx/tools/test_tool_utils.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/backend/onyx/tools/utils.py b/backend/onyx/tools/utils.py index b05a3505121..833e11d5102 100644 --- a/backend/onyx/tools/utils.py +++ b/backend/onyx/tools/utils.py @@ -8,6 +8,7 @@ from onyx.db.connector import check_connectors_exist from onyx.db.document import check_docs_exist from onyx.db.models import LLMProvider from onyx.llm.llm_provider_options import ANTHROPIC_PROVIDER_NAME +from onyx.llm.llm_provider_options import BEDROCK_PROVIDER_NAME from onyx.llm.utils import find_model_obj from onyx.llm.utils import get_model_map from onyx.natural_language_processing.utils import BaseTokenizer @@ -35,6 +36,10 @@ def explicit_tool_calling_supported(model_provider: str, model_name: str) -> boo model_supports and model_provider != ANTHROPIC_PROVIDER_NAME and model_name not in litellm.anthropic_models + and ( + model_provider != BEDROCK_PROVIDER_NAME + or not any(name in model_name for name in litellm.anthropic_models) + ) ) diff --git a/backend/tests/unit/onyx/tools/test_tool_utils.py b/backend/tests/unit/onyx/tools/test_tool_utils.py index 6128c2e30f5..419bd7ca98e 100644 --- a/backend/tests/unit/onyx/tools/test_tool_utils.py +++ b/backend/tests/unit/onyx/tools/test_tool_utils.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest from onyx.llm.llm_provider_options import ANTHROPIC_PROVIDER_NAME +from onyx.llm.llm_provider_options import BEDROCK_PROVIDER_NAME from onyx.tools.utils import explicit_tool_calling_supported @@ -32,6 +33,40 @@ from onyx.tools.utils import explicit_tool_calling_supported # === Anthropic Scenarios (expected False due to base support being False) === # Provider is Anthropic, base model does NOT claim FC support (ANTHROPIC_PROVIDER_NAME, "claude-2.1", False, [], False), + # === Bedrock Scenarios === + # Bedrock provider with model name containing anthropic model name as substring -> False + ( + BEDROCK_PROVIDER_NAME, + "anthropic.claude-3-opus-20240229-v1:0", + True, + ["claude-3-opus-20240229"], + False, + ), + # Bedrock provider with model name containing different anthropic model name as substring -> False + ( + BEDROCK_PROVIDER_NAME, + "aws-anthropic-claude-3-haiku-20240307", + True, + ["claude-3-haiku-20240307"], + False, + ), + # Bedrock provider with model name NOT containing any anthropic model name as substring -> True + ( + BEDROCK_PROVIDER_NAME, + "amazon.titan-text-express-v1", + True, + ["claude-3-opus-20240229", "claude-3-haiku-20240307"], + True, + ), + # Bedrock provider with model name NOT containing any anthropic model + # name as substring, but base model doesn't support FC -> False + ( + BEDROCK_PROVIDER_NAME, + "amazon.titan-text-express-v1", + False, + ["claude-3-opus-20240229", "claude-3-haiku-20240307"], + False, + ), # === Non-Anthropic Scenarios === # Non-Anthropic provider, base model claims FC support -> True ("openai", "gpt-4o", True, [], True), @@ -73,6 +108,9 @@ def test_explicit_tool_calling_supported( We don't want to provide that list of tools because our UI doesn't support sequential tool calling yet for (a) and just looks bad for (b), so for now we just treat anthropic models as non-tool-calling. + + Additionally, for Bedrock provider, any model containing an anthropic model name as a + substring should also return False for the same reasons. """ mock_find_model_obj.return_value = { "supports_function_calling": mock_model_supports_fc