From dc4304889bb67bc616a664a5dbaa468b93d75370 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Feb 2025 17:24:32 +0000 Subject: [PATCH] Refactor Google embedder error handling and expand test coverage Co-Authored-By: Joe Moura --- src/crewai/crew.py | 29 ++++--- src/crewai/utilities/constants.py | 1 + tests/knowledge/test_google_embedder.py | 105 +++++++++++++++++++++++- 3 files changed, 116 insertions(+), 19 deletions(-) diff --git a/src/crewai/crew.py b/src/crewai/crew.py index 2b5484b1c..ab8b59ed1 100644 --- a/src/crewai/crew.py +++ b/src/crewai/crew.py @@ -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] + [ diff --git a/src/crewai/utilities/constants.py b/src/crewai/utilities/constants.py index 9ff10f1d4..97c4232c2 100644 --- a/src/crewai/utilities/constants.py +++ b/src/crewai/utilities/constants.py @@ -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`" diff --git a/tests/knowledge/test_google_embedder.py b/tests/knowledge/test_google_embedder.py index 7a7fbf333..392c5a111 100644 --- a/tests/knowledge/test_google_embedder.py +++ b/tests/knowledge/test_google_embedder.py @@ -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" + )