CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-12 21:01:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
requirements-dev.txt
View File

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

7
scrapegraphai/graphs/abstract_graph.py
View File

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

16
scrapegraphai/nodes/rag_node.py
View File

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