From 6435a419d7fdb488e46bb972d31c54f681d9b4c5 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sat, 26 Apr 2025 19:06:19 +0000
Subject: [PATCH] Fix import error handling for optional chromadb dependency

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/knowledge/storage/knowledge_storage.py | 15 +++++----------
 src/crewai/memory/storage/rag_storage.py          | 13 +++++--------
 tests/storage/test_optional_chromadb.py           |  2 --
 3 files changed, 10 insertions(+), 20 deletions(-)

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)