fix tool calling for bedrock claude models (#4761)

* fix tool calling for bedrock claude models

* unit test

* fix unit test
This commit is contained in:
Evan Lohn
2025-05-22 18:13:18 -07:00
committed by GitHub
parent c8e65cce1e
commit ca20e527fc
2 changed files with 43 additions and 0 deletions

View File

@ -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)
)
)

View File

@ -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