fix(cache): correctly pass the node arguments and logging

This commit is contained in:
Marco Perini 2024-06-11 22:56:09 +02:00
parent 543b48764a
commit c881f64209
3 changed files with 15 additions and 10 deletions

View File

@ -1,4 +1,4 @@
sphinx==7.1.2
furo==2024.5.6
pytest==8.0.0
burr[start]==0.19.1
burr[start]==0.22.1

View File

@ -76,6 +76,7 @@ class AbstractGraph(ABC):
self.headless = True if config is None else config.get(
"headless", True)
self.loader_kwargs = config.get("loader_kwargs", {})
self.cache_path = config.get("cache_path", False)
# Create the graph
self.graph = self._create_graph()
@ -91,15 +92,13 @@ class AbstractGraph(ABC):
else:
set_verbosity_warning()
self.headless = True if config is None else config.get("headless", True)
self.loader_kwargs = config.get("loader_kwargs", {})
common_params = {
"headless": self.headless,
"verbose": self.verbose,
"loader_kwargs": self.loader_kwargs,
"llm_model": self.llm_model,
"embedder_model": self.embedder_model
"embedder_model": self.embedder_model,
"cache_path": self.cache_path,
}
self.set_common_params(common_params, overwrite=False)

View File

@ -51,6 +51,7 @@ class RAGNode(BaseNode):
self.verbose = (
False if node_config is None else node_config.get("verbose", False)
)
self.cache_path = node_config.get("cache_path", False)
def execute(self, state: dict) -> dict:
"""
@ -99,15 +100,20 @@ class RAGNode(BaseNode):
)
embeddings = self.embedder_model
folder_name = self.node_config.get("cache", "cache")
folder_name = self.node_config.get("cache_path", "cache")
if self.node_config.get("cache", False) and not os.path.exists(folder_name):
if self.node_config.get("cache_path", False) and not os.path.exists(folder_name):
index = FAISS.from_documents(chunked_docs, embeddings)
os.makedirs(folder_name)
index.save_local(folder_name)
if self.node_config.get("cache", False) and os.path.exists(folder_name):
index = FAISS.load_local(folder_path=folder_name, embeddings=embeddings)
self.logger.info("--- (indexes saved to cache) ---")
elif self.node_config.get("cache_path", False) and os.path.exists(folder_name):
index = FAISS.load_local(folder_path=folder_name,
embeddings=embeddings,
allow_dangerous_deserialization=True)
self.logger.info("--- (indexes loaded from cache) ---")
else:
index = FAISS.from_documents(chunked_docs, embeddings)