Fix import error handling for optional chromadb dependency

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-04-26 19:06:19 +00:00
parent 9f57e266b8
commit 6435a419d7
3 changed files with 10 additions and 20 deletions

View File

@@ -6,13 +6,6 @@ import os
import shutil import shutil
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
if TYPE_CHECKING:
import chromadb
import chromadb.errors
from chromadb.api import ClientAPI
from chromadb.api.types import OneOrMany
from chromadb.config import Settings
else:
try: try:
import chromadb import chromadb
import chromadb.errors import chromadb.errors
@@ -20,7 +13,9 @@ else:
from chromadb.api.types import OneOrMany from chromadb.api.types import OneOrMany
from chromadb.config import Settings from chromadb.config import Settings
except ImportError: except ImportError:
chromadb = None 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.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage
from crewai.utilities import EmbeddingConfigurator from crewai.utilities import EmbeddingConfigurator

View File

@@ -6,17 +6,14 @@ import shutil
import uuid import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional from typing import TYPE_CHECKING, Any, Dict, List, Optional
if TYPE_CHECKING:
import chromadb
from chromadb.api import ClientAPI
from chromadb.config import Settings
else:
try: try:
import chromadb import chromadb
from chromadb.api import ClientAPI from chromadb.api import ClientAPI
from chromadb.config import Settings from chromadb.config import Settings
except ImportError: except ImportError:
chromadb = None 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.memory.storage.base_rag_storage import BaseRAGStorage
from crewai.utilities import EmbeddingConfigurator from crewai.utilities import EmbeddingConfigurator

View File

@@ -12,7 +12,6 @@ class TestOptionalChromadb(unittest.TestCase):
with patch.dict(sys.modules, {"chromadb": None}): with patch.dict(sys.modules, {"chromadb": None}):
with pytest.raises(ImportError) as excinfo: with pytest.raises(ImportError) as excinfo:
from crewai.memory.storage.rag_storage import RAGStorage from crewai.memory.storage.rag_storage import RAGStorage
storage = RAGStorage(type="test")
assert "ChromaDB is not installed" in str(excinfo.value) 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 patch.dict(sys.modules, {"chromadb": None}):
with pytest.raises(ImportError) as excinfo: with pytest.raises(ImportError) as excinfo:
from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage
storage = KnowledgeStorage()
assert "ChromaDB is not installed" in str(excinfo.value) assert "ChromaDB is not installed" in str(excinfo.value)