From fbfddb61daca69f4cf9f3a8fe6466227be30b52f Mon Sep 17 00:00:00 2001 From: "Rossdan Craig rossdan@lastmileai.dev" <> Date: Sun, 14 Jan 2024 02:55:22 -0500 Subject: [PATCH] Implement STREAM_AICONFIG_CHUNK action This is the biggest part of this diff stack. The functionality is still exactly the same as before, but now we can call it directly instead of as a sub-action within `"CONSOLIDATE_AICONFIG"`. I think this is much cleaner. Another benefit is this removes the need for having to check the `isRunning` flags within the consolidated `"RUN_PROMPT"` sub-action, since we now know that the only time we can call that is through the beginning and the end (non-streaming end). Next PR I am going to split this into two separate Actions: 1. `"RUN_PROMPT_START"` 2. `"RUN_PROMPT_SUCCESS"` ## Test Plan Streaming and non-streaming models still work. Can cancel, it'll end, etc https://github.com/lastmile-ai/aiconfig/assets/151060367/cf29ee99-9ede-4c5e-99a7-a7f7816adfe1 --- .../client/src/components/AIConfigEditor.tsx | 10 +-- .../client/src/components/aiconfigReducer.ts | 86 +++++-------------- 2 files changed, 24 insertions(+), 72 deletions(-) diff --git a/python/src/aiconfig/editor/client/src/components/AIConfigEditor.tsx b/python/src/aiconfig/editor/client/src/components/AIConfigEditor.tsx index 90fe1180d..c76641ea1 100644 --- a/python/src/aiconfig/editor/client/src/components/AIConfigEditor.tsx +++ b/python/src/aiconfig/editor/client/src/components/AIConfigEditor.tsx @@ -638,15 +638,9 @@ export default function EditorContainer({ }); } else if (event.type === "aiconfig_chunk") { dispatch({ - type: "CONSOLIDATE_AICONFIG", - action: { - type: "STREAM_AICONFIG_CHUNK", - id: promptId, - cancellationToken, - // Keep the prompt running state until the end of streaming - isRunning: true, - }, + type: "STREAM_AICONFIG_CHUNK", config: event.data, + cancellationToken, }); } else if (event.type === "stop_streaming") { // Pass this event at the end of streaming to signal diff --git a/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts b/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts index 218b6a114..645a1fff9 100644 --- a/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts +++ b/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts @@ -29,7 +29,6 @@ export type MutateAIConfigAction = export type ConsolidateAIConfigSubAction = | AddPromptAction | RunPromptAction - | StreamAIConfigChunkAction | UpdatePromptInputAction; export type ConsolidateAIConfigAction = { @@ -82,9 +81,8 @@ export type SetNameAction = { export type StreamAIConfigChunkAction = { type: "STREAM_AICONFIG_CHUNK"; - id: string; + config: AIConfig; cancellationToken?: string; - isRunning?: boolean; }; export type StreamOutputChunkAction = { @@ -201,57 +199,16 @@ function reduceConsolidateAIConfig( consolidatePrompt ); } + // Next PR: Split "RUN_PROMPT" into two actions: + // 1) "RUN_PROMPT_START" + // 2) "RUN_PROMPT_SUCCESS" + // 3) (Already exists) "RUN_PROMPT_ERROR" case "RUN_PROMPT": { - // Note: If we are calling "RUN_PROMPT" directly as a dispatched event - // type, we automatically set the state there to `isRunning` for that - // prompt. That logic does not happen here, it happens in - // `aiconfigReducer`. - // If we are calling "RUN_PROMPT" indirectly via the action of a - // "CONSOLIDATE_AICONFIG" dispatch, we end up here. We need to check - // if we actually want to set the prompt state to `isRunning` - const isRunning = action.isRunning ?? false; const stateWithUpdatedRunningPromptId = { ...state, _ui: { ...state._ui, - runningPromptId: isRunning ? action.id : undefined, - }, - }; - return reduceReplacePrompt( - stateWithUpdatedRunningPromptId, - action.id, - (prompt) => { - const responsePrompt = responseConfig.prompts.find( - (resPrompt) => resPrompt.name === prompt.name - ); - - const outputs = responsePrompt?.outputs ?? prompt.outputs; - - return { - ...prompt, - _ui: { - ...prompt._ui, - isRunning, - }, - outputs, - }; - } - ); - } - case "STREAM_AICONFIG_CHUNK": { - // Note: If we are calling "RUN_PROMPT" directly as a dispatched event - // type, we automatically set the state there to `isRunning` for that - // prompt. That logic does not happen here, it happens in - // `aiconfigReducer`. - // If we are calling "RUN_PROMPT" indirectly via the action of a - // "CONSOLIDATE_AICONFIG" dispatch, we end up here. We need to check - // if we actually want to set the prompt state to `isRunning` - const isRunning = action.isRunning ?? false; - const stateWithUpdatedRunningPromptId = { - ...state, - _ui: { - ...state._ui, - runningPromptId: isRunning ? action.id : undefined, + runningPromptId: undefined, }, }; return reduceReplacePrompt( @@ -268,7 +225,7 @@ function reduceConsolidateAIConfig( ...prompt, _ui: { ...prompt._ui, - isRunning, + isRunning: false, }, outputs, }; @@ -393,21 +350,22 @@ export default function aiconfigReducer( }; } case "STREAM_AICONFIG_CHUNK": { - const runningState = { - ...dirtyState, - _ui: { - ...dirtyState._ui, - runningPromptId: action.id, - }, + const replaceOutput = (statePrompt: ClientPrompt) => { + const responsePrompt = action.config.prompts.find( + (resPrompt) => resPrompt.name === statePrompt.name + ); + return { + // Note: Don't need to set `isRunning` or `cancellationToken` + // because we already call RUN_PROMPT earlier in `onRunPrompt` + ...statePrompt, + outputs: responsePrompt?.outputs, + } as ClientPrompt; }; - return reduceReplacePrompt(runningState, action.id, (prompt) => ({ - ...prompt, - _ui: { - ...prompt._ui, - cancellationToken: action.cancellationToken, - isRunning: true, - }, - })); + return reduceReplacePrompt( + dirtyState, + dirtyState._ui.runningPromptId as string, + replaceOutput + ); } case "STREAM_OUTPUT_CHUNK": { return reduceReplacePrompt(dirtyState, action.id, (prompt) => ({