diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index 72240e2b6..b67be51af 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/src/crewai/knowledge/storage/knowledge_storage.py @@ -114,8 +114,20 @@ class KnowledgeStorage(BaseKnowledgeStorage): settings=Settings(allow_reset=True), ) - self.app.reset() - shutil.rmtree(base_path) + # If a specific collection name is provided, try to delete just that collection + if self.collection_name: + try: + collection_name = f"knowledge_{self.collection_name}" + self.app.delete_collection(name=collection_name) + Logger(verbose=True).log("info", f"Collection '{collection_name}' has been reset.") + except Exception: + # If deletion of specific collection fails, fall back to resetting all + self.app.reset() + shutil.rmtree(base_path) + else: + self.app.reset() + shutil.rmtree(base_path) + self.app = None self.collection = None @@ -164,15 +176,19 @@ class KnowledgeStorage(BaseKnowledgeStorage): ids=filtered_ids, ) except chromadb.errors.InvalidDimensionException as e: + collection_reset_cmd = "`crewai reset-memories -a`" + if self.collection_name: + collection_reset_cmd = f"`crewai reset-memories --knowledge` (for collection: {self.collection_name})" + Logger(verbose=True).log( "error", - "Embedding dimension mismatch. This usually happens when mixing different embedding models. Try resetting the collection using `crewai reset-memories -a`", + f"Embedding dimension mismatch. This usually happens when mixing different embedding models. Try resetting using {collection_reset_cmd}", "red", ) raise ValueError( - "Embedding dimension mismatch. Make sure you're using the same embedding model " - "across all operations with this collection." - "Try resetting the collection using `crewai reset-memories -a`" + f"Embedding dimension mismatch. Make sure you're using the same embedding model " + f"across all operations with this collection. " + f"Try resetting using {collection_reset_cmd}" ) from e except Exception as e: Logger(verbose=True).log("error", f"Failed to upsert documents: {e}", "red") diff --git a/tests/knowledge/knowledge_storage_test.py b/tests/knowledge/knowledge_storage_test.py new file mode 100644 index 000000000..adf1c9f13 --- /dev/null +++ b/tests/knowledge/knowledge_storage_test.py @@ -0,0 +1,55 @@ +import os +import pytest +from unittest.mock import MagicMock, patch + +from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage + + +class TestKnowledgeStorage: + @patch("crewai.knowledge.storage.knowledge_storage.chromadb") + @patch("crewai.knowledge.storage.knowledge_storage.shutil") + def test_reset_with_default_collection(self, mock_shutil, mock_chromadb): + # Setup + mock_app = MagicMock() + mock_chromadb.PersistentClient.return_value = mock_app + + # Execute + storage = KnowledgeStorage() + storage.reset() + + # Verify + mock_app.reset.assert_called_once() + mock_shutil.rmtree.assert_called_once() + + @patch("crewai.knowledge.storage.knowledge_storage.chromadb") + @patch("crewai.knowledge.storage.knowledge_storage.shutil") + def test_reset_with_custom_collection(self, mock_shutil, mock_chromadb): + # Setup + mock_app = MagicMock() + mock_chromadb.PersistentClient.return_value = mock_app + + # Execute + storage = KnowledgeStorage(collection_name="custom_collection") + storage.reset() + + # Verify + mock_app.delete_collection.assert_called_once_with(name="knowledge_custom_collection") + mock_app.reset.assert_not_called() + mock_shutil.rmtree.assert_not_called() + + @patch("crewai.knowledge.storage.knowledge_storage.chromadb") + @patch("crewai.knowledge.storage.knowledge_storage.shutil") + def test_reset_with_custom_collection_fallback(self, mock_shutil, mock_chromadb): + # Setup + mock_app = MagicMock() + mock_app.delete_collection.side_effect = Exception("Collection not found") + mock_chromadb.PersistentClient.return_value = mock_app + + # Execute + storage = KnowledgeStorage(collection_name="custom_collection") + storage.reset() + + # Verify + mock_app.delete_collection.assert_called_once_with(name="knowledge_custom_collection") + mock_app.reset.assert_called_once() + mock_shutil.rmtree.assert_called_once()