feat: lazy import chromadb

lib/crewai/src/crewai/knowledge/storage/knowledge_storage.py

@@ -1,21 +1,24 @@
+from __future__ import annotations
+
 import logging
 import traceback
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 import warnings
 
 from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage
-from crewai.rag.chromadb.config import ChromaDBConfig
-from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
 from crewai.rag.config.utils import get_rag_client
 from crewai.rag.core.base_client import BaseClient
 from crewai.rag.core.base_embeddings_provider import BaseEmbeddingsProvider
 from crewai.rag.embeddings.factory import build_embedder
-from crewai.rag.embeddings.types import ProviderSpec
 from crewai.rag.factory import create_client
 from crewai.rag.types import BaseRecord, SearchResult
 from crewai.utilities.logger import Logger
 
 
+if TYPE_CHECKING:
+    from crewai.rag.embeddings.types import ProviderSpec
+
+
 class KnowledgeStorage(BaseKnowledgeStorage):
     """
     Extends Storage to handle embeddings for memory entries, improving
@@ -30,6 +33,9 @@ class KnowledgeStorage(BaseKnowledgeStorage):
         | None = None,
         collection_name: str | None = None,
     ) -> None:
+        from crewai.rag.chromadb.config import ChromaDBConfig
+        from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
+
         self.collection_name = collection_name
         self._client: BaseClient | None = None
 

lib/crewai/src/crewai/memory/storage/rag_storage.py

@@ -5,8 +5,6 @@ import traceback
 from typing import TYPE_CHECKING, Any, cast
 import warnings
 
-from crewai.rag.chromadb.config import ChromaDBConfig
-from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
 from crewai.rag.config.utils import get_rag_client
 from crewai.rag.embeddings.factory import build_embedder
 from crewai.rag.factory import create_client
@@ -37,6 +35,9 @@ class RAGStorage(BaseRAGStorage):
         crew: Crew | None = None,
         path: str | None = None,
     ) -> None:
+        from crewai.rag.chromadb.config import ChromaDBConfig
+        from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
+
         super().__init__(type, allow_reset, embedder_config, crew)
         crew_agents = crew.agents if crew else []
         sanitized_roles = [self._sanitize_role(agent.role) for agent in crew_agents]

lib/crewai/src/crewai/rag/chromadb/config.py

@@ -1,11 +1,12 @@
 """ChromaDB configuration model."""
 
+from __future__ import annotations
+
 from dataclasses import field
 import os
-from typing import Literal, cast
+from typing import TYPE_CHECKING, Literal, cast
 import warnings
 
-from chromadb.config import Settings
 from pydantic.dataclasses import dataclass as pyd_dataclass
 
 from crewai.rag.chromadb.constants import (
@@ -13,10 +14,15 @@ from crewai.rag.chromadb.constants import (
     DEFAULT_STORAGE_PATH,
     DEFAULT_TENANT,
 )
-from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
 from crewai.rag.config.base import BaseRagConfig
 
 
+if TYPE_CHECKING:
+    from chromadb.config import Settings
+
+    from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
+
+
 warnings.filterwarnings(
     "ignore",
     message=".*Mixing V1 models and V2 models.*",
@@ -37,6 +43,8 @@ def _default_settings() -> Settings:
     Returns:
         Settings with persistent storage and reset enabled.
     """
+    from chromadb.config import Settings
+
     return Settings(
         persist_directory=DEFAULT_STORAGE_PATH,
         allow_reset=True,
@@ -54,6 +62,8 @@ def _default_embedding_function() -> ChromaEmbeddingFunctionWrapper:
         OpenAIEmbeddingFunction,
     )
 
+    from crewai.rag.chromadb.types import ChromaEmbeddingFunctionWrapper
+
     return cast(
         ChromaEmbeddingFunctionWrapper,
         OpenAIEmbeddingFunction(

lib/crewai/src/crewai/rag/chromadb/factory.py

@@ -1,13 +1,18 @@
 """Factory functions for creating ChromaDB clients."""
 
+from __future__ import annotations
+
 from hashlib import md5
 import os
+from typing import TYPE_CHECKING
 
 from chromadb import PersistentClient
 import portalocker
 
-from crewai.rag.chromadb.client import ChromaDBClient
+
-from crewai.rag.chromadb.config import ChromaDBConfig
+if TYPE_CHECKING:
+    from crewai.rag.chromadb.client import ChromaDBClient
+    from crewai.rag.chromadb.config import ChromaDBConfig
 
 
 def create_client(config: ChromaDBConfig) -> ChromaDBClient:
@@ -22,6 +27,7 @@ def create_client(config: ChromaDBConfig) -> ChromaDBClient:
     Notes:
         Need to update to use chromadb.Client to support more client types in the near future.
     """
+    from crewai.rag.chromadb.client import ChromaDBClient
 
     persist_dir = config.settings.persist_directory
     os.makedirs(persist_dir, exist_ok=True)