Refactor Google embedder error handling and expand test coverage

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-02-28 17:24:32 +00:00
parent 7f89aeef84
commit dc4304889b
3 changed files with 116 additions and 19 deletions

View File

@@ -315,22 +315,9 @@ class Crew(BaseModel):
)
except ValueError as e:
# Check if the error is related to missing google-generativeai package
if "The Google Generative AI python package is not installed" in str(e):
self._logger.log(
"error",
"Google AI Studio embedder requires the google-generativeai package. "
"Please install it with `pip install google-generativeai`",
color="red",
)
else:
self._logger.log(
"warning", f"Failed to init knowledge: {e}", color="yellow"
)
self._handle_value_error(e)
except Exception as e:
self._logger.log(
"warning", f"Failed to init knowledge: {e}", color="yellow"
)
self._log_init_error(e)
return self
@model_validator(mode="after")
@@ -485,6 +472,18 @@ class Crew(BaseModel):
)
return self
def _handle_value_error(self, e: ValueError) -> None:
"""Handle ValueError exceptions during knowledge initialization."""
from crewai.utilities.constants import GOOGLE_EMBEDDER_PACKAGE_ERROR_MSG
if "The Google Generative AI python package is not installed" in str(e):
self._logger.log("error", GOOGLE_EMBEDDER_PACKAGE_ERROR_MSG, color="red")
else:
self._log_init_error(e)
def _log_init_error(self, e: Exception) -> None:
"""Log initialization errors."""
self._logger.log("warning", f"Failed to init knowledge: {e}", color="yellow")
@property
def key(self) -> str:
source = [agent.key for agent in self.agents] + [

View File

@@ -5,3 +5,4 @@ KNOWLEDGE_DIRECTORY = "knowledge"
MAX_LLM_RETRY = 3
MAX_FILE_NAME_LENGTH = 255
EMITTER_COLOR = "bold_blue"
GOOGLE_EMBEDDER_PACKAGE_ERROR_MSG = "Google AI Studio embedder requires the google-generativeai package. Please install it with `pip install google-generativeai`"

View File

@@ -1,7 +1,11 @@
from unittest.mock import MagicMock, patch
import pytest
from unittest.mock import patch, MagicMock
from crewai import Crew, Agent, Task
from crewai import Agent, Crew, Task
from crewai.knowledge.source.json_knowledge_source import JSONKnowledgeSource
from crewai.utilities.constants import GOOGLE_EMBEDDER_PACKAGE_ERROR_MSG
def test_google_embedder_missing_package():
"""Test that a helpful error message is displayed when google-generativeai is not installed."""
@@ -45,7 +49,100 @@ def test_google_embedder_missing_package():
# Verify that the error message was logged correctly
mock_log.assert_any_call(
"error",
"Google AI Studio embedder requires the google-generativeai package. "
"Please install it with `pip install google-generativeai`",
GOOGLE_EMBEDDER_PACKAGE_ERROR_MSG,
color="red"
)
def test_google_embedder_invalid_api_key():
"""Test that a warning is logged when an invalid API key is provided."""
# Create a simple agent and task
agent = Agent(
role="Test Agent",
goal="Test the knowledge component",
backstory="I am a test agent",
)
task = Task(
description="Test task",
expected_output="Test output",
agent=agent,
)
# Create a mock JSONKnowledgeSource
json_knowledge_source = MagicMock(spec=JSONKnowledgeSource)
# Mock the GoogleGenerativeAiEmbeddingFunction to raise an exception for invalid API key
with patch("chromadb.utils.embedding_functions.google_embedding_function.GoogleGenerativeAiEmbeddingFunction.__init__") as mock_init:
mock_init.side_effect = ValueError("Invalid API key")
# Mock the logger to capture the warning message
with patch("crewai.utilities.logger.Logger.log") as mock_log:
# Create a crew with Google embedder
crew = Crew(
agents=[agent],
tasks=[task],
verbose=True,
knowledge_sources=[json_knowledge_source],
embedder={
"provider": "google",
"config": {
"api_key": "invalid-api-key",
"model": 'models/embedding-001'
}
}
)
# Verify that the warning message was logged correctly
mock_log.assert_any_call(
"warning",
"Failed to init knowledge: Invalid API key",
color="yellow"
)
def test_google_embedder_invalid_model():
"""Test that a warning is logged when an invalid model is provided."""
# Create a simple agent and task
agent = Agent(
role="Test Agent",
goal="Test the knowledge component",
backstory="I am a test agent",
)
task = Task(
description="Test task",
expected_output="Test output",
agent=agent,
)
# Create a mock JSONKnowledgeSource
json_knowledge_source = MagicMock(spec=JSONKnowledgeSource)
# Mock the GoogleGenerativeAiEmbeddingFunction to raise an exception for invalid model
with patch("chromadb.utils.embedding_functions.google_embedding_function.GoogleGenerativeAiEmbeddingFunction.__init__") as mock_init:
mock_init.side_effect = ValueError("Invalid model name")
# Mock the logger to capture the warning message
with patch("crewai.utilities.logger.Logger.log") as mock_log:
# Create a crew with Google embedder
crew = Crew(
agents=[agent],
tasks=[task],
verbose=True,
knowledge_sources=[json_knowledge_source],
embedder={
"provider": "google",
"config": {
"api_key": "fake-api-key",
"model": 'invalid-model'
}
}
)
# Verify that the warning message was logged correctly
mock_log.assert_any_call(
"warning",
"Failed to init knowledge: Invalid model name",
color="yellow"
)