andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-08 15:48:29 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Enhance error handling and type annotations for optional chromadb dependency

Browse Source 
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-03-11 05:52:17 +00:00
parent dcf8e8d6e7
commit 538f1a2a84
2 changed files with 48 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/crewai/knowledge/storage/knowledge_storage.py
View File

@@ -49,9 +49,10 @@ class KnowledgeStorage(BaseKnowledgeStorage):
search efficiency.
"""
collection = None # Type will be chromadb.Collection when available
collection: Optional[Any] = None # Will be chromadb.Collection when available
collection_name: Optional[str] = "knowledge"
app = None # Type will be ClientAPI when available
app: Optional[Any] = None # Will be ClientAPI when available
embedder: Optional[Any] = None
def __init__(
self,
@@ -68,6 +69,13 @@ class KnowledgeStorage(BaseKnowledgeStorage):
filter: Optional[dict] = None,
score_threshold: float = 0.35,
) -> List[Dict[str, Any]]:
if not CHROMADB_AVAILABLE:
logging.warning(
"ChromaDB is not installed. Search functionality limited. "
"Install with 'pip install crewai[chromadb]' to enable full functionality."
)
return []
try:
with suppress_logging():
if self.collection:
@@ -89,8 +97,8 @@ class KnowledgeStorage(BaseKnowledgeStorage):
return results
else:
return []
except (ImportError, NameError, AttributeError, Exception):
# Return empty results if chromadb is not available or collection is not initialized
except (ImportError, NameError, AttributeError, Exception) as e:
logging.error(f"Error during knowledge search operation: {str(e)}")
return []
def initialize_knowledge_storage(self):
@@ -216,8 +224,15 @@ class KnowledgeStorage(BaseKnowledgeStorage):
OpenAIEmbeddingFunction,
)
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
logging.warning(
"OPENAI_API_KEY not found. Vector operations will be limited."
)
return None
return OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
api_key=api_key, model_name="text-embedding-3-small"
)
except ImportError:
import logging

36
src/crewai/memory/storage/rag_storage.py
View File

@@ -66,6 +66,15 @@ class RAGStorage(BaseRAGStorage):
self.embedder_config = configurator.configure_embedder(self.embedder_config)
def _initialize_app(self):
if not CHROMADB_AVAILABLE:
logging.warning(
"ChromaDB is not installed. RAG storage functionality will be limited. "
"Install with 'pip install crewai[chromadb]' to enable full functionality."
)
self.app = None
self.collection = None
return
try:
import chromadb
from chromadb.config import Settings
@@ -82,16 +91,13 @@ class RAGStorage(BaseRAGStorage):
self.collection = self.app.get_collection(
name=self.type, embedding_function=self.embedder_config
)
except Exception:
except Exception as e:
logging.info(f"Collection not found, creating new one: {str(e)}")
self.collection = self.app.create_collection(
name=self.type, embedding_function=self.embedder_config
)
except ImportError:
import logging
logging.warning(
"ChromaDB is not installed. RAG storage functionality will be limited. "
"Install with 'pip install crewai[chromadb]' to enable full functionality."
)
except Exception as e:
logging.error(f"Failed to initialize ChromaDB: {str(e)}")
self.app = None
self.collection = None
@@ -134,6 +140,13 @@ class RAGStorage(BaseRAGStorage):
filter: Optional[dict] = None,
score_threshold: float = 0.35,
) -> List[Any]:
if not CHROMADB_AVAILABLE:
logging.warning(
"ChromaDB is not installed. Search functionality limited. "
"Install with 'pip install crewai[chromadb]' to enable full functionality."
)
return []
if not hasattr(self, "app"):
self._initialize_app()
@@ -198,8 +211,15 @@ class RAGStorage(BaseRAGStorage):
OpenAIEmbeddingFunction,
)
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
logging.warning(
"OPENAI_API_KEY not found. Vector operations will be limited."
)
return None
return OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
api_key=api_key, model_name="text-embedding-3-small"
)
except ImportError:
import logging