Fix Vertex AI embeddings URL typo (publishers/goole -> publishers/google)

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-04-25 20:51:36 +00:00
parent b2969e9441
commit 1f2def2cbe
4 changed files with 141 additions and 10 deletions

View File

@@ -0,0 +1,36 @@
import pytest
from unittest.mock import patch, MagicMock
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
from crewai.utilities.embedding_functions import FixedGoogleVertexEmbeddingFunction
class TestEmbeddingConfigurator:
@pytest.fixture
def embedding_configurator(self):
return EmbeddingConfigurator()
def test_configure_vertexai(self, embedding_configurator):
with patch('crewai.utilities.embedding_functions.FixedGoogleVertexEmbeddingFunction') as mock_class:
mock_instance = MagicMock()
mock_class.return_value = mock_instance
config = {
"provider": "vertexai",
"config": {
"api_key": "test-key",
"model": "test-model",
"project_id": "test-project",
"region": "test-region"
}
}
result = embedding_configurator.configure_embedder(config)
mock_class.assert_called_once_with(
model_name="test-model",
api_key="test-key",
project_id="test-project",
region="test-region"
)
assert result == mock_instance

View File

@@ -0,0 +1,51 @@
import pytest
import requests
from unittest.mock import patch, MagicMock
from crewai.utilities.embedding_functions import FixedGoogleVertexEmbeddingFunction
class TestFixedGoogleVertexEmbeddingFunction:
@pytest.fixture
def embedding_function(self):
with patch('requests.post') as mock_post:
mock_response = MagicMock()
mock_response.json.return_value = {"predictions": [[0.1, 0.2, 0.3]]}
mock_post.return_value = mock_response
function = FixedGoogleVertexEmbeddingFunction(
model_name="test-model",
api_key="test-key"
)
yield function, mock_post
if hasattr(function, '_original_post'):
requests.post = function._original_post
def test_url_correction(self, embedding_function):
function, mock_post = embedding_function
typo_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/goole/models/test-model:predict"
expected_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/test-model:predict"
with patch.object(function, '_original_post') as mock_original_post:
mock_response = MagicMock()
mock_response.json.return_value = {"predictions": [[0.1, 0.2, 0.3]]}
mock_original_post.return_value = mock_response
response = function._patched_post(typo_url, json={})
mock_original_post.assert_called_once()
call_args = mock_original_post.call_args
assert call_args[0][0] == expected_url
def test_embedding_call(self, embedding_function):
function, mock_post = embedding_function
embeddings = function(["test text"])
mock_post.assert_called_once()
assert isinstance(embeddings, list)