From 538f1a2a84c61e90b9a470eb79a5b709f04f841c Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 11 Mar 2025 05:52:17 +0000
Subject: [PATCH] Enhance error handling and type annotations for optional
 chromadb dependency

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 .../knowledge/storage/knowledge_storage.py    | 25 ++++++++++---
 src/crewai/memory/storage/rag_storage.py      | 36 ++++++++++++++-----
 2 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py
index 059cbd8cc..a1e1670f3 100644
--- a/src/crewai/knowledge/storage/knowledge_storage.py
+++ b/src/crewai/knowledge/storage/knowledge_storage.py
@@ -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
diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py
index b6d4baf83..7c00c123d 100644
--- a/src/crewai/memory/storage/rag_storage.py
+++ b/src/crewai/memory/storage/rag_storage.py
@@ -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