mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
Fix #2464: Fix embedding dimension mismatch with reset-memories command
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -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")
|
||||
|
||||
55
tests/knowledge/knowledge_storage_test.py
Normal file
55
tests/knowledge/knowledge_storage_test.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user