diff --git a/src/crewai/utilities/embedding_configurator.py b/src/crewai/utilities/embedding_configurator.py index e523b60f0..7b6e6f27f 100644 --- a/src/crewai/utilities/embedding_configurator.py +++ b/src/crewai/utilities/embedding_configurator.py @@ -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) diff --git a/tests/utilities/test_embedding_configurator.py b/tests/utilities/test_embedding_configurator.py new file mode 100644 index 000000000..7c64c2457 --- /dev/null +++ b/tests/utilities/test_embedding_configurator.py @@ -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)