diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index 20a06f05d..036f72554 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/src/crewai/knowledge/storage/knowledge_storage.py @@ -6,21 +6,16 @@ import os import shutil from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -if TYPE_CHECKING: +try: import chromadb import chromadb.errors from chromadb.api import ClientAPI from chromadb.api.types import OneOrMany from chromadb.config import Settings -else: - try: - import chromadb - import chromadb.errors - from chromadb.api import ClientAPI - from chromadb.api.types import OneOrMany - from chromadb.config import Settings - except ImportError: - chromadb = None +except ImportError: + raise ImportError( + "ChromaDB is not installed. Please install it with `pip install crewai[chromadb]`." + ) from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage from crewai.utilities import EmbeddingConfigurator diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index 87e580f24..8aa9bf98a 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -6,17 +6,14 @@ import shutil import uuid from typing import TYPE_CHECKING, Any, Dict, List, Optional -if TYPE_CHECKING: +try: import chromadb from chromadb.api import ClientAPI from chromadb.config import Settings -else: - try: - import chromadb - from chromadb.api import ClientAPI - from chromadb.config import Settings - except ImportError: - chromadb = None +except ImportError: + raise ImportError( + "ChromaDB is not installed. Please install it with `pip install crewai[chromadb]`." + ) from crewai.memory.storage.base_rag_storage import BaseRAGStorage from crewai.utilities import EmbeddingConfigurator diff --git a/tests/storage/test_optional_chromadb.py b/tests/storage/test_optional_chromadb.py index f00950bca..f5d005ce5 100644 --- a/tests/storage/test_optional_chromadb.py +++ b/tests/storage/test_optional_chromadb.py @@ -12,7 +12,6 @@ class TestOptionalChromadb(unittest.TestCase): with patch.dict(sys.modules, {"chromadb": None}): with pytest.raises(ImportError) as excinfo: from crewai.memory.storage.rag_storage import RAGStorage - storage = RAGStorage(type="test") assert "ChromaDB is not installed" in str(excinfo.value) @@ -21,6 +20,5 @@ class TestOptionalChromadb(unittest.TestCase): with patch.dict(sys.modules, {"chromadb": None}): with pytest.raises(ImportError) as excinfo: from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage - storage = KnowledgeStorage() assert "ChromaDB is not installed" in str(excinfo.value)