Add ChromaDBRequiredError class and improve error handling

Co-Authored-By: João <joao@crewai.com>

src/crewai/knowledge/storage/knowledge_storage.py

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

src/crewai/memory/storage/rag_storage.py

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

src/crewai/utilities/embedding_configurator.py

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

src/crewai/utilities/errors/__init__.py

@@ -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)