diff --git a/lib/crewai/src/crewai/knowledge/storage/knowledge_storage.py b/lib/crewai/src/crewai/knowledge/storage/knowledge_storage.py index 055763f7f..a1ad11929 100644 --- a/lib/crewai/src/crewai/knowledge/storage/knowledge_storage.py +++ b/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 diff --git a/lib/crewai/src/crewai/memory/storage/rag_storage.py b/lib/crewai/src/crewai/memory/storage/rag_storage.py index b45cde55a..08f8ea2ca 100644 --- a/lib/crewai/src/crewai/memory/storage/rag_storage.py +++ b/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] diff --git a/lib/crewai/src/crewai/rag/chromadb/config.py b/lib/crewai/src/crewai/rag/chromadb/config.py index 49a8b22ff..67b421b3e 100644 --- a/lib/crewai/src/crewai/rag/chromadb/config.py +++ b/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( diff --git a/lib/crewai/src/crewai/rag/chromadb/factory.py b/lib/crewai/src/crewai/rag/chromadb/factory.py index 933da10a2..48edc82e2 100644 --- a/lib/crewai/src/crewai/rag/chromadb/factory.py +++ b/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)