mirror of
https://github.com/crewAIInc/crewAI.git
synced 2025-12-16 04:18:35 +00:00
Compare commits
2 Commits
1.2.1
...
devin/1740
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc4304889b | ||
|
|
7f89aeef84 |
@@ -314,10 +314,10 @@ class Crew(BaseModel):
|
||||
collection_name="crew",
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
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")
|
||||
@@ -472,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] + [
|
||||
|
||||
@@ -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`"
|
||||
|
||||
148
tests/knowledge/test_google_embedder.py
Normal file
148
tests/knowledge/test_google_embedder.py
Normal file
@@ -0,0 +1,148 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
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."""
|
||||
# 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 the import error
|
||||
with patch("chromadb.utils.embedding_functions.google_embedding_function.GoogleGenerativeAiEmbeddingFunction.__init__") as mock_init:
|
||||
mock_init.side_effect = ValueError("The Google Generative AI python package is not installed. Please install it with `pip install google-generativeai`")
|
||||
|
||||
# Mock the logger to capture the error 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": 'models/embedding-001'
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Verify that the error message was logged correctly
|
||||
mock_log.assert_any_call(
|
||||
"error",
|
||||
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"
|
||||
)
|
||||
Reference in New Issue
Block a user