feat: add reasoning integration

This commit is contained in:
Marco Vinciguerra 2024-09-27 17:45:44 +02:00
parent 857f28dba0
commit b2822f620a
5 changed files with 81 additions and 6 deletions

View File

@ -0,0 +1,46 @@
"""
Basic example of scraping pipeline using SmartScraper
"""
import os
import json
from dotenv import load_dotenv
from scrapegraphai.graphs import SmartScraperGraph
from scrapegraphai.utils import prettify_exec_info
load_dotenv()
# ************************************************
# Define the configuration for the graph
# ************************************************
graph_config = {
"llm": {
"api_key": os.getenv("OPENAI_API_KEY"),
"model": "openai/gpt-4o",
},
"reasoning": True,
"verbose": True,
"headless": False,
}
# ************************************************
# Create the SmartScraperGraph instance and run it
# ************************************************
smart_scraper_graph = SmartScraperGraph(
prompt="List me what does the company do, the name and a contact email.",
source="https://scrapegraphai.com/",
config=graph_config
)
result = smart_scraper_graph.run()
print(json.dumps(result, indent=4))
# ************************************************
# Get graph execution info
# ************************************************
graph_exec_info = smart_scraper_graph.get_execution_info()
print(prettify_exec_info(graph_exec_info))

View File

@ -9,6 +9,7 @@ from .abstract_graph import AbstractGraph
from ..nodes import ( from ..nodes import (
FetchNode, FetchNode,
ParseNode, ParseNode,
ReasoningNode,
GenerateAnswerNode GenerateAnswerNode
) )
@ -88,6 +89,33 @@ class SmartScraperGraph(AbstractGraph):
} }
) )
if self.config.get("reasoning"):
reasoning_node = ReasoningNode(
input="user_prompt & (relevant_chunks | parsed_doc | doc)",
output=["answer"],
node_config={
"llm_model": self.llm_model,
"additional_info": self.config.get("additional_info"),
"schema": self.schema,
}
)
return BaseGraph(
nodes=[
fetch_node,
parse_node,
reasoning_node,
generate_answer_node,
],
edges=[
(fetch_node, parse_node),
(parse_node, reasoning_node),
(reasoning_node, generate_answer_node)
],
entry_point=fetch_node,
graph_name=self.__class__.__name__
)
return BaseGraph( return BaseGraph(
nodes=[ nodes=[
fetch_node, fetch_node,

View File

@ -26,4 +26,4 @@ from .concat_answers_node import ConcatAnswersNode
from .prompt_refiner_node import PromptRefinerNode from .prompt_refiner_node import PromptRefinerNode
from .html_analyzer_node import HtmlAnalyzerNode from .html_analyzer_node import HtmlAnalyzerNode
from .generate_code_node import GenerateCodeNode from .generate_code_node import GenerateCodeNode
from .reasoning_node import ReasoningNode from .reasoning_node import ReasoningNode

View File

@ -50,12 +50,13 @@ class ReasoningNode(BaseNode):
) )
self.additional_info = node_config.get("additional_info", None) self.additional_info = node_config.get("additional_info", None)
self.output_schema = node_config.get("schema") self.output_schema = node_config.get("schema")
def execute(self, state: dict) -> dict: def execute(self, state: dict) -> dict:
""" """
Generate a refined prompt for the reasoning task based on the user's input and the JSON schema. Generate a refined prompt for the reasoning task based
on the user's input and the JSON schema.
Args: Args:
state (dict): The current state of the graph. The input keys will be used state (dict): The current state of the graph. The input keys will be used
@ -70,11 +71,11 @@ class ReasoningNode(BaseNode):
""" """
self.logger.info(f"--- Executing {self.node_name} Node ---") self.logger.info(f"--- Executing {self.node_name} Node ---")
user_prompt = state['user_prompt'] user_prompt = state['user_prompt']
self.simplefied_schema = transform_schema(self.output_schema.schema()) self.simplefied_schema = transform_schema(self.output_schema.schema())
if self.additional_info is not None: if self.additional_info is not None:
prompt = PromptTemplate( prompt = PromptTemplate(
template=TEMPLATE_REASONING_WITH_CONTEXT, template=TEMPLATE_REASONING_WITH_CONTEXT,

View File

@ -31,7 +31,7 @@ This analysis will be used to instruct an LLM that has the HTML content in its c
**Reasoning Output**: **Reasoning Output**:
[Your detailed analysis based on the above instructions] [Your detailed analysis based on the above instructions]
""" """
TEMPLATE_REASONING_WITH_CONTEXT = """ TEMPLATE_REASONING_WITH_CONTEXT = """
**Task**: Analyze the user's request and the provided JSON schema to guide an LLM in extracting information directly from a markdown file previously parsed froma a HTML file. **Task**: Analyze the user's request and the provided JSON schema to guide an LLM in extracting information directly from a markdown file previously parsed froma a HTML file.