Skip to content

Commit fff1232

Browse files
committed
add rag node
1 parent 5bda918 commit fff1232

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

scrapegraphai/graphs/pdf_scraper_graph.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
from ..nodes import (
1212
FetchNode,
13+
RAGNode,
1314
GenerateAnswerPDFNode
1415
)
1516

@@ -63,7 +64,15 @@ def _create_graph(self) -> BaseGraph:
6364
input='pdf | pdf_dir',
6465
output=["doc"],
6566
)
66-
67+
68+
rag_node = RAGNode(
69+
input="user_prompt & (parsed_doc | doc)",
70+
output=["relevant_chunks"],
71+
node_config={
72+
"llm_model": self.llm_model,
73+
"embedder_model": self.embedder_model
74+
}
75+
)
6776
generate_answer_node_pdf = GenerateAnswerPDFNode(
6877
input="user_prompt & (relevant_chunks | doc)",
6978
output=["answer"],
@@ -76,10 +85,12 @@ def _create_graph(self) -> BaseGraph:
7685
return BaseGraph(
7786
nodes=[
7887
fetch_node,
88+
rag_node,
7989
generate_answer_node_pdf,
8090
],
8191
edges=[
82-
(fetch_node, generate_answer_node_pdf)
92+
(fetch_node, rag_node),
93+
(rag_node, generate_answer_node_pdf)
8394
],
8495
entry_point=fetch_node
8596
)
@@ -95,4 +106,4 @@ def run(self) -> str:
95106
inputs = {"user_prompt": self.prompt, self.input_key: self.source}
96107
self.final_state, self.execution_info = self.graph.execute(inputs)
97108

98-
return self.final_state.get("answer", "No answer found.")
109+
return self.final_state.get("answer", "No answer found.")

0 commit comments

Comments
 (0)