diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index 27e7c04b9..969dc1f21 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/src/crewai/knowledge/storage/knowledge_storage.py @@ -6,20 +6,27 @@ import os import shutil from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -try: +if TYPE_CHECKING: 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 - ClientAPI = None - OneOrMany = None - Settings = None - raise ImportError( - "ChromaDB is not installed. Please install it with `pip install crewai[chromadb]`." - ) +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 + ClientAPI = None + OneOrMany = None + Settings = 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.utilities import EmbeddingConfigurator diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index 8aa9bf98a..827784e31 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -6,14 +6,22 @@ import shutil import uuid from typing import TYPE_CHECKING, Any, Dict, List, Optional -try: +if TYPE_CHECKING: import chromadb from chromadb.api import ClientAPI from chromadb.config import Settings -except ImportError: - raise ImportError( - "ChromaDB is not installed. Please install it with `pip install crewai[chromadb]`." - ) +else: + try: + import chromadb + from chromadb.api import ClientAPI + from chromadb.config import Settings + except ImportError: + chromadb = None + ClientAPI = None + Settings = 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.utilities import EmbeddingConfigurator