[editor] Streaming Outputs (#798)

# [editor] Streaming Outputs

Add client-side handling for streaming outputs. We currently derive
whether the request should stream by checking settings. Not a huge fan
of this since:
- it couples client with config runtime logic for getting the 'stream'
value
- assumes 'stream' setting is used by the models

This is sufficient for now, I think. Note that I do think it's nice to
allow the editor user to toggle between stream/no stream since it will
be useful for debugging their model parser

## Testing:
- Error handling: added `return HttpResponseWithAIConfig(message="No
AIConfig loaded", code=400, aiconfig=None).to_flask_format()` in /run
with 'stream' set to true and ensured error is shown as notification and
run state is stopped on the client
- Non-streaming model (dall-e) run correctly
- Streaming works for models with stream setting set to tru/false:


https://github.com/lastmile-ai/aiconfig/assets/5060851/5242f49f-0003-45a6-af19-1e29c611eca4

python/src/aiconfig/editor/client/src/Editor.tsx

@@ -1,11 +1,19 @@
 import EditorContainer, {
  AIConfigCallbacks,
+ RunPromptStreamCallback,
 } from "./components/EditorContainer";
 import { Flex, Loader, MantineProvider, Image } from "@mantine/core";
-import { AIConfig, InferenceSettings, JSONObject, Prompt } from "aiconfig";
+import {
+ AIConfig,
+ InferenceSettings,
+ JSONObject,
+ Output,
+ Prompt,
+} from "aiconfig";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { ufetch } from "ufetch";
 import { ROUTE_TABLE } from "./utils/api";
+import { streamingApi } from "./utils/oboeHelpers";
 
 export default function Editor() {
  const [aiconfig, setAiConfig] = useState<AIConfig | undefined>();
@@ -66,6 +74,30 @@ export default function Editor() {
  });
  }, []);
 
+ const runPromptStream = useCallback(
+ async (promptName: string, onStream: RunPromptStreamCallback) => {
+ await streamingApi(
+ {
+ url: ROUTE_TABLE.RUN_PROMPT,
+ method: "POST",
+ body: {
+ prompt_name: promptName,
+ stream: true,
+ },
+ },
+ "output_chunk",
+ (data) => {
+ onStream({ type: "output_chunk", data: data as Output });
+ },
+ "aiconfig",
+ (data) => {
+ onStream({ type: "aiconfig", data: data as AIConfig });
+ }
+ );
+ },
+ []
+ );
+
  const updatePrompt = useCallback(
  async (promptName: string, promptData: Prompt) => {
  return await ufetch.post(ROUTE_TABLE.UPDATE_PROMPT, {
@@ -124,6 +156,7 @@ export default function Editor() {
  getModels,
  getServerStatus,
  runPrompt,
+ runPromptStream,
  save,
  setConfigDescription,
  setConfigName,
@@ -137,6 +170,7 @@ export default function Editor() {
  getModels,
  getServerStatus,
  runPrompt,
+ runPromptStream,
  save,
  setConfigDescription,
  setConfigName,
@@ -168,7 +202,7 @@ export default function Editor() {
  deg: 45,
  },
 
- globalStyles: (theme) => ({
+ globalStyles: () => ({
  ".editorBackground": {
  background:
  "radial-gradient(ellipse at top,#08122d,#030712),radial-gradient(ellipse at bottom,#030712,#030712)",

python/src/aiconfig/editor/client/src/components/EditorContainer.tsx

@@ -14,6 +14,7 @@ import {
  AIConfig,
  InferenceSettings,
  JSONObject,
+ Output,
  Prompt,
  PromptInput,
 } from "aiconfig";
@@ -36,6 +37,7 @@ import AddPromptButton from "./prompt/AddPromptButton";
 import {
  getDefaultNewPromptName,
  getPrompt,
+ isStreamingSupported,
 } from "../utils/aiconfigStateUtils";
 import { debounce, uniqueId } from "lodash";
 import PromptMenuButton from "./prompt/PromptMenuButton";
@@ -56,6 +58,18 @@ type Props = {
  callbacks: AIConfigCallbacks;
 };
 
+type RunPromptStreamEvent =
+ | {
+ type: "output_chunk";
+ data: Output;
+ }
+ | {
+ type: "aiconfig";
+ data: AIConfig;
+ };
+
+export type RunPromptStreamCallback = (event: RunPromptStreamEvent) => void;
+
 export type AIConfigCallbacks = {
  addPrompt: (
  promptName: string,
@@ -66,6 +80,10 @@ export type AIConfigCallbacks = {
  getModels: (search: string) => Promise<string[]>;
  getServerStatus?: () => Promise<{ status: "OK" | "ERROR" }>;
  runPrompt: (promptName: string) => Promise<{ aiconfig: AIConfig }>;
+ runPromptStream: (
+ promptName: string,
+ onStream: RunPromptStreamCallback
+ ) => Promise<void>;
  save: (aiconfig: AIConfig) => Promise<void>;
  setConfigDescription: (description: string) => Promise<void>;
  setConfigName: (name: string) => Promise<void>;
@@ -526,6 +544,8 @@ export default function EditorContainer({
  );
 
  const runPromptCallback = callbacks.runPrompt;
+ const runPromptStreamCallback = callbacks.runPromptStream;
+
  const onRunPrompt = useCallback(
  async (promptId: string) => {
  const action: AIConfigReducerAction = {
@@ -540,14 +560,36 @@ export default function EditorContainer({
  if (!statePrompt) {
  throw new Error(`Could not find prompt with id ${promptId}`);
  }
- const promptName = statePrompt.name;
- const serverConfigRes = await runPromptCallback(promptName);
 
- dispatch({
- type: "CONSOLIDATE_AICONFIG",
- action,
- config: serverConfigRes.aiconfig,
- });
+ const promptName = statePrompt.name;
+ const isStream = isStreamingSupported(statePrompt, stateRef.current);
+
+ if (isStream) {
+ await runPromptStreamCallback(promptName, (event) => {
+ if (event.type === "output_chunk") {
+ dispatch({
+ type: "STREAM_OUTPUT_CHUNK",
+ id: promptId,
+ output: event.data,
+ });
+ } else if (event.type === "aiconfig") {
+ dispatch({
+ type: "CONSOLIDATE_AICONFIG",
+ action,
+ config: event.data,
+ });
+ }
+ });
+ return;
+ } else {
+ const serverConfigRes = await runPromptCallback(promptName);
+
+ dispatch({
+ type: "CONSOLIDATE_AICONFIG",
+ action,
+ config: serverConfigRes.aiconfig,
+ });
+ }
  } catch (err: unknown) {
  const message = (err as RequestCallbackError).message ?? null;
 
@@ -567,7 +609,7 @@ export default function EditorContainer({
  });
  }
  },
- [runPromptCallback]
+ [runPromptCallback, runPromptStreamCallback]
  );
 
  const setNameCallback = callbacks.setConfigName;

python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts

@@ -1,6 +1,6 @@
 import { ClientAIConfig, ClientPrompt } from "../shared/types";
 import { getPromptModelName } from "../utils/promptUtils";
-import { AIConfig, JSONObject, PromptInput } from "aiconfig";
+import { AIConfig, JSONObject, Output, PromptInput } from "aiconfig";
 
 export type AIConfigReducerAction =
  | MutateAIConfigAction
@@ -14,6 +14,7 @@ export type MutateAIConfigAction =
  | RunPromptAction
  | SetDescriptionAction
  | SetNameAction
+ | StreamOutputChunkAction
  | UpdatePromptInputAction
  | UpdatePromptNameAction
  | UpdatePromptModelAction
@@ -63,6 +64,12 @@ export type SetNameAction = {
  name: string;
 };
 
+export type StreamOutputChunkAction = {
+ type: "STREAM_OUTPUT_CHUNK";
+ id: string;
+ output: Output;
+};
+
 export type UpdatePromptInputAction = {
  type: "UPDATE_PROMPT_INPUT";
  id: string;
@@ -255,6 +262,12 @@ export default function aiconfigReducer(
  name: action.name,
  };
  }
+ case "STREAM_OUTPUT_CHUNK": {
+ return reduceReplacePrompt(dirtyState, action.id, (prompt) => ({
+ ...prompt,
+ outputs: [action.output],
+ }));
+ }
  case "UPDATE_PROMPT_INPUT": {
  return reduceReplaceInput(dirtyState, action.id, () => action.input);
  }

python/src/aiconfig/editor/client/src/utils/aiconfigStateUtils.ts

@@ -1,5 +1,6 @@
 import { AIConfig } from "aiconfig";
 import { ClientAIConfig, ClientPrompt } from "../shared/types";
+import { getPromptModelName } from "./promptUtils";
 
 export function getDefaultNewPromptName(aiconfig: AIConfig): string {
  const existingNames = aiconfig.prompts.map((prompt) => prompt.name);
@@ -16,3 +17,31 @@ export function getPrompt(
 ): ClientPrompt | undefined {
  return aiconfig.prompts.find((prompt) => prompt._ui.id === id);
 }
+
+// TODO: This is pretty hacky. Streaming is actually part of AIConfig runtime and not necessarily part of model settings,
+// let alone required to be defined by 'stream' boolean... Ideally we should treat everything as stream but this should work for now.
+export function isStreamingSupported(
+ prompt: ClientPrompt,
+ config: ClientAIConfig
+): boolean {
+ const promptModelSettings =
+ prompt.metadata?.model && typeof prompt.metadata.model !== "string"
+ ? prompt.metadata.model?.settings
+ : undefined;
+ if (promptModelSettings) {
+ if (promptModelSettings?.stream === true) {
+ return true;
+ } else if (promptModelSettings?.stream === false) {
+ return false;
+ }
+ }
+
+ const promptModelName = getPromptModelName(prompt);
+ if (promptModelName) {
+ const globalModelSettings =
+ config.metadata?.models?.[promptModelName]?.settings;
+ return globalModelSettings?.stream === true;
+ }
+
+ return false;
+}

python/src/aiconfig/editor/client/src/utils/oboeHelpers.ts

@@ -7,11 +7,11 @@ import oboe, { Options } from "oboe";
 export async function streamingApi<T>(
  headers: Options,
  on: string = "*",
- fn: (data: any) => void,
+ fn: (data: unknown) => void,
  on2?: string,
- fn2?: (data: any) => void,
+ fn2?: (data: unknown) => void,
  on3?: string,
- fn3?: (data: any) => void
+ fn3?: (data: unknown) => void
 ): Promise<T> {
  return new Promise((resolve, reject) => {
  if (fn2 && on2 && fn3 && on3) {