Skip to content

Commit 654a042

Browse files
committed
feat(burr-node): working burr bridge
1 parent cfaf7ee commit 654a042

File tree

7 files changed

+15
-66
lines changed

7 files changed

+15
-66
lines changed

examples/openai/burr_integration_openai.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
entry_point=fetch_node,
9191
use_burr=True,
9292
burr_config={
93+
"project_name": "smart-scraper-graph",
9394
"app_instance_id": str(uuid.uuid4()),
9495
"inputs": {
9596
"llm_model": graph_config["llm"].get("model", "gpt-3.5-turbo"),
@@ -101,9 +102,9 @@
101102
# Execute the graph
102103
# ************************************************
103104

104-
result, execution_info = graph.execute({
105-
"user_prompt": "Describe the content",
106-
"url": "https://example.com/"
105+
result, exec_info = graph.execute({
106+
"user_prompt": "List me all the projects with their description",
107+
"url": "https://perinim.github.io/projects/"
107108
})
108109

109110
# get the answer from the result

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ dependencies = [
2929
"playwright==1.43.0",
3030
"google==3.0.0",
3131
"yahoo-search-py==0.3",
32-
"burr[start]"
32+
"burr[start]==0.17.1"
3333
]
3434

3535
license = "MIT"

requirements-dev.lock

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
# features: []
77
# all-features: false
88
# with-sources: false
9+
# generate-hashes: false
910

1011
-e file:.
1112
aiofiles==23.2.1

requirements.lock

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
# features: []
77
# all-features: false
88
# with-sources: false
9+
# generate-hashes: false
910

1011
-e file:.
1112
aiofiles==23.2.1

scrapegraphai/graphs/base_graph.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ def execute(self, initial_state: dict) -> Tuple[dict, list]:
164164
self.initial_state = initial_state
165165
if self.use_burr:
166166
bridge = BurrBridge(self, self.burr_config)
167-
return bridge.execute(initial_state)
167+
result = bridge.execute(initial_state)
168+
return (result["_state"], [])
168169
else:
169170
return self._execute_standard(initial_state)

scrapegraphai/graphs/smart_scraper_graph

Lines changed: 0 additions & 16 deletions
This file was deleted.

scrapegraphai/integrations/burr_bridge.py

Lines changed: 6 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88

99
from burr import tracking
1010
from burr.core import Application, ApplicationBuilder, State, Action, default
11-
from burr.core.action import action
1211
from burr.lifecycle import PostRunStepHook, PreRunStepHook
1312

1413

@@ -40,7 +39,7 @@ def reads(self) -> list[str]:
4039
return parse_boolean_expression(self.node.input)
4140

4241
def run(self, state: State, **run_kwargs) -> dict:
43-
node_inputs = {key: state[key] for key in self.reads}
42+
node_inputs = {key: state[key] for key in self.reads if key in state}
4443
result_state = self.node.execute(node_inputs, **run_kwargs)
4544
return result_state
4645

@@ -49,7 +48,7 @@ def writes(self) -> list[str]:
4948
return self.node.output
5049

5150
def update(self, result: dict, state: State) -> State:
52-
return state.update(**state)
51+
return state.update(**result)
5352

5453

5554
def parse_boolean_expression(expression: str) -> List[str]:
@@ -92,7 +91,8 @@ class BurrBridge:
9291
def __init__(self, base_graph, burr_config):
9392
self.base_graph = base_graph
9493
self.burr_config = burr_config
95-
self.tracker = tracking.LocalTrackingClient(project="smart-scraper-graph")
94+
self.project_name = burr_config.get("project_name", "default-project")
95+
self.tracker = tracking.LocalTrackingClient(project=self.project_name)
9696
self.app_instance_id = burr_config.get("app_instance_id", "default-instance")
9797
self.burr_inputs = burr_config.get("inputs", {})
9898
self.burr_app = None
@@ -111,7 +111,7 @@ def _initialize_burr_app(self, initial_state: Dict[str, Any] = {}) -> Applicatio
111111
actions = self._create_actions()
112112
transitions = self._create_transitions()
113113
hooks = [PrintLnHook()]
114-
burr_state = self._convert_state_to_burr(initial_state)
114+
burr_state = State(initial_state)
115115

116116
app = (
117117
ApplicationBuilder()
@@ -136,32 +136,10 @@ def _create_actions(self) -> Dict[str, Any]:
136136

137137
actions = {}
138138
for node in self.base_graph.nodes:
139-
action_func = self._create_action(node)
139+
action_func = BurrNodeBridge(node)
140140
actions[node.node_name] = action_func
141141
return actions
142142

143-
def _create_action(self, node) -> Any:
144-
"""
145-
Create a Burr action function from a base graph node.
146-
147-
Args:
148-
node (Node): The base graph node to convert to a Burr action.
149-
150-
Returns:
151-
function: The Burr action function.
152-
"""
153-
154-
# @action(reads=parse_boolean_expression(node.input), writes=node.output)
155-
# def dynamic_action(state: State, **kwargs):
156-
# node_inputs = {key: state[key] for key in self._parse_boolean_expression(node.input)}
157-
# result_state = node.execute(node_inputs, **kwargs)
158-
# return result_state, state.update(**result_state)
159-
#
160-
# return dynamic_action
161-
# import pdb
162-
# pdb.set_trace()
163-
return BurrNodeBridge(node)
164-
165143
def _create_transitions(self) -> List[Tuple[str, str, Any]]:
166144
"""
167145
Create Burr transitions from the base graph edges.
@@ -175,22 +153,6 @@ def _create_transitions(self) -> List[Tuple[str, str, Any]]:
175153
transitions.append((from_node, to_node, default))
176154
return transitions
177155

178-
def _convert_state_to_burr(self, state: Dict[str, Any]) -> State:
179-
"""
180-
Convert a dictionary state to a Burr state.
181-
182-
Args:
183-
state (dict): The dictionary state to convert.
184-
185-
Returns:
186-
State: The Burr state instance.
187-
"""
188-
189-
burr_state = State()
190-
for key, value in state.items():
191-
setattr(burr_state, key, value)
192-
return burr_state
193-
194156
def _convert_state_from_burr(self, burr_state: State) -> Dict[str, Any]:
195157
"""
196158
Convert a Burr state to a dictionary state.
@@ -223,7 +185,6 @@ def execute(self, initial_state: Dict[str, Any] = {}) -> Dict[str, Any]:
223185
# TODO: to fix final nodes detection
224186
final_nodes = [self.burr_app.graph.actions[-1].name]
225187

226-
# TODO: fix inputs
227188
last_action, result, final_state = self.burr_app.run(
228189
halt_after=final_nodes,
229190
inputs=self.burr_inputs

0 commit comments

Comments
 (0)