andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-10 16:48:30 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Add ChromaDBRequiredError class and improve error handling

Browse Source 
Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
Devin AI
2025-05-30 09:10:35 +00:00
parent 7b129fc847
commit f0b1cc23f4
4 changed files with 66 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
src/crewai/knowledge/storage/knowledge_storage.py
View File

@@ -24,6 +24,7 @@ from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage
from crewai.utilities import EmbeddingConfigurator from crewai.utilities import EmbeddingConfigurator
from crewai.utilities.chromadb import sanitize_collection_name from crewai.utilities.chromadb import sanitize_collection_name
from crewai.utilities.constants import KNOWLEDGE_DIRECTORY from crewai.utilities.constants import KNOWLEDGE_DIRECTORY
from crewai.utilities.errors import ChromaDBRequiredError
from crewai.utilities.logger import Logger from crewai.utilities.logger import Logger
from crewai.utilities.paths import db_storage_path from crewai.utilities.paths import db_storage_path
@@ -71,10 +72,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
score_threshold: float = 0.35, score_threshold: float = 0.35,
) -> List[Dict[str, Any]]: ) -> List[Dict[str, Any]]:
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
with suppress_logging(): with suppress_logging():
if self.collection: if self.collection:
@@ -99,10 +97,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
def initialize_knowledge_storage(self): def initialize_knowledge_storage(self):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
base_path = os.path.join(db_storage_path(), "knowledge") base_path = os.path.join(db_storage_path(), "knowledge")
try: try:
@@ -129,17 +124,11 @@ class KnowledgeStorage(BaseKnowledgeStorage):
except Exception: except Exception:
raise Exception("Failed to create or get collection") raise Exception("Failed to create or get collection")
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
def reset(self): def reset(self):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
base_path = os.path.join(db_storage_path(), KNOWLEDGE_DIRECTORY) base_path = os.path.join(db_storage_path(), KNOWLEDGE_DIRECTORY)
try: try:
@@ -154,10 +143,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
self.app = None self.app = None
self.collection = None self.collection = None
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
def save( def save(
self, self,
@@ -165,10 +151,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
metadata: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, metadata: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
): ):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
if not self.collection: if not self.collection:
raise Exception("Collection not initialized") raise Exception("Collection not initialized")
@@ -210,10 +193,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
ids=filtered_ids, ids=filtered_ids,
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
except Exception as e: except Exception as e:
if HAS_CHROMADB and isinstance(e, chromadb.errors.InvalidDimensionException): if HAS_CHROMADB and isinstance(e, chromadb.errors.InvalidDimensionException):
Logger(verbose=True).log( Logger(verbose=True).log(
@@ -232,10 +212,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
def _create_default_embedding_function(self): def _create_default_embedding_function(self):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.openai_embedding_function import ( from chromadb.utils.embedding_functions.openai_embedding_function import (
@@ -246,10 +223,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small" api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("knowledge storage")
"ChromaDB is required for knowledge storage features. "
"Please install it with 'pip install crewai[storage]'"
)
def _set_embedder_config(self, embedder: Optional[Dict[str, Any]] = None) -> None: def _set_embedder_config(self, embedder: Optional[Dict[str, Any]] = None) -> None:
"""Set the embedding configuration for the knowledge storage. """Set the embedding configuration for the knowledge storage.

23
src/crewai/memory/storage/rag_storage.py
View File

@@ -4,7 +4,7 @@ import logging
import os import os
import shutil import shutil
import uuid import uuid
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional
try: try:
from chromadb.api import ClientAPI from chromadb.api import ClientAPI
@@ -16,6 +16,7 @@ except ImportError:
from crewai.memory.storage.base_rag_storage import BaseRAGStorage from crewai.memory.storage.base_rag_storage import BaseRAGStorage
from crewai.utilities import EmbeddingConfigurator from crewai.utilities import EmbeddingConfigurator
from crewai.utilities.constants import MAX_FILE_NAME_LENGTH from crewai.utilities.constants import MAX_FILE_NAME_LENGTH
from crewai.utilities.errors import ChromaDBRequiredError
from crewai.utilities.paths import db_storage_path from crewai.utilities.paths import db_storage_path
@@ -66,10 +67,7 @@ class RAGStorage(BaseRAGStorage):
def _initialize_app(self): def _initialize_app(self):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("memory storage")
"ChromaDB is required for memory storage features. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
import chromadb import chromadb
@@ -92,10 +90,7 @@ class RAGStorage(BaseRAGStorage):
name=self.type, embedding_function=self.embedder_config name=self.type, embedding_function=self.embedder_config
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("memory storage")
"ChromaDB is required for memory storage features. "
"Please install it with 'pip install crewai[storage]'"
)
def _sanitize_role(self, role: str) -> str: def _sanitize_role(self, role: str) -> str:
""" """
@@ -183,10 +178,7 @@ class RAGStorage(BaseRAGStorage):
def _create_default_embedding_function(self): def _create_default_embedding_function(self):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("memory storage")
"ChromaDB is required for memory storage features. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.openai_embedding_function import ( from chromadb.utils.embedding_functions.openai_embedding_function import (
@@ -197,7 +189,4 @@ class RAGStorage(BaseRAGStorage):
api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small" api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("memory storage")
"ChromaDB is required for memory storage features. "
"Please install it with 'pip install crewai[storage]'"
)

126
src/crewai/utilities/embedding_configurator.py
View File

@@ -1,20 +1,28 @@
import os import os
from typing import Any, Dict, Optional, Union, cast, Protocol from typing import Any, Dict, Optional, cast, Protocol, TypeVar, Sequence
from crewai.utilities.errors import ChromaDBRequiredError
T = TypeVar('T')
try: try:
from chromadb import Documents, EmbeddingFunction, Embeddings from chromadb import Documents, EmbeddingFunction as ChromaEmbeddingFunction, Embeddings
from chromadb.api.types import validate_embedding_function from chromadb.api.types import validate_embedding_function
HAS_CHROMADB = True HAS_CHROMADB = True
EmbeddingFunction = ChromaEmbeddingFunction
except ImportError: except ImportError:
HAS_CHROMADB = False HAS_CHROMADB = False
class EmbeddingFunction(Protocol): class EmbeddingFunction(Protocol[T]):
def __call__(self, input: Any) -> Any: ... """Protocol for embedding functions when ChromaDB is not available."""
def __call__(self, input: Sequence[str]) -> Sequence[Sequence[float]]: ...
Documents = Any Documents = Any
Embeddings = Any Embeddings = Any
def validate_embedding_function(func: Any) -> None: def validate_embedding_function(func: Any) -> None:
"""Stub for validate_embedding_function when ChromaDB is not available."""
pass pass
@@ -40,10 +48,7 @@ class EmbeddingConfigurator:
) -> EmbeddingFunction: ) -> EmbeddingFunction:
"""Configures and returns an embedding function based on the provided config.""" """Configures and returns an embedding function based on the provided config."""
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
if embedder_config is None: if embedder_config is None:
return self._create_default_embedding_function() return self._create_default_embedding_function()
@@ -67,10 +72,7 @@ class EmbeddingConfigurator:
@staticmethod @staticmethod
def _create_default_embedding_function(): def _create_default_embedding_function():
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.openai_embedding_function import ( from chromadb.utils.embedding_functions.openai_embedding_function import (
@@ -81,18 +83,12 @@ class EmbeddingConfigurator:
api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small" api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_openai(config, model_name): def _configure_openai(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.openai_embedding_function import ( from chromadb.utils.embedding_functions.openai_embedding_function import (
@@ -111,18 +107,12 @@ class EmbeddingConfigurator:
organization_id=config.get("organization_id", None), organization_id=config.get("organization_id", None),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_azure(config, model_name): def _configure_azure(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.openai_embedding_function import ( from chromadb.utils.embedding_functions.openai_embedding_function import (
@@ -141,18 +131,12 @@ class EmbeddingConfigurator:
organization_id=config.get("organization_id"), organization_id=config.get("organization_id"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_ollama(config, model_name): def _configure_ollama(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.ollama_embedding_function import ( from chromadb.utils.embedding_functions.ollama_embedding_function import (
@@ -164,18 +148,12 @@ class EmbeddingConfigurator:
model_name=model_name, model_name=model_name,
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_vertexai(config, model_name): def _configure_vertexai(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.google_embedding_function import ( from chromadb.utils.embedding_functions.google_embedding_function import (
@@ -189,18 +167,12 @@ class EmbeddingConfigurator:
region=config.get("region"), region=config.get("region"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_google(config, model_name): def _configure_google(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.google_embedding_function import ( from chromadb.utils.embedding_functions.google_embedding_function import (
@@ -213,18 +185,12 @@ class EmbeddingConfigurator:
task_type=config.get("task_type"), task_type=config.get("task_type"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_cohere(config, model_name): def _configure_cohere(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.cohere_embedding_function import ( from chromadb.utils.embedding_functions.cohere_embedding_function import (
@@ -236,18 +202,12 @@ class EmbeddingConfigurator:
api_key=config.get("api_key"), api_key=config.get("api_key"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_voyageai(config, model_name): def _configure_voyageai(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.voyageai_embedding_function import ( from chromadb.utils.embedding_functions.voyageai_embedding_function import (
@@ -259,18 +219,12 @@ class EmbeddingConfigurator:
api_key=config.get("api_key"), api_key=config.get("api_key"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_bedrock(config, model_name): def _configure_bedrock(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.amazon_bedrock_embedding_function import ( from chromadb.utils.embedding_functions.amazon_bedrock_embedding_function import (
@@ -283,18 +237,12 @@ class EmbeddingConfigurator:
kwargs["model_name"] = model_name kwargs["model_name"] = model_name
return AmazonBedrockEmbeddingFunction(**kwargs) return AmazonBedrockEmbeddingFunction(**kwargs)
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_huggingface(config, model_name): def _configure_huggingface(config, model_name):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
try: try:
from chromadb.utils.embedding_functions.huggingface_embedding_function import ( from chromadb.utils.embedding_functions.huggingface_embedding_function import (
@@ -305,10 +253,7 @@ class EmbeddingConfigurator:
url=config.get("api_url"), url=config.get("api_url"),
) )
except ImportError: except ImportError:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
@staticmethod @staticmethod
def _configure_watson(config, model_name): def _configure_watson(config, model_name):
@@ -352,10 +297,7 @@ class EmbeddingConfigurator:
@staticmethod @staticmethod
def _configure_custom(config): def _configure_custom(config):
if not HAS_CHROMADB: if not HAS_CHROMADB:
raise ImportError( raise ChromaDBRequiredError("embedding functionality")
"ChromaDB is required for embedding functionality. "
"Please install it with 'pip install crewai[storage]'"
)
custom_embedder = config.get("embedder") custom_embedder = config.get("embedder")
if isinstance(custom_embedder, EmbeddingFunction): if isinstance(custom_embedder, EmbeddingFunction):

16
src/crewai/utilities/errors/__init__.py Normal file
View File

@@ -0,0 +1,16 @@
"""Custom error classes for CrewAI."""
class ChromaDBRequiredError(ImportError):
"""Error raised when ChromaDB is required but not installed."""
def __init__(self, feature: str):
"""Initialize the error with a specific feature name.
Args:
feature: The name of the feature that requires ChromaDB.
"""
message = (
f"ChromaDB is required for {feature} features. "
"Please install it with 'pip install crewai[storage]'"
)
super().__init__(message)