mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-09 21:19:20 +08:00
222 lines
7.4 KiB
Python
222 lines
7.4 KiB
Python
"""
|
|
Bridge class to integrate Burr into ScrapeGraphAI graphs
|
|
[Burr](https://github.com/DAGWorks-Inc/burr)
|
|
"""
|
|
import re
|
|
import uuid
|
|
from hashlib import md5
|
|
from typing import Any, Dict, List, Tuple
|
|
import inspect
|
|
try:
|
|
import burr
|
|
from burr import tracking
|
|
from burr.core import (Application, ApplicationBuilder,
|
|
State, Action, default, ApplicationContext)
|
|
from burr.lifecycle import PostRunStepHook, PreRunStepHook
|
|
except ImportError:
|
|
raise ImportError("""burr package is not installed.
|
|
Please install it with 'pip install scrapegraphai[burr]'""")
|
|
|
|
|
|
class PrintLnHook(PostRunStepHook, PreRunStepHook):
|
|
"""
|
|
Hook to print the action name before and after it is executed.
|
|
"""
|
|
|
|
def pre_run_step(self, *, state: "State", action: "Action", **future_kwargs: Any):
|
|
print(f"Starting action: {action.name}")
|
|
|
|
def post_run_step(self, *, state: "State", action: "Action", **future_kwargs: Any):
|
|
print(f"Finishing action: {action.name}")
|
|
|
|
|
|
class BurrNodeBridge(Action):
|
|
"""Bridge class to convert a base graph node to a Burr action.
|
|
This is nice because we can dynamically declare
|
|
the inputs/outputs (and not rely on function-parsing).
|
|
"""
|
|
|
|
def __init__(self, node):
|
|
"""Instantiates a BurrNodeBridge object.
|
|
"""
|
|
super(BurrNodeBridge, self).__init__()
|
|
self.node = node
|
|
|
|
@property
|
|
def reads(self) -> list[str]:
|
|
return parse_boolean_expression(self.node.input)
|
|
|
|
def run(self, state: State, **run_kwargs) -> dict:
|
|
node_inputs = {key: state[key] for key in self.reads if key in state}
|
|
result_state = self.node.execute(node_inputs, **run_kwargs)
|
|
return result_state
|
|
|
|
@property
|
|
def writes(self) -> list[str]:
|
|
return self.node.output
|
|
|
|
def update(self, result: dict, state: State) -> State:
|
|
return state.update(**result)
|
|
|
|
def get_source(self) -> str:
|
|
return inspect.getsource(self.node.__class__)
|
|
|
|
|
|
def parse_boolean_expression(expression: str) -> List[str]:
|
|
"""
|
|
Parse a boolean expression to extract the keys
|
|
used in the expression, without boolean operators.
|
|
|
|
Args:
|
|
expression (str): The boolean expression to parse.
|
|
|
|
Returns:
|
|
list: A list of unique keys used in the expression.
|
|
"""
|
|
|
|
# Use regular expression to extract all unique keys
|
|
keys = re.findall(r'\w+', expression)
|
|
return list(set(keys)) # Remove duplicates
|
|
|
|
|
|
class BurrBridge:
|
|
"""
|
|
Bridge class to integrate Burr into ScrapeGraphAI graphs.
|
|
|
|
Args:
|
|
base_graph (BaseGraph): The base graph to convert to a Burr application.
|
|
burr_config (dict): Configuration parameters for the Burr application.
|
|
|
|
Attributes:
|
|
base_graph (BaseGraph): The base graph to convert to a Burr application.
|
|
burr_config (dict): Configuration parameters for the Burr application.
|
|
tracker (LocalTrackingClient): The tracking client for the Burr application.
|
|
app_instance_id (str): The instance ID for the Burr application.
|
|
burr_inputs (dict): The inputs for the Burr application.
|
|
burr_app (Application): The Burr application instance.
|
|
|
|
Example:
|
|
>>> burr_bridge = BurrBridge(base_graph, burr_config)
|
|
>>> result = burr_bridge.execute(initial_state={"input_key": "input_value"})
|
|
"""
|
|
|
|
def __init__(self, base_graph, burr_config):
|
|
self.base_graph = base_graph
|
|
self.burr_config = burr_config
|
|
self.project_name = burr_config.get("project_name", "scrapegraph_project")
|
|
self.app_instance_id = burr_config.get("app_instance_id", "default-instance")
|
|
self.burr_inputs = burr_config.get("inputs", {})
|
|
self.burr_app = None
|
|
|
|
def _initialize_burr_app(self, initial_state: Dict[str, Any] = None) -> Application:
|
|
"""
|
|
Initialize a Burr application from the base graph.
|
|
|
|
Args:
|
|
initial_state (dict): The initial state of the Burr application.
|
|
|
|
Returns:
|
|
Application: The Burr application instance.
|
|
"""
|
|
if initial_state is None:
|
|
initial_state = {}
|
|
|
|
actions = self._create_actions()
|
|
transitions = self._create_transitions()
|
|
hooks = [PrintLnHook()]
|
|
burr_state = State(initial_state)
|
|
application_context = ApplicationContext.get()
|
|
builder = (
|
|
ApplicationBuilder()
|
|
.with_actions(**actions)
|
|
.with_transitions(*transitions)
|
|
.with_entrypoint(self.base_graph.entry_point)
|
|
.with_state(**burr_state)
|
|
.with_identifiers(app_id=str(uuid.uuid4())) # TODO -- grab this from state
|
|
.with_hooks(*hooks)
|
|
)
|
|
if application_context is not None:
|
|
builder = (
|
|
builder
|
|
.with_tracker(
|
|
application_context.tracker.copy() if application_context.tracker is not None else None
|
|
)
|
|
.with_spawning_parent(
|
|
application_context.app_id,
|
|
application_context.sequence_id,
|
|
application_context.partition_key,
|
|
)
|
|
)
|
|
else:
|
|
# This is the case in which nothing is spawning it
|
|
# in this case, we want to create a new tracker from scratch
|
|
builder = builder.with_tracker(tracking.LocalTrackingClient(project=self.project_name))
|
|
return builder.build()
|
|
|
|
def _create_actions(self) -> Dict[str, Any]:
|
|
"""
|
|
Create Burr actions from the base graph nodes.
|
|
|
|
Returns:
|
|
dict: A dictionary of Burr actions with the node name
|
|
as keys and the action functions as values.
|
|
"""
|
|
|
|
actions = {}
|
|
for node in self.base_graph.nodes:
|
|
action_func = BurrNodeBridge(node)
|
|
actions[node.node_name] = action_func
|
|
return actions
|
|
|
|
def _create_transitions(self) -> List[Tuple[str, str, Any]]:
|
|
"""
|
|
Create Burr transitions from the base graph edges.
|
|
|
|
Returns:
|
|
list: A list of tuples representing the transitions between Burr actions.
|
|
"""
|
|
|
|
transitions = []
|
|
for from_node, to_node in self.base_graph.edges.items():
|
|
transitions.append((from_node, to_node, default))
|
|
return transitions
|
|
|
|
def _convert_state_from_burr(self, burr_state: State) -> Dict[str, Any]:
|
|
"""
|
|
Convert a Burr state to a dictionary state.
|
|
|
|
Args:
|
|
burr_state (State): The Burr state to convert.
|
|
|
|
Returns:
|
|
dict: The dictionary state instance.
|
|
"""
|
|
|
|
state = {}
|
|
for key in burr_state.__dict__.keys():
|
|
state[key] = getattr(burr_state, key)
|
|
return state
|
|
|
|
def execute(self, initial_state: Dict[str, Any] = {}) -> Dict[str, Any]:
|
|
"""
|
|
Execute the Burr application with the given initial state.
|
|
|
|
Args:
|
|
initial_state (dict): The initial state to pass to the Burr application.
|
|
|
|
Returns:
|
|
dict: The final state of the Burr application.
|
|
"""
|
|
|
|
self.burr_app = self._initialize_burr_app(initial_state)
|
|
|
|
# TODO: to fix final nodes detection
|
|
final_nodes = [self.burr_app.graph.actions[-1].name]
|
|
|
|
last_action, result, final_state = self.burr_app.run(
|
|
halt_after=final_nodes,
|
|
inputs=self.burr_inputs
|
|
)
|
|
|
|
return self._convert_state_from_burr(final_state)
|