89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
import asyncio
|
|
|
|
from neo4j_graphrag.experimental.pipeline.component import Component, DataModel
|
|
from neo4j_graphrag.experimental.pipeline.pipeline import Pipeline
|
|
from neo4j_graphrag.experimental.pipeline.notification import EventType, Event
|
|
from neo4j_graphrag.experimental.pipeline.types.context import RunContext
|
|
|
|
|
|
# Define some example components with progress notifications
|
|
class OutputModel(DataModel):
|
|
result: int
|
|
|
|
|
|
class SlowAdder(Component):
|
|
"""A component that slowly adds numbers and reports progress"""
|
|
|
|
def __init__(self, number: int) -> None:
|
|
self.number = number
|
|
|
|
async def run_with_context(self, context_: RunContext, value: int) -> OutputModel:
|
|
# Simulate work with progress updates
|
|
for i in range(value):
|
|
await asyncio.sleep(0.5) # Simulate work
|
|
await context_.notify(
|
|
message=f"Added {i + 1}/{value}",
|
|
data={"current": i + 1, "total": value},
|
|
)
|
|
return OutputModel(result=value + self.number)
|
|
|
|
|
|
class SlowMultiplier(Component):
|
|
"""A component that slowly multiplies numbers and reports progress"""
|
|
|
|
def __init__(self, multiplier: int) -> None:
|
|
self.multiplier = multiplier
|
|
|
|
async def run_with_context(self, context_: RunContext, value: int) -> OutputModel:
|
|
# Simulate work with progress updates
|
|
for i in range(3): # Always do 3 steps
|
|
await asyncio.sleep(0.7) # Simulate work
|
|
await context_.notify(
|
|
message=f"Multiplication step {i + 1}/3",
|
|
data={"step": i + 1, "total": 3},
|
|
)
|
|
return OutputModel(result=value * self.multiplier)
|
|
|
|
|
|
async def callback(event: Event) -> None:
|
|
await asyncio.sleep(0.1)
|
|
|
|
|
|
async def main() -> None:
|
|
# Create pipeline
|
|
pipeline = Pipeline(callback=callback)
|
|
|
|
# Add components
|
|
pipeline.add_component(SlowAdder(number=3), "adder")
|
|
pipeline.add_component(SlowMultiplier(multiplier=2), "multiplier")
|
|
|
|
# Connect components
|
|
pipeline.connect("adder", "multiplier", {"value": "adder.result"})
|
|
|
|
print("\n=== Running pipeline with streaming ===")
|
|
# Run pipeline with streaming - see events as they happen
|
|
async for event in pipeline.stream(
|
|
{"adder": {"value": 2}},
|
|
raise_exception=False, # default is True
|
|
):
|
|
if event.event_type == EventType.PIPELINE_STARTED:
|
|
print("Stream: Pipeline started!")
|
|
elif event.event_type == EventType.PIPELINE_FINISHED:
|
|
print(f"Stream: Pipeline finished! Final results: {event.payload}")
|
|
elif event.event_type == EventType.PIPELINE_FAILED:
|
|
print(f"Stream: Pipeline failed with message: {event.message}")
|
|
elif event.event_type == EventType.TASK_STARTED:
|
|
print(
|
|
f"Stream: Task {event.task_name} started with inputs: {event.payload}" # type: ignore
|
|
)
|
|
elif event.event_type == EventType.TASK_PROGRESS:
|
|
print(f"Stream: Task {event.task_name} progress - {event.message}") # type: ignore
|
|
elif event.event_type == EventType.TASK_FINISHED:
|
|
print(
|
|
f"Stream: Task {event.task_name} finished with result: {event.payload}" # type: ignore
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|