From 6b87d22a7084897d243f47da5a9218786994d72f Mon Sep 17 00:00:00 2001 From: Carlos Souza Date: Mon, 9 Dec 2024 13:30:51 -0500 Subject: [PATCH] Fix disk I/O error when resetting short-term memory. (#1724) * Fix disk I/O error when resetting short-term memory. Reset chromadb client and nullifies references before removing directory. * Nit for clarity * did the same for knowledge_storage * cleanup * cleanup order * Cleanup after the rm of the directories --------- Co-authored-by: Lorenze Jay Co-authored-by: Lorenze Jay <63378463+lorenzejay@users.noreply.github.com> --- src/crewai/knowledge/storage/knowledge_storage.py | 14 +++++++++----- src/crewai/memory/storage/rag_storage.py | 4 +++- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index e0e36f4a7..acbbe75a1 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/src/crewai/knowledge/storage/knowledge_storage.py @@ -3,6 +3,7 @@ import hashlib import io import logging import os +import shutil from typing import Any, Dict, List, Optional, Union, cast import chromadb @@ -15,6 +16,7 @@ from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage from crewai.utilities import EmbeddingConfigurator from crewai.utilities.logger import Logger from crewai.utilities.paths import db_storage_path +from crewai.utilities.constants import KNOWLEDGE_DIRECTORY @contextlib.contextmanager @@ -105,15 +107,17 @@ class KnowledgeStorage(BaseKnowledgeStorage): raise Exception("Failed to create or get collection") def reset(self): - if self.app: - self.app.reset() - else: - base_path = os.path.join(db_storage_path(), "knowledge") + base_path = os.path.join(db_storage_path(), KNOWLEDGE_DIRECTORY) + if not self.app: self.app = chromadb.PersistentClient( path=base_path, settings=Settings(allow_reset=True), ) - self.app.reset() + + self.app.reset() + shutil.rmtree(base_path) + self.app = None + self.collection = None def save( self, diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index bf40aee96..fd4c77838 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -150,9 +150,11 @@ class RAGStorage(BaseRAGStorage): def reset(self) -> None: try: - shutil.rmtree(f"{db_storage_path()}/{self.type}") if self.app: self.app.reset() + shutil.rmtree(f"{db_storage_path()}/{self.type}") + self.app = None + self.collection = None except Exception as e: if "attempt to write a readonly database" in str(e): # Ignore this specific error