run-llama/LlamaIndexTS
1
0
Fork 0
mirror of https://github.com/run-llama/LlamaIndexTS.git synced 2026-08-26 12:24:44 -04:00
Code Issues 138 Packages Projects Releases 10 Wiki Activity

feat(workflow): allow send event with no output (#1479)

Browse Source
This commit is contained in:
Alex Yang
2024-11-14 00:17:53 -08:00
committed by GitHub
parent 1d470363df
commit ee20c44d9b
4 changed files with 130 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/perfect-zebras-live.md
+6
View File
@@ -0,0 +1,6 @@
---
"@llamaindex/unit-test": patch
"@llamaindex/workflow": patch
---
feat(workflow): allow send event with no output
packages/workflow/src/workflow-context.ts
+103 -96
View File
@@ -13,19 +13,18 @@ export type StepHandler<
AnyWorkflowEventConstructor | StartEventConstructor,
...(AnyWorkflowEventConstructor | StopEventConstructor)[],
] = [AnyWorkflowEventConstructor | StartEventConstructor],
Out extends [
AnyWorkflowEventConstructor | StartEventConstructor,
...(AnyWorkflowEventConstructor | StopEventConstructor)[],
] = [AnyWorkflowEventConstructor | StopEventConstructor],
Out extends (AnyWorkflowEventConstructor | StopEventConstructor)[] = [],
> = (
context: HandlerContext<Data>,
...events: {
[K in keyof Inputs]: InstanceType<Inputs[K]>;
}
) => Promise<
{
[K in keyof Out]: InstanceType<Out[K]>;
}[number]
Out extends []
? void
: {
[K in keyof Out]: InstanceType<Out[K]>;
}[number]
>;
export type ReadonlyStepMap<Data> = ReadonlyMap<
@@ -275,7 +274,7 @@ export class WorkflowContext<Start = string, Stop = string, Data = unknown>
*/
#createStreamEvents(): AsyncIterableIterator<WorkflowEvent<unknown>> {
const isPendingEvents = new WeakSet<WorkflowEvent<unknown>>();
const pendingTasks = new Set<Promise<WorkflowEvent<unknown>>>();
const pendingTasks = new Set<Promise<WorkflowEvent<unknown> | void>>();
const enqueuedEvents = new Set<WorkflowEvent<unknown>>();
const stream = new ReadableStream<WorkflowEvent<unknown>>({
start: async (controller) => {
@@ -325,102 +324,104 @@ export class WorkflowContext<Start = string, Stop = string, Data = unknown>
}
const [steps, inputsMap, outputsMap] =
this.#getStepFunction(event);
const nextEventPromises: Promise<WorkflowEvent<unknown>>[] = [
...steps,
]
.map((step) => {
const inputs = [...(inputsMap.get(step) ?? [])];
const acceptableInputs: WorkflowEvent<unknown>[] =
this.#pendingInputQueue.filter((event) =>
inputs.some((input) => event instanceof input),
);
const events: WorkflowEvent<unknown>[] = flattenEvents(
inputs,
[event, ...acceptableInputs],
);
// remove the event from the queue, in case of infinite loop
events.forEach((event) => {
const protocolIdx = this.#queue.findIndex(
(protocol) =>
protocol.type === "event" &&
protocol.event === event,
);
if (protocolIdx !== -1) {
this.#queue.splice(protocolIdx, 1);
}
});
if (events.length !== inputs.length) {
if (this.#verbose) {
console.log(
`Not enough inputs for step ${step.name}, waiting for more events`,
const nextEventPromises: Promise<WorkflowEvent<unknown> | void>[] =
[...steps]
.map((step) => {
const inputs = [...(inputsMap.get(step) ?? [])];
const acceptableInputs: WorkflowEvent<unknown>[] =
this.#pendingInputQueue.filter((event) =>
inputs.some((input) => event instanceof input),
);
}
// not enough to run the step, push back to the queue
this.#sendEvent(event);
isPendingEvents.add(event);
return null;
}
if (isPendingEvents.has(event)) {
isPendingEvents.delete(event);
}
if (this.#verbose) {
console.log(
`Running step ${step.name} with inputs ${events}`,
const events: WorkflowEvent<unknown>[] = flattenEvents(
inputs,
[event, ...acceptableInputs],
);
}
const data = this.data;
return (step as StepHandler<Data>)
.call(
null,
{
get data() {
return data;
},
sendEvent: this.#sendEvent,
requireEvent: this.#requireEvent,
},
// @ts-expect-error IDK why
...events.sort((a, b) => {
const aIndex = inputs.indexOf(
a.constructor as AnyWorkflowEventConstructor,
);
const bIndex = inputs.indexOf(
b.constructor as AnyWorkflowEventConstructor,
);
return aIndex - bIndex;
}),
)
.then((nextEvent) => {
// remove the event from the queue, in case of infinite loop
events.forEach((event) => {
const protocolIdx = this.#queue.findIndex(
(protocol) =>
protocol.type === "event" &&
protocol.event === event,
);
if (protocolIdx !== -1) {
this.#queue.splice(protocolIdx, 1);
}
});
if (events.length !== inputs.length) {
if (this.#verbose) {
console.log(
`Step ${step.name} completed, next event is ${nextEvent}`,
`Not enough inputs for step ${step.name}, waiting for more events`,
);
}
const outputs = outputsMap.get(step) ?? [];
if (
!outputs.some(
(output) => nextEvent.constructor === output,
)
) {
if (this.#strict) {
const error = Error(
`Step ${step.name} returned an unexpected output event ${nextEvent}`,
// not enough to run the step, push back to the queue
this.#sendEvent(event);
isPendingEvents.add(event);
return null;
}
if (isPendingEvents.has(event)) {
isPendingEvents.delete(event);
}
if (this.#verbose) {
console.log(
`Running step ${step.name} with inputs ${events}`,
);
}
const data = this.data;
return (step as StepHandler<Data>)
.call(
null,
{
get data() {
return data;
},
sendEvent: this.#sendEvent,
requireEvent: this.#requireEvent,
},
// @ts-expect-error IDK why
...events.sort((a, b) => {
const aIndex = inputs.indexOf(
a.constructor as AnyWorkflowEventConstructor,
);
controller.error(error);
} else {
console.warn(
`Step ${step.name} returned an unexpected output event ${nextEvent}`,
const bIndex = inputs.indexOf(
b.constructor as AnyWorkflowEventConstructor,
);
return aIndex - bIndex;
}),
)
.then((nextEvent: void | WorkflowEvent<unknown>) => {
if (nextEvent === undefined) {
return;
}
if (this.#verbose) {
console.log(
`Step ${step.name} completed, next event is ${nextEvent}`,
);
}
}
if (!(nextEvent instanceof StopEvent)) {
this.#pendingInputQueue.unshift(nextEvent);
this.#sendEvent(nextEvent);
}
return nextEvent;
});
})
.filter((promise) => promise !== null);
const outputs = outputsMap.get(step) ?? [];
if (
!outputs.some(
(output) => nextEvent.constructor === output,
)
) {
if (this.#strict) {
const error = Error(
`Step ${step.name} returned an unexpected output event ${nextEvent}`,
);
controller.error(error);
} else {
console.warn(
`Step ${step.name} returned an unexpected output event ${nextEvent}`,
);
}
}
if (!(nextEvent instanceof StopEvent)) {
this.#pendingInputQueue.unshift(nextEvent);
this.#sendEvent(nextEvent);
}
return nextEvent;
});
})
.filter((promise) => promise !== null);
nextEventPromises.forEach((promise) => {
pendingTasks.add(promise);
promise
@@ -433,6 +434,9 @@ export class WorkflowContext<Start = string, Stop = string, Data = unknown>
});
Promise.race(nextEventPromises)
.then((fastestNextEvent) => {
if (fastestNextEvent === undefined) {
return;
}
if (!enqueuedEvents.has(fastestNextEvent)) {
controller.enqueue(fastestNextEvent);
enqueuedEvents.add(fastestNextEvent);
@@ -441,7 +445,10 @@ export class WorkflowContext<Start = string, Stop = string, Data = unknown>
})
.then(async (fastestNextEvent) =>
Promise.all(nextEventPromises).then((nextEvents) => {
for (const nextEvent of nextEvents) {
const events = nextEvents.filter(
(event) => event !== undefined,
);
for (const nextEvent of events) {
// do not enqueue the same event twice
if (fastestNextEvent !== nextEvent) {
if (!enqueuedEvents.has(nextEvent)) {
packages/workflow/src/workflow.ts
+6 -7
View File
@@ -57,10 +57,7 @@ export class Workflow<ContextData, Start, Stop> {
AnyWorkflowEventConstructor | StartEventConstructor,
...(AnyWorkflowEventConstructor | StopEventConstructor)[],
],
const Out extends [
AnyWorkflowEventConstructor | StopEventConstructor,
...(AnyWorkflowEventConstructor | StopEventConstructor)[],
],
const Out extends (AnyWorkflowEventConstructor | StopEventConstructor)[],
>(
parameters: StepParameters<In, Out>,
stepFn: (
@@ -69,9 +66,11 @@ export class Workflow<ContextData, Start, Stop> {
[K in keyof In]: InstanceType<In[K]>;
}
) => Promise<
{
[K in keyof Out]: InstanceType<Out[K]>;
}[number]
Out extends []
? void
: {
[K in keyof Out]: InstanceType<Out[K]>;
}[number]
>,
): this {
const { inputs, outputs } = parameters;
unit/workflow/workflow.test.ts
+15
View File
@@ -531,6 +531,21 @@ describe("workflow basic", () => {
const result = await myWorkflow.run("start");
expect(result.data).toBe("query result");
});
test("allow output with send event", async () => {
const myFlow = new Workflow<unknown, string, string>({ verbose: true });
myFlow.addStep(
{
inputs: [StartEvent<string>],
outputs: [],
},
async (context, ev) => {
context.sendEvent(new StopEvent(`Hello ${ev.data}!`));
},
);
const result = myFlow.run("world");
expect((await result).data).toBe("Hello world!");
});
});
describe("workflow event loop", () => {