From f0b1cc23f4a998d1e245419434332a16f1dc1162 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 May 2025 09:10:35 +0000 Subject: [PATCH] Add ChromaDBRequiredError class and improve error handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: João --- .../knowledge/storage/knowledge_storage.py | 46 ++----- src/crewai/memory/storage/rag_storage.py | 23 +--- .../utilities/embedding_configurator.py | 126 +++++------------- src/crewai/utilities/errors/__init__.py | 16 +++ 4 files changed, 66 insertions(+), 145 deletions(-) create mode 100644 src/crewai/utilities/errors/__init__.py diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index 105067cae..8286b0480 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/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. diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index 18cfa358d..110f1ccee 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/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") diff --git a/src/crewai/utilities/embedding_configurator.py b/src/crewai/utilities/embedding_configurator.py index d21ffc3cb..8596ecd58 100644 --- a/src/crewai/utilities/embedding_configurator.py +++ b/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): diff --git a/src/crewai/utilities/errors/__init__.py b/src/crewai/utilities/errors/__init__.py new file mode 100644 index 000000000..5ffb32421 --- /dev/null +++ b/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)