mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
Fix import error handling for optional chromadb dependency
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user