fix: integration with html_mode

This commit is contained in:
Marco Vinciguerra 2024-09-27 18:31:42 +02:00
parent 7783bfbbec
commit f87ffa1d8d

View File

@ -70,7 +70,6 @@ class SmartScraperGraph(AbstractGraph):
"scrape_do": self.config.get("scrape_do") "scrape_do": self.config.get("scrape_do")
} }
) )
generate_answer_node = GenerateAnswerNode( generate_answer_node = GenerateAnswerNode(
input="user_prompt & (relevant_chunks | parsed_doc | doc)", input="user_prompt & (relevant_chunks | parsed_doc | doc)",
@ -82,14 +81,15 @@ class SmartScraperGraph(AbstractGraph):
} }
) )
if self.config.get("html_mode") is not True: if self.config.get("html_mode") is False:
parse_node = ParseNode( parse_node = ParseNode(
input="doc", input="doc",
output=["parsed_doc"], output=["parsed_doc"],
node_config={ node_config={
"llm_model": self.llm_model, "llm_model": self.llm_model,
"chunk_size": self.model_token "chunk_size": self.model_token
}
)
if self.config.get("reasoning"): if self.config.get("reasoning"):
reasoning_node = ReasoningNode( reasoning_node = ReasoningNode(
@ -102,17 +102,17 @@ class SmartScraperGraph(AbstractGraph):
} }
) )
if self.config.get("html_mode") is False and self.config.get("reasoning") is True:
return BaseGraph( return BaseGraph(
nodes=[ nodes=[
fetch_node, fetch_node,
parse_node, parse_node,
reasoning_node, reasoning_node,
generate_answer_node, generate_answer_node,
], ],
edges=[ edges=[
(fetch_node, parse_node), (fetch_node, parse_node),
(parse_node, generate_answer_node)
(parse_node, reasoning_node), (parse_node, reasoning_node),
(reasoning_node, generate_answer_node) (reasoning_node, generate_answer_node)
], ],
@ -120,19 +120,49 @@ class SmartScraperGraph(AbstractGraph):
graph_name=self.__class__.__name__ graph_name=self.__class__.__name__
) )
elif self.config.get("html_mode") is True and self.config.get("reasoning") is True:
return BaseGraph(
nodes=[
fetch_node,
reasoning_node,
generate_answer_node,
],
edges=[
(fetch_node, reasoning_node),
(reasoning_node, generate_answer_node)
],
entry_point=fetch_node,
graph_name=self.__class__.__name__
)
elif self.config.get("html_mode") is True and self.config.get("reasoning") is False:
return BaseGraph(
nodes=[
fetch_node,
generate_answer_node,
],
edges=[
(fetch_node, generate_answer_node)
],
entry_point=fetch_node,
graph_name=self.__class__.__name__
)
return BaseGraph( return BaseGraph(
nodes=[ nodes=[
fetch_node, fetch_node,
parse_node,
generate_answer_node, generate_answer_node,
], ],
edges=[ edges=[
(fetch_node, generate_answer_node) (fetch_node, parse_node),
(parse_node, generate_answer_node)
], ],
entry_point=fetch_node, entry_point=fetch_node,
graph_name=self.__class__.__name__ graph_name=self.__class__.__name__
) )
def run(self) -> str: def run(self) -> str:
""" """
Executes the scraping process and returns the answer to the prompt. Executes the scraping process and returns the answer to the prompt.