Fix #2688: Add operation_name parameter to AWS Bedrock embedder configuration

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-04-25 17:59:40 +00:00
parent b2969e9441
commit fe33f90bb2
2 changed files with 22 additions and 1 deletions

View File

@@ -156,7 +156,7 @@ class EmbeddingConfigurator:
)
# Allow custom model_name override with backwards compatibility
kwargs = {"session": config.get("session")}
kwargs = {"session": config.get("session"), "operation_name": "InvokeModel"}
if model_name is not None:
kwargs["model_name"] = model_name
return AmazonBedrockEmbeddingFunction(**kwargs)

View File

@@ -0,0 +1,21 @@
import unittest
from unittest.mock import patch, MagicMock
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
class TestEmbeddingConfigurator(unittest.TestCase):
@patch("chromadb.utils.embedding_functions.amazon_bedrock_embedding_function.AmazonBedrockEmbeddingFunction")
def test_configure_bedrock(self, mock_bedrock_embedder):
"""Test that the Bedrock embedder is configured correctly."""
config = {"session": MagicMock()}
model_name = "amazon.titan-embed-text-v1"
embedder = EmbeddingConfigurator()._configure_bedrock(config, model_name)
mock_bedrock_embedder.assert_called_once_with(
session=config["session"],
model_name=model_name,
operation_name="InvokeModel",
)
self.assertEqual(embedder, mock_bedrock_embedder.return_value)