mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
Fix #2688: Add operation_name parameter to AWS Bedrock embedder configuration
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
21
tests/utilities/test_embedding_configurator.py
Normal file
21
tests/utilities/test_embedding_configurator.py
Normal 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)
|
||||
Reference in New Issue
Block a user