mirror of
https://github.com/langchain-ai/langgraphjs.git
synced 2026-07-22 17:15:25 -04:00
feat(sdk): useStream with custom transport (#1675)
This commit is contained in:
David Duong
@@ -0,0 +1,5 @@
|
||||
---
|
||||
"@langchain/langgraph-sdk": patch
|
||||
---
|
||||
|
||||
Add `transport` option to useStream, allowing custom endpoints, that emit compatible Server-Sent Events to be used with `useStream`.
|
||||
@@ -1,2 +1,10 @@
|
||||
export { useStream } from "./stream.js";
|
||||
export type { MessageMetadata, UseStream, UseStreamOptions } from "./types.js";
|
||||
export { FetchStreamTransport } from "./stream.custom.js";
|
||||
export type {
|
||||
MessageMetadata,
|
||||
UseStream,
|
||||
UseStreamOptions,
|
||||
UseStreamCustom,
|
||||
UseStreamCustomOptions,
|
||||
UseStreamTransport,
|
||||
} from "./types.js";
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
/* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
|
||||
|
||||
"use client";
|
||||
|
||||
import { useState, useSyncExternalStore } from "react";
|
||||
import { EventStreamEvent, StreamManager } from "./manager.js";
|
||||
import type {
|
||||
BagTemplate,
|
||||
GetUpdateType,
|
||||
GetCustomEventType,
|
||||
GetInterruptType,
|
||||
RunCallbackMeta,
|
||||
GetConfigurableType,
|
||||
UseStreamCustomOptions,
|
||||
UseStreamCustom,
|
||||
UseStreamTransport,
|
||||
CustomSubmitOptions,
|
||||
} from "./types.js";
|
||||
import type { Message } from "../types.messages.js";
|
||||
import { MessageTupleManager } from "./messages.js";
|
||||
import { Interrupt } from "../schema.js";
|
||||
import { BytesLineDecoder, SSEDecoder } from "../utils/sse.js";
|
||||
import { IterableReadableStream } from "../utils/stream.js";
|
||||
import { Command } from "../types.js";
|
||||
|
||||
interface FetchStreamTransportOptions {
|
||||
/**
|
||||
* The URL of the API to use.
|
||||
*/
|
||||
apiUrl: string;
|
||||
|
||||
/**
|
||||
* Default headers to send with requests.
|
||||
*/
|
||||
defaultHeaders?: HeadersInit;
|
||||
|
||||
/**
|
||||
* Specify a custom fetch implementation.
|
||||
*/
|
||||
fetch?: typeof fetch | ((...args: any[]) => any); // eslint-disable-line @typescript-eslint/no-explicit-any
|
||||
|
||||
/**
|
||||
* Callback that is called before the request is made.
|
||||
*/
|
||||
onRequest?: (
|
||||
url: string,
|
||||
init: RequestInit
|
||||
) => Promise<RequestInit> | RequestInit;
|
||||
}
|
||||
|
||||
export class FetchStreamTransport<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends BagTemplate = BagTemplate
|
||||
> implements UseStreamTransport<StateType, Bag>
|
||||
{
|
||||
constructor(private readonly options: FetchStreamTransportOptions) {}
|
||||
|
||||
async stream(payload: {
|
||||
input: GetUpdateType<Bag, StateType> | null | undefined;
|
||||
context: GetConfigurableType<Bag> | undefined;
|
||||
command: Command | undefined;
|
||||
signal: AbortSignal;
|
||||
}): Promise<AsyncGenerator<{ id?: string; event: string; data: unknown }>> {
|
||||
const { signal, ...body } = payload;
|
||||
|
||||
let requestInit: RequestInit = {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
...this.options.defaultHeaders,
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
signal,
|
||||
};
|
||||
|
||||
if (this.options.onRequest) {
|
||||
requestInit = await this.options.onRequest(
|
||||
this.options.apiUrl,
|
||||
requestInit
|
||||
);
|
||||
}
|
||||
const fetchFn = this.options.fetch ?? fetch;
|
||||
|
||||
const response = await fetchFn(this.options.apiUrl, requestInit);
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to stream: ${response.statusText}`);
|
||||
}
|
||||
|
||||
const stream = (
|
||||
response.body || new ReadableStream({ start: (ctrl) => ctrl.close() })
|
||||
)
|
||||
.pipeThrough(BytesLineDecoder())
|
||||
.pipeThrough(SSEDecoder());
|
||||
|
||||
return IterableReadableStream.fromReadableStream(stream);
|
||||
}
|
||||
}
|
||||
|
||||
export function useStreamCustom<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
ConfigurableType?: Record<string, unknown>;
|
||||
InterruptType?: unknown;
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(
|
||||
options: UseStreamCustomOptions<StateType, Bag>
|
||||
): UseStreamCustom<StateType, Bag> {
|
||||
type UpdateType = GetUpdateType<Bag, StateType>;
|
||||
type CustomType = GetCustomEventType<Bag>;
|
||||
type InterruptType = GetInterruptType<Bag>;
|
||||
type ConfigurableType = GetConfigurableType<Bag>;
|
||||
|
||||
const [messageManager] = useState(() => new MessageTupleManager());
|
||||
const [stream] = useState(
|
||||
() => new StreamManager<StateType, Bag>(messageManager)
|
||||
);
|
||||
|
||||
useSyncExternalStore(
|
||||
stream.subscribe,
|
||||
stream.getSnapshot,
|
||||
stream.getSnapshot
|
||||
);
|
||||
|
||||
const getMessages = (value: StateType): Message[] => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return Array.isArray(value[messagesKey]) ? value[messagesKey] : [];
|
||||
};
|
||||
|
||||
const setMessages = (current: StateType, messages: Message[]): StateType => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return { ...current, [messagesKey]: messages };
|
||||
};
|
||||
|
||||
const historyValues = options.initialValues ?? ({} as StateType);
|
||||
|
||||
const stop = () => stream.stop(historyValues, { onStop: options.onStop });
|
||||
|
||||
const submit = async (
|
||||
values: UpdateType | null | undefined,
|
||||
submitOptions?: CustomSubmitOptions<StateType, ConfigurableType>
|
||||
) => {
|
||||
let callbackMeta: RunCallbackMeta | undefined;
|
||||
|
||||
stream.setStreamValues(() => {
|
||||
if (submitOptions?.optimisticValues != null) {
|
||||
return {
|
||||
...historyValues,
|
||||
...(typeof submitOptions.optimisticValues === "function"
|
||||
? submitOptions.optimisticValues(historyValues)
|
||||
: submitOptions.optimisticValues),
|
||||
};
|
||||
}
|
||||
|
||||
return { ...historyValues };
|
||||
});
|
||||
|
||||
await stream.start(
|
||||
async (signal: AbortSignal) =>
|
||||
options.transport.stream({
|
||||
input: values,
|
||||
context: submitOptions?.context,
|
||||
command: submitOptions?.command,
|
||||
signal,
|
||||
}) as Promise<
|
||||
AsyncGenerator<EventStreamEvent<StateType, UpdateType, CustomType>>
|
||||
>,
|
||||
{
|
||||
getMessages,
|
||||
setMessages,
|
||||
|
||||
initialValues: {} as StateType,
|
||||
callbacks: options,
|
||||
|
||||
onSuccess: () => undefined,
|
||||
onError(error) {
|
||||
options.onError?.(error, callbackMeta);
|
||||
},
|
||||
}
|
||||
);
|
||||
};
|
||||
|
||||
return {
|
||||
get values() {
|
||||
return stream.values ?? ({} as StateType);
|
||||
},
|
||||
|
||||
error: stream.error,
|
||||
isLoading: stream.isLoading,
|
||||
|
||||
stop,
|
||||
submit,
|
||||
|
||||
get interrupt(): Interrupt<InterruptType> | undefined {
|
||||
if (
|
||||
stream.values != null &&
|
||||
"__interrupt__" in stream.values &&
|
||||
Array.isArray(stream.values.__interrupt__)
|
||||
) {
|
||||
const valueInterrupts = stream.values.__interrupt__;
|
||||
if (valueInterrupts.length === 0) return { when: "breakpoint" };
|
||||
if (valueInterrupts.length === 1) return valueInterrupts[0];
|
||||
|
||||
// TODO: fix the typing of interrupts if multiple interrupts are returned
|
||||
return valueInterrupts as Interrupt<InterruptType>;
|
||||
}
|
||||
|
||||
return undefined;
|
||||
},
|
||||
|
||||
get messages() {
|
||||
if (!stream.values) return [];
|
||||
return getMessages(stream.values);
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,652 @@
|
||||
/* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
|
||||
|
||||
"use client";
|
||||
|
||||
import {
|
||||
type RefObject,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
useSyncExternalStore,
|
||||
} from "react";
|
||||
import { findLast, unique } from "./utils.js";
|
||||
import { StreamError } from "./errors.js";
|
||||
import { getBranchContext } from "./branching.js";
|
||||
import { EventStreamEvent, StreamManager } from "./manager.js";
|
||||
import type {
|
||||
BagTemplate,
|
||||
UseStreamOptions,
|
||||
UseStream,
|
||||
GetUpdateType,
|
||||
GetCustomEventType,
|
||||
GetInterruptType,
|
||||
GetConfigurableType,
|
||||
RunCallbackMeta,
|
||||
SubmitOptions,
|
||||
MessageMetadata,
|
||||
} from "./types.js";
|
||||
import { Client, getClientConfigHash } from "../client.js";
|
||||
import type { Message } from "../types.messages.js";
|
||||
import type { Interrupt, ThreadState } from "../schema.js";
|
||||
import type { StreamMode } from "../types.stream.js";
|
||||
import { MessageTupleManager } from "./messages.js";
|
||||
|
||||
function fetchHistory<StateType extends Record<string, unknown>>(
|
||||
client: Client,
|
||||
threadId: string,
|
||||
options?: { limit?: boolean | number }
|
||||
) {
|
||||
if (options?.limit === false) {
|
||||
return client.threads.getState<StateType>(threadId).then((state) => {
|
||||
if (state.checkpoint == null) return [];
|
||||
return [state];
|
||||
});
|
||||
}
|
||||
|
||||
const limit = typeof options?.limit === "number" ? options.limit : 10;
|
||||
return client.threads.getHistory<StateType>(threadId, { limit });
|
||||
}
|
||||
|
||||
function useThreadHistory<StateType extends Record<string, unknown>>(
|
||||
threadId: string | undefined | null,
|
||||
client: Client,
|
||||
limit: boolean | number,
|
||||
clearCallbackRef: RefObject<(() => void) | undefined>,
|
||||
submittingRef: RefObject<boolean>,
|
||||
onErrorRef: RefObject<((error: unknown) => void) | undefined>
|
||||
) {
|
||||
const [history, setHistory] = useState<ThreadState<StateType>[] | undefined>(
|
||||
undefined
|
||||
);
|
||||
const [isLoading, setIsLoading] = useState(() => {
|
||||
if (threadId == null) return false;
|
||||
return true;
|
||||
});
|
||||
const [error, setError] = useState<unknown | undefined>(undefined);
|
||||
|
||||
const clientHash = getClientConfigHash(client);
|
||||
const clientRef = useRef(client);
|
||||
clientRef.current = client;
|
||||
|
||||
const fetcher = useCallback(
|
||||
(
|
||||
threadId: string | undefined | null
|
||||
): Promise<ThreadState<StateType>[]> => {
|
||||
if (threadId != null) {
|
||||
const client = clientRef.current;
|
||||
|
||||
setIsLoading(true);
|
||||
return fetchHistory<StateType>(client, threadId, {
|
||||
limit,
|
||||
})
|
||||
.then(
|
||||
(history) => {
|
||||
setHistory(history);
|
||||
return history;
|
||||
},
|
||||
(error) => {
|
||||
setError(error);
|
||||
onErrorRef.current?.(error);
|
||||
return Promise.reject(error);
|
||||
}
|
||||
)
|
||||
.finally(() => {
|
||||
setIsLoading(false);
|
||||
});
|
||||
}
|
||||
|
||||
setHistory(undefined);
|
||||
setError(undefined);
|
||||
setIsLoading(false);
|
||||
|
||||
clearCallbackRef.current?.();
|
||||
return Promise.resolve([]);
|
||||
},
|
||||
[clearCallbackRef, onErrorRef, limit]
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (submittingRef.current) return;
|
||||
void fetcher(threadId);
|
||||
}, [fetcher, submittingRef, clientHash, limit, threadId]);
|
||||
|
||||
return {
|
||||
data: history,
|
||||
isLoading,
|
||||
error,
|
||||
mutate: (mutateId?: string) => fetcher(mutateId ?? threadId),
|
||||
};
|
||||
}
|
||||
|
||||
const useControllableThreadId = (options?: {
|
||||
threadId?: string | null;
|
||||
onThreadId?: (threadId: string) => void;
|
||||
}): [string | null, (threadId: string) => void] => {
|
||||
const [localThreadId, _setLocalThreadId] = useState<string | null>(
|
||||
options?.threadId ?? null
|
||||
);
|
||||
|
||||
const onThreadIdRef = useRef(options?.onThreadId);
|
||||
onThreadIdRef.current = options?.onThreadId;
|
||||
|
||||
const onThreadId = useCallback((threadId: string) => {
|
||||
_setLocalThreadId(threadId);
|
||||
onThreadIdRef.current?.(threadId);
|
||||
}, []);
|
||||
|
||||
if (!options || !("threadId" in options)) {
|
||||
return [localThreadId, onThreadId];
|
||||
}
|
||||
|
||||
return [options.threadId ?? null, onThreadId];
|
||||
};
|
||||
|
||||
export function useStreamLGP<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
ConfigurableType?: Record<string, unknown>;
|
||||
InterruptType?: unknown;
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(options: UseStreamOptions<StateType, Bag>): UseStream<StateType, Bag> {
|
||||
type UpdateType = GetUpdateType<Bag, StateType>;
|
||||
type CustomType = GetCustomEventType<Bag>;
|
||||
type InterruptType = GetInterruptType<Bag>;
|
||||
type ConfigurableType = GetConfigurableType<Bag>;
|
||||
|
||||
const reconnectOnMountRef = useRef(options.reconnectOnMount);
|
||||
const runMetadataStorage = useMemo(() => {
|
||||
if (typeof window === "undefined") return null;
|
||||
const storage = reconnectOnMountRef.current;
|
||||
if (storage === true) return window.sessionStorage;
|
||||
if (typeof storage === "function") return storage();
|
||||
return null;
|
||||
}, []);
|
||||
|
||||
const client = useMemo(
|
||||
() =>
|
||||
options.client ??
|
||||
new Client({
|
||||
apiUrl: options.apiUrl,
|
||||
apiKey: options.apiKey,
|
||||
callerOptions: options.callerOptions,
|
||||
defaultHeaders: options.defaultHeaders,
|
||||
}),
|
||||
[
|
||||
options.client,
|
||||
options.apiKey,
|
||||
options.apiUrl,
|
||||
options.callerOptions,
|
||||
options.defaultHeaders,
|
||||
]
|
||||
);
|
||||
|
||||
const [messageManager] = useState(() => new MessageTupleManager());
|
||||
const [stream] = useState(
|
||||
() => new StreamManager<StateType, Bag>(messageManager)
|
||||
);
|
||||
|
||||
useSyncExternalStore(
|
||||
stream.subscribe,
|
||||
stream.getSnapshot,
|
||||
stream.getSnapshot
|
||||
);
|
||||
|
||||
const [threadId, onThreadId] = useControllableThreadId(options);
|
||||
const trackStreamModeRef = useRef<Exclude<StreamMode, "messages">[]>([]);
|
||||
|
||||
const trackStreamMode = useCallback(
|
||||
(...mode: Exclude<StreamMode, "messages">[]) => {
|
||||
const ref = trackStreamModeRef.current;
|
||||
for (const m of mode) {
|
||||
if (!ref.includes(m)) ref.push(m);
|
||||
}
|
||||
},
|
||||
[]
|
||||
);
|
||||
|
||||
const hasUpdateListener = options.onUpdateEvent != null;
|
||||
const hasCustomListener = options.onCustomEvent != null;
|
||||
const hasLangChainListener = options.onLangChainEvent != null;
|
||||
const hasDebugListener = options.onDebugEvent != null;
|
||||
const hasCheckpointListener = options.onCheckpointEvent != null;
|
||||
const hasTaskListener = options.onTaskEvent != null;
|
||||
|
||||
const callbackStreamMode = useMemo(() => {
|
||||
const modes: Exclude<StreamMode, "messages">[] = [];
|
||||
if (hasUpdateListener) modes.push("updates");
|
||||
if (hasCustomListener) modes.push("custom");
|
||||
if (hasLangChainListener) modes.push("events");
|
||||
if (hasDebugListener) modes.push("debug");
|
||||
if (hasCheckpointListener) modes.push("checkpoints");
|
||||
if (hasTaskListener) modes.push("tasks");
|
||||
return modes;
|
||||
}, [
|
||||
hasUpdateListener,
|
||||
hasCustomListener,
|
||||
hasLangChainListener,
|
||||
hasDebugListener,
|
||||
hasCheckpointListener,
|
||||
hasTaskListener,
|
||||
]);
|
||||
|
||||
const clearCallbackRef = useRef<() => void>(null!);
|
||||
clearCallbackRef.current = stream.clear;
|
||||
|
||||
const submittingRef = useRef(false);
|
||||
submittingRef.current = stream.isLoading;
|
||||
|
||||
const onErrorRef = useRef<
|
||||
((error: unknown, run?: RunCallbackMeta) => void) | undefined
|
||||
>(undefined);
|
||||
onErrorRef.current = options.onError;
|
||||
|
||||
const historyLimit =
|
||||
typeof options.fetchStateHistory === "object" &&
|
||||
options.fetchStateHistory != null
|
||||
? options.fetchStateHistory.limit ?? false
|
||||
: options.fetchStateHistory ?? false;
|
||||
|
||||
const history = useThreadHistory<StateType>(
|
||||
threadId,
|
||||
client,
|
||||
historyLimit,
|
||||
clearCallbackRef,
|
||||
submittingRef,
|
||||
onErrorRef
|
||||
);
|
||||
|
||||
const getMessages = (value: StateType): Message[] => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return Array.isArray(value[messagesKey]) ? value[messagesKey] : [];
|
||||
};
|
||||
|
||||
const setMessages = (current: StateType, messages: Message[]): StateType => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return { ...current, [messagesKey]: messages };
|
||||
};
|
||||
|
||||
const [branch, setBranch] = useState<string>("");
|
||||
const branchContext = getBranchContext(branch, history.data);
|
||||
|
||||
const historyValues =
|
||||
branchContext.threadHead?.values ??
|
||||
options.initialValues ??
|
||||
({} as StateType);
|
||||
|
||||
const historyError = (() => {
|
||||
const error = branchContext.threadHead?.tasks?.at(-1)?.error;
|
||||
if (error == null) return undefined;
|
||||
try {
|
||||
const parsed = JSON.parse(error) as unknown;
|
||||
if (StreamError.isStructuredError(parsed)) return new StreamError(parsed);
|
||||
return parsed;
|
||||
} catch {
|
||||
// do nothing
|
||||
}
|
||||
return error;
|
||||
})();
|
||||
|
||||
const messageMetadata = (() => {
|
||||
const alreadyShown = new Set<string>();
|
||||
return getMessages(historyValues).map(
|
||||
(message, idx): Omit<MessageMetadata<StateType>, "streamMetadata"> => {
|
||||
const messageId = message.id ?? idx;
|
||||
|
||||
// Find the first checkpoint where the message was seen
|
||||
const firstSeenState = findLast(history.data ?? [], (state) =>
|
||||
getMessages(state.values)
|
||||
.map((m, idx) => m.id ?? idx)
|
||||
.includes(messageId)
|
||||
);
|
||||
|
||||
const checkpointId = firstSeenState?.checkpoint?.checkpoint_id;
|
||||
let branch =
|
||||
checkpointId != null
|
||||
? branchContext.branchByCheckpoint[checkpointId]
|
||||
: undefined;
|
||||
if (!branch?.branch?.length) branch = undefined;
|
||||
|
||||
// serialize branches
|
||||
const optionsShown = branch?.branchOptions?.flat(2).join(",");
|
||||
if (optionsShown) {
|
||||
if (alreadyShown.has(optionsShown)) branch = undefined;
|
||||
alreadyShown.add(optionsShown);
|
||||
}
|
||||
|
||||
return {
|
||||
messageId: messageId.toString(),
|
||||
firstSeenState,
|
||||
|
||||
branch: branch?.branch,
|
||||
branchOptions: branch?.branchOptions,
|
||||
};
|
||||
}
|
||||
);
|
||||
})();
|
||||
|
||||
const stop = () => stream.stop(historyValues, { onStop: options.onStop });
|
||||
|
||||
// --- TRANSPORT ---
|
||||
const submit = async (
|
||||
values: UpdateType | null | undefined,
|
||||
submitOptions?: SubmitOptions<StateType, ConfigurableType>
|
||||
) => {
|
||||
// Unbranch things
|
||||
const checkpointId = submitOptions?.checkpoint?.checkpoint_id;
|
||||
setBranch(
|
||||
checkpointId != null
|
||||
? branchContext.branchByCheckpoint[checkpointId]?.branch ?? ""
|
||||
: ""
|
||||
);
|
||||
|
||||
stream.setStreamValues(() => {
|
||||
if (submitOptions?.optimisticValues != null) {
|
||||
return {
|
||||
...historyValues,
|
||||
...(typeof submitOptions.optimisticValues === "function"
|
||||
? submitOptions.optimisticValues(historyValues)
|
||||
: submitOptions.optimisticValues),
|
||||
};
|
||||
}
|
||||
|
||||
return { ...historyValues };
|
||||
});
|
||||
|
||||
// When `fetchStateHistory` is requested, thus we assume that branching
|
||||
// is enabled. We then need to include the implicit branch.
|
||||
const includeImplicitBranch =
|
||||
historyLimit === true || typeof historyLimit === "number";
|
||||
|
||||
let callbackMeta: RunCallbackMeta | undefined;
|
||||
let rejoinKey: `lg:stream:${string}` | undefined;
|
||||
let usableThreadId = threadId;
|
||||
|
||||
await stream.start(
|
||||
async (signal: AbortSignal) => {
|
||||
if (!usableThreadId) {
|
||||
const thread = await client.threads.create({
|
||||
threadId: submitOptions?.threadId,
|
||||
metadata: submitOptions?.metadata,
|
||||
});
|
||||
onThreadId(thread.thread_id);
|
||||
usableThreadId = thread.thread_id;
|
||||
}
|
||||
|
||||
if (!usableThreadId) {
|
||||
throw new Error("Failed to obtain valid thread ID.");
|
||||
}
|
||||
|
||||
const streamMode = unique([
|
||||
...(submitOptions?.streamMode ?? []),
|
||||
...trackStreamModeRef.current,
|
||||
...callbackStreamMode,
|
||||
]);
|
||||
|
||||
let checkpoint =
|
||||
submitOptions?.checkpoint ??
|
||||
(includeImplicitBranch
|
||||
? branchContext.threadHead?.checkpoint
|
||||
: undefined) ??
|
||||
undefined;
|
||||
|
||||
// Avoid specifying a checkpoint if user explicitly set it to null
|
||||
if (submitOptions?.checkpoint === null) checkpoint = undefined;
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
|
||||
// @ts-expect-error
|
||||
if (checkpoint != null) delete checkpoint.thread_id;
|
||||
const streamResumable =
|
||||
submitOptions?.streamResumable ?? !!runMetadataStorage;
|
||||
|
||||
return client.runs.stream(usableThreadId, options.assistantId, {
|
||||
input: values as Record<string, unknown>,
|
||||
config: submitOptions?.config,
|
||||
context: submitOptions?.context,
|
||||
command: submitOptions?.command,
|
||||
|
||||
interruptBefore: submitOptions?.interruptBefore,
|
||||
interruptAfter: submitOptions?.interruptAfter,
|
||||
metadata: submitOptions?.metadata,
|
||||
multitaskStrategy: submitOptions?.multitaskStrategy,
|
||||
onCompletion: submitOptions?.onCompletion,
|
||||
onDisconnect:
|
||||
submitOptions?.onDisconnect ??
|
||||
(streamResumable ? "continue" : "cancel"),
|
||||
|
||||
signal,
|
||||
|
||||
checkpoint,
|
||||
streamMode,
|
||||
streamSubgraphs: submitOptions?.streamSubgraphs,
|
||||
streamResumable,
|
||||
durability: submitOptions?.durability,
|
||||
onRunCreated(params) {
|
||||
callbackMeta = {
|
||||
run_id: params.run_id,
|
||||
thread_id: params.thread_id ?? usableThreadId!,
|
||||
};
|
||||
|
||||
if (runMetadataStorage) {
|
||||
rejoinKey = `lg:stream:${usableThreadId}`;
|
||||
runMetadataStorage.setItem(rejoinKey, callbackMeta.run_id);
|
||||
}
|
||||
|
||||
options.onCreated?.(callbackMeta);
|
||||
},
|
||||
}) as AsyncGenerator<
|
||||
EventStreamEvent<StateType, UpdateType, CustomType>
|
||||
>;
|
||||
},
|
||||
{
|
||||
getMessages,
|
||||
setMessages,
|
||||
|
||||
initialValues: historyValues,
|
||||
callbacks: options,
|
||||
|
||||
async onSuccess() {
|
||||
if (rejoinKey) runMetadataStorage?.removeItem(rejoinKey);
|
||||
const shouldRefetch =
|
||||
// We're expecting the whole thread state in onFinish
|
||||
options.onFinish != null ||
|
||||
// We're fetching history, thus we need the latest checkpoint
|
||||
// to ensure we're not accidentally submitting to a wrong branch
|
||||
includeImplicitBranch;
|
||||
|
||||
if (shouldRefetch) {
|
||||
const newHistory = await history.mutate(usableThreadId!);
|
||||
const lastHead = newHistory.at(0);
|
||||
if (lastHead) {
|
||||
// We now have the latest update from /history
|
||||
// Thus we can clear the local stream state
|
||||
options.onFinish?.(lastHead, callbackMeta);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
},
|
||||
onError(error) {
|
||||
options.onError?.(error, callbackMeta);
|
||||
},
|
||||
}
|
||||
);
|
||||
};
|
||||
|
||||
const joinStream = async (
|
||||
runId: string,
|
||||
lastEventId?: string,
|
||||
joinOptions?: { streamMode?: StreamMode | StreamMode[] }
|
||||
) => {
|
||||
// eslint-disable-next-line no-param-reassign
|
||||
lastEventId ??= "-1";
|
||||
if (!threadId) return;
|
||||
|
||||
const callbackMeta: RunCallbackMeta = {
|
||||
thread_id: threadId,
|
||||
run_id: runId,
|
||||
};
|
||||
|
||||
await stream.start(
|
||||
async (signal: AbortSignal) => {
|
||||
return client.runs.joinStream(threadId, runId, {
|
||||
signal,
|
||||
lastEventId,
|
||||
streamMode: joinOptions?.streamMode,
|
||||
}) as AsyncGenerator<
|
||||
EventStreamEvent<StateType, UpdateType, CustomType>
|
||||
>;
|
||||
},
|
||||
{
|
||||
getMessages,
|
||||
setMessages,
|
||||
|
||||
initialValues: historyValues,
|
||||
callbacks: options,
|
||||
async onSuccess() {
|
||||
runMetadataStorage?.removeItem(`lg:stream:${threadId}`);
|
||||
const newHistory = await history.mutate(threadId);
|
||||
const lastHead = newHistory.at(0);
|
||||
if (lastHead) options.onFinish?.(lastHead, callbackMeta);
|
||||
},
|
||||
onError(error) {
|
||||
options.onError?.(error, callbackMeta);
|
||||
},
|
||||
}
|
||||
);
|
||||
};
|
||||
|
||||
const reconnectKey = useMemo(() => {
|
||||
if (!runMetadataStorage || stream.isLoading) return undefined;
|
||||
if (typeof window === "undefined") return undefined;
|
||||
const runId = runMetadataStorage?.getItem(`lg:stream:${threadId}`);
|
||||
if (!runId) return undefined;
|
||||
return { runId, threadId };
|
||||
}, [runMetadataStorage, stream.isLoading, threadId]);
|
||||
|
||||
const shouldReconnect = !!runMetadataStorage;
|
||||
const reconnectRef = useRef({ threadId, shouldReconnect });
|
||||
|
||||
const joinStreamRef = useRef<typeof joinStream>(joinStream);
|
||||
joinStreamRef.current = joinStream;
|
||||
|
||||
useEffect(() => {
|
||||
// reset shouldReconnect when switching threads
|
||||
if (reconnectRef.current.threadId !== threadId) {
|
||||
reconnectRef.current = { threadId, shouldReconnect };
|
||||
}
|
||||
}, [threadId, shouldReconnect]);
|
||||
|
||||
useEffect(() => {
|
||||
if (reconnectKey && reconnectRef.current.shouldReconnect) {
|
||||
reconnectRef.current.shouldReconnect = false;
|
||||
void joinStreamRef.current?.(reconnectKey.runId);
|
||||
}
|
||||
}, [reconnectKey]);
|
||||
// --- END TRANSPORT ---
|
||||
|
||||
const error = stream.error ?? historyError ?? history.error;
|
||||
const values = stream.values ?? historyValues;
|
||||
|
||||
return {
|
||||
get values() {
|
||||
trackStreamMode("values");
|
||||
return values;
|
||||
},
|
||||
|
||||
client,
|
||||
assistantId: options.assistantId,
|
||||
|
||||
error,
|
||||
isLoading: stream.isLoading,
|
||||
|
||||
stop,
|
||||
submit,
|
||||
|
||||
joinStream,
|
||||
|
||||
branch,
|
||||
setBranch,
|
||||
|
||||
get history() {
|
||||
if (historyLimit === false) {
|
||||
throw new Error(
|
||||
"`fetchStateHistory` must be set to `true` to use `history`"
|
||||
);
|
||||
}
|
||||
|
||||
return branchContext.flatHistory;
|
||||
},
|
||||
|
||||
isThreadLoading: history.isLoading && history.data == null,
|
||||
|
||||
get experimental_branchTree() {
|
||||
if (historyLimit === false) {
|
||||
throw new Error(
|
||||
"`fetchStateHistory` must be set to `true` to use `experimental_branchTree`"
|
||||
);
|
||||
}
|
||||
|
||||
return branchContext.branchTree;
|
||||
},
|
||||
|
||||
get interrupt() {
|
||||
if (
|
||||
values != null &&
|
||||
"__interrupt__" in values &&
|
||||
Array.isArray(values.__interrupt__)
|
||||
) {
|
||||
const valueInterrupts = values.__interrupt__;
|
||||
if (valueInterrupts.length === 0) return { when: "breakpoint" };
|
||||
if (valueInterrupts.length === 1) return valueInterrupts[0];
|
||||
|
||||
// TODO: fix the typing of interrupts if multiple interrupts are returned
|
||||
return valueInterrupts;
|
||||
}
|
||||
|
||||
// If we're deferring to old interrupt detection logic, don't show the interrupt if the stream is loading
|
||||
if (stream.isLoading) return undefined;
|
||||
|
||||
const interrupts = branchContext.threadHead?.tasks?.at(-1)?.interrupts;
|
||||
if (interrupts == null || interrupts.length === 0) {
|
||||
// check if there's a next task present
|
||||
const next = branchContext.threadHead?.next ?? [];
|
||||
if (!next.length || error != null) return undefined;
|
||||
return { when: "breakpoint" };
|
||||
}
|
||||
|
||||
// Return only the current interrupt
|
||||
return interrupts.at(-1) as Interrupt<InterruptType> | undefined;
|
||||
},
|
||||
|
||||
get messages() {
|
||||
trackStreamMode("messages-tuple", "values");
|
||||
return getMessages(values);
|
||||
},
|
||||
|
||||
getMessagesMetadata(
|
||||
message: Message,
|
||||
index?: number
|
||||
): MessageMetadata<StateType> | undefined {
|
||||
trackStreamMode("values");
|
||||
|
||||
const streamMetadata = messageManager.get(message.id)?.metadata;
|
||||
const historyMetadata = messageMetadata?.find(
|
||||
(m) => m.messageId === (message.id ?? index)
|
||||
);
|
||||
|
||||
if (streamMetadata != null || historyMetadata != null) {
|
||||
return {
|
||||
...historyMetadata,
|
||||
streamMetadata,
|
||||
} as MessageMetadata<StateType>;
|
||||
}
|
||||
|
||||
return undefined;
|
||||
},
|
||||
};
|
||||
}
|
||||
+57
-636
@@ -1,148 +1,30 @@
|
||||
/* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
|
||||
|
||||
"use client";
|
||||
|
||||
import { useState } from "react";
|
||||
import { useStreamLGP } from "./stream.lgp.js";
|
||||
import { useStreamCustom } from "./stream.custom.js";
|
||||
import {
|
||||
type RefObject,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
useSyncExternalStore,
|
||||
} from "react";
|
||||
import { findLast, unique } from "./utils.js";
|
||||
import { StreamError } from "./errors.js";
|
||||
import { getBranchContext } from "./branching.js";
|
||||
import { EventStreamEvent, StreamManager } from "./manager.js";
|
||||
import type {
|
||||
BagTemplate,
|
||||
UseStreamOptions,
|
||||
UseStream,
|
||||
GetUpdateType,
|
||||
GetCustomEventType,
|
||||
GetInterruptType,
|
||||
GetConfigurableType,
|
||||
RunCallbackMeta,
|
||||
SubmitOptions,
|
||||
MessageMetadata,
|
||||
UseStreamCustom,
|
||||
UseStreamCustomOptions,
|
||||
UseStreamOptions,
|
||||
} from "./types.js";
|
||||
import { Client, getClientConfigHash } from "../client.js";
|
||||
import type { Message } from "../types.messages.js";
|
||||
import type { Interrupt, ThreadState } from "../schema.js";
|
||||
import type { StreamMode } from "../types.stream.js";
|
||||
import { MessageTupleManager } from "./messages.js";
|
||||
|
||||
function fetchHistory<StateType extends Record<string, unknown>>(
|
||||
client: Client,
|
||||
threadId: string,
|
||||
options?: { limit?: boolean | number }
|
||||
) {
|
||||
if (options?.limit === false) {
|
||||
return client.threads.getState<StateType>(threadId).then((state) => {
|
||||
if (state.checkpoint == null) return [];
|
||||
return [state];
|
||||
});
|
||||
}
|
||||
|
||||
const limit = typeof options?.limit === "number" ? options.limit : 10;
|
||||
return client.threads.getHistory<StateType>(threadId, { limit });
|
||||
function isCustomOptions<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
ConfigurableType?: Record<string, unknown>;
|
||||
InterruptType?: unknown;
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(
|
||||
options:
|
||||
| UseStreamOptions<StateType, Bag>
|
||||
| UseStreamCustomOptions<StateType, Bag>
|
||||
): options is UseStreamCustomOptions<StateType, Bag> {
|
||||
return "transport" in options;
|
||||
}
|
||||
|
||||
function useThreadHistory<StateType extends Record<string, unknown>>(
|
||||
threadId: string | undefined | null,
|
||||
client: Client,
|
||||
limit: boolean | number,
|
||||
clearCallbackRef: RefObject<(() => void) | undefined>,
|
||||
submittingRef: RefObject<boolean>,
|
||||
onErrorRef: RefObject<((error: unknown) => void) | undefined>
|
||||
) {
|
||||
const [history, setHistory] = useState<ThreadState<StateType>[] | undefined>(
|
||||
undefined
|
||||
);
|
||||
const [isLoading, setIsLoading] = useState(() => {
|
||||
if (threadId == null) return false;
|
||||
return true;
|
||||
});
|
||||
const [error, setError] = useState<unknown | undefined>(undefined);
|
||||
|
||||
const clientHash = getClientConfigHash(client);
|
||||
const clientRef = useRef(client);
|
||||
clientRef.current = client;
|
||||
|
||||
const fetcher = useCallback(
|
||||
(
|
||||
threadId: string | undefined | null
|
||||
): Promise<ThreadState<StateType>[]> => {
|
||||
if (threadId != null) {
|
||||
const client = clientRef.current;
|
||||
|
||||
setIsLoading(true);
|
||||
return fetchHistory<StateType>(client, threadId, {
|
||||
limit,
|
||||
})
|
||||
.then(
|
||||
(history) => {
|
||||
setHistory(history);
|
||||
return history;
|
||||
},
|
||||
(error) => {
|
||||
setError(error);
|
||||
onErrorRef.current?.(error);
|
||||
return Promise.reject(error);
|
||||
}
|
||||
)
|
||||
.finally(() => {
|
||||
setIsLoading(false);
|
||||
});
|
||||
}
|
||||
|
||||
setHistory(undefined);
|
||||
setError(undefined);
|
||||
setIsLoading(false);
|
||||
|
||||
clearCallbackRef.current?.();
|
||||
return Promise.resolve([]);
|
||||
},
|
||||
[clearCallbackRef, onErrorRef, limit]
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (submittingRef.current) return;
|
||||
void fetcher(threadId);
|
||||
}, [fetcher, submittingRef, clientHash, limit, threadId]);
|
||||
|
||||
return {
|
||||
data: history,
|
||||
isLoading,
|
||||
error,
|
||||
mutate: (mutateId?: string) => fetcher(mutateId ?? threadId),
|
||||
};
|
||||
}
|
||||
|
||||
const useControllableThreadId = (options?: {
|
||||
threadId?: string | null;
|
||||
onThreadId?: (threadId: string) => void;
|
||||
}): [string | null, (threadId: string) => void] => {
|
||||
const [localThreadId, _setLocalThreadId] = useState<string | null>(
|
||||
options?.threadId ?? null
|
||||
);
|
||||
|
||||
const onThreadIdRef = useRef(options?.onThreadId);
|
||||
onThreadIdRef.current = options?.onThreadId;
|
||||
|
||||
const onThreadId = useCallback((threadId: string) => {
|
||||
_setLocalThreadId(threadId);
|
||||
onThreadIdRef.current?.(threadId);
|
||||
}, []);
|
||||
|
||||
if (!options || !("threadId" in options)) {
|
||||
return [localThreadId, onThreadId];
|
||||
}
|
||||
|
||||
return [options.threadId ?? null, onThreadId];
|
||||
};
|
||||
|
||||
export function useStream<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
@@ -151,502 +33,41 @@ export function useStream<
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(options: UseStreamOptions<StateType, Bag>): UseStream<StateType, Bag> {
|
||||
type UpdateType = GetUpdateType<Bag, StateType>;
|
||||
type CustomType = GetCustomEventType<Bag>;
|
||||
type InterruptType = GetInterruptType<Bag>;
|
||||
type ConfigurableType = GetConfigurableType<Bag>;
|
||||
|
||||
const reconnectOnMountRef = useRef(options.reconnectOnMount);
|
||||
const runMetadataStorage = useMemo(() => {
|
||||
if (typeof window === "undefined") return null;
|
||||
const storage = reconnectOnMountRef.current;
|
||||
if (storage === true) return window.sessionStorage;
|
||||
if (typeof storage === "function") return storage();
|
||||
return null;
|
||||
}, []);
|
||||
|
||||
const client = useMemo(
|
||||
() =>
|
||||
options.client ??
|
||||
new Client({
|
||||
apiUrl: options.apiUrl,
|
||||
apiKey: options.apiKey,
|
||||
callerOptions: options.callerOptions,
|
||||
defaultHeaders: options.defaultHeaders,
|
||||
}),
|
||||
[
|
||||
options.client,
|
||||
options.apiKey,
|
||||
options.apiUrl,
|
||||
options.callerOptions,
|
||||
options.defaultHeaders,
|
||||
]
|
||||
);
|
||||
|
||||
const [messageManager] = useState(() => new MessageTupleManager());
|
||||
const [stream] = useState(
|
||||
() => new StreamManager<StateType, Bag>(messageManager)
|
||||
);
|
||||
|
||||
useSyncExternalStore(
|
||||
stream.subscribe,
|
||||
stream.getSnapshot,
|
||||
stream.getSnapshot
|
||||
);
|
||||
|
||||
const [threadId, onThreadId] = useControllableThreadId(options);
|
||||
const trackStreamModeRef = useRef<Exclude<StreamMode, "messages">[]>([]);
|
||||
|
||||
const trackStreamMode = useCallback(
|
||||
(...mode: Exclude<StreamMode, "messages">[]) => {
|
||||
const ref = trackStreamModeRef.current;
|
||||
for (const m of mode) {
|
||||
if (!ref.includes(m)) ref.push(m);
|
||||
}
|
||||
},
|
||||
[]
|
||||
);
|
||||
|
||||
const hasUpdateListener = options.onUpdateEvent != null;
|
||||
const hasCustomListener = options.onCustomEvent != null;
|
||||
const hasLangChainListener = options.onLangChainEvent != null;
|
||||
const hasDebugListener = options.onDebugEvent != null;
|
||||
const hasCheckpointListener = options.onCheckpointEvent != null;
|
||||
const hasTaskListener = options.onTaskEvent != null;
|
||||
|
||||
const callbackStreamMode = useMemo(() => {
|
||||
const modes: Exclude<StreamMode, "messages">[] = [];
|
||||
if (hasUpdateListener) modes.push("updates");
|
||||
if (hasCustomListener) modes.push("custom");
|
||||
if (hasLangChainListener) modes.push("events");
|
||||
if (hasDebugListener) modes.push("debug");
|
||||
if (hasCheckpointListener) modes.push("checkpoints");
|
||||
if (hasTaskListener) modes.push("tasks");
|
||||
return modes;
|
||||
}, [
|
||||
hasUpdateListener,
|
||||
hasCustomListener,
|
||||
hasLangChainListener,
|
||||
hasDebugListener,
|
||||
hasCheckpointListener,
|
||||
hasTaskListener,
|
||||
]);
|
||||
|
||||
const clearCallbackRef = useRef<() => void>(null!);
|
||||
clearCallbackRef.current = stream.clear;
|
||||
|
||||
const submittingRef = useRef(false);
|
||||
submittingRef.current = stream.isLoading;
|
||||
|
||||
const onErrorRef = useRef<
|
||||
((error: unknown, run?: RunCallbackMeta) => void) | undefined
|
||||
>(undefined);
|
||||
onErrorRef.current = options.onError;
|
||||
|
||||
const historyLimit =
|
||||
typeof options.fetchStateHistory === "object" &&
|
||||
options.fetchStateHistory != null
|
||||
? options.fetchStateHistory.limit ?? false
|
||||
: options.fetchStateHistory ?? false;
|
||||
|
||||
const history = useThreadHistory<StateType>(
|
||||
threadId,
|
||||
client,
|
||||
historyLimit,
|
||||
clearCallbackRef,
|
||||
submittingRef,
|
||||
onErrorRef
|
||||
);
|
||||
|
||||
const getMessages = (value: StateType): Message[] => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return Array.isArray(value[messagesKey]) ? value[messagesKey] : [];
|
||||
};
|
||||
|
||||
const setMessages = (current: StateType, messages: Message[]): StateType => {
|
||||
const messagesKey = options.messagesKey ?? "messages";
|
||||
return { ...current, [messagesKey]: messages };
|
||||
};
|
||||
|
||||
const [branch, setBranch] = useState<string>("");
|
||||
const branchContext = getBranchContext(branch, history.data);
|
||||
|
||||
const historyValues =
|
||||
branchContext.threadHead?.values ??
|
||||
options.initialValues ??
|
||||
({} as StateType);
|
||||
|
||||
const historyError = (() => {
|
||||
const error = branchContext.threadHead?.tasks?.at(-1)?.error;
|
||||
if (error == null) return undefined;
|
||||
try {
|
||||
const parsed = JSON.parse(error) as unknown;
|
||||
if (StreamError.isStructuredError(parsed)) return new StreamError(parsed);
|
||||
return parsed;
|
||||
} catch {
|
||||
// do nothing
|
||||
}
|
||||
return error;
|
||||
})();
|
||||
|
||||
const messageMetadata = (() => {
|
||||
const alreadyShown = new Set<string>();
|
||||
return getMessages(historyValues).map(
|
||||
(message, idx): Omit<MessageMetadata<StateType>, "streamMetadata"> => {
|
||||
const messageId = message.id ?? idx;
|
||||
|
||||
// Find the first checkpoint where the message was seen
|
||||
const firstSeenState = findLast(history.data ?? [], (state) =>
|
||||
getMessages(state.values)
|
||||
.map((m, idx) => m.id ?? idx)
|
||||
.includes(messageId)
|
||||
);
|
||||
|
||||
const checkpointId = firstSeenState?.checkpoint?.checkpoint_id;
|
||||
let branch =
|
||||
checkpointId != null
|
||||
? branchContext.branchByCheckpoint[checkpointId]
|
||||
: undefined;
|
||||
if (!branch?.branch?.length) branch = undefined;
|
||||
|
||||
// serialize branches
|
||||
const optionsShown = branch?.branchOptions?.flat(2).join(",");
|
||||
if (optionsShown) {
|
||||
if (alreadyShown.has(optionsShown)) branch = undefined;
|
||||
alreadyShown.add(optionsShown);
|
||||
}
|
||||
|
||||
return {
|
||||
messageId: messageId.toString(),
|
||||
firstSeenState,
|
||||
|
||||
branch: branch?.branch,
|
||||
branchOptions: branch?.branchOptions,
|
||||
};
|
||||
}
|
||||
);
|
||||
})();
|
||||
|
||||
const stop = () => stream.stop(historyValues, { onStop: options.onStop });
|
||||
|
||||
// --- TRANSPORT ---
|
||||
const submit = async (
|
||||
values: UpdateType | null | undefined,
|
||||
submitOptions?: SubmitOptions<StateType, ConfigurableType>
|
||||
) => {
|
||||
// Unbranch things
|
||||
const checkpointId = submitOptions?.checkpoint?.checkpoint_id;
|
||||
setBranch(
|
||||
checkpointId != null
|
||||
? branchContext.branchByCheckpoint[checkpointId]?.branch ?? ""
|
||||
: ""
|
||||
);
|
||||
|
||||
stream.setStreamValues(() => {
|
||||
if (submitOptions?.optimisticValues != null) {
|
||||
return {
|
||||
...historyValues,
|
||||
...(typeof submitOptions.optimisticValues === "function"
|
||||
? submitOptions.optimisticValues(historyValues)
|
||||
: submitOptions.optimisticValues),
|
||||
};
|
||||
}
|
||||
|
||||
return { ...historyValues };
|
||||
});
|
||||
|
||||
// When `fetchStateHistory` is requested, thus we assume that branching
|
||||
// is enabled. We then need to include the implicit branch.
|
||||
const includeImplicitBranch =
|
||||
historyLimit === true || typeof historyLimit === "number";
|
||||
|
||||
let callbackMeta: RunCallbackMeta | undefined;
|
||||
let rejoinKey: `lg:stream:${string}` | undefined;
|
||||
let usableThreadId = threadId;
|
||||
|
||||
await stream.start(
|
||||
async (signal: AbortSignal) => {
|
||||
if (!usableThreadId) {
|
||||
const thread = await client.threads.create({
|
||||
threadId: submitOptions?.threadId,
|
||||
metadata: submitOptions?.metadata,
|
||||
});
|
||||
onThreadId(thread.thread_id);
|
||||
usableThreadId = thread.thread_id;
|
||||
}
|
||||
|
||||
if (!usableThreadId) {
|
||||
throw new Error("Failed to obtain valid thread ID.");
|
||||
}
|
||||
|
||||
const streamMode = unique([
|
||||
...(submitOptions?.streamMode ?? []),
|
||||
...trackStreamModeRef.current,
|
||||
...callbackStreamMode,
|
||||
]);
|
||||
|
||||
let checkpoint =
|
||||
submitOptions?.checkpoint ??
|
||||
(includeImplicitBranch
|
||||
? branchContext.threadHead?.checkpoint
|
||||
: undefined) ??
|
||||
undefined;
|
||||
|
||||
// Avoid specifying a checkpoint if user explicitly set it to null
|
||||
if (submitOptions?.checkpoint === null) checkpoint = undefined;
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
|
||||
// @ts-expect-error
|
||||
if (checkpoint != null) delete checkpoint.thread_id;
|
||||
const streamResumable =
|
||||
submitOptions?.streamResumable ?? !!runMetadataStorage;
|
||||
|
||||
return client.runs.stream(usableThreadId, options.assistantId, {
|
||||
input: values as Record<string, unknown>,
|
||||
config: submitOptions?.config,
|
||||
context: submitOptions?.context,
|
||||
command: submitOptions?.command,
|
||||
|
||||
interruptBefore: submitOptions?.interruptBefore,
|
||||
interruptAfter: submitOptions?.interruptAfter,
|
||||
metadata: submitOptions?.metadata,
|
||||
multitaskStrategy: submitOptions?.multitaskStrategy,
|
||||
onCompletion: submitOptions?.onCompletion,
|
||||
onDisconnect:
|
||||
submitOptions?.onDisconnect ??
|
||||
(streamResumable ? "continue" : "cancel"),
|
||||
|
||||
signal,
|
||||
|
||||
checkpoint,
|
||||
streamMode,
|
||||
streamSubgraphs: submitOptions?.streamSubgraphs,
|
||||
streamResumable,
|
||||
durability: submitOptions?.durability,
|
||||
onRunCreated(params) {
|
||||
callbackMeta = {
|
||||
run_id: params.run_id,
|
||||
thread_id: params.thread_id ?? usableThreadId!,
|
||||
};
|
||||
|
||||
if (runMetadataStorage) {
|
||||
rejoinKey = `lg:stream:${usableThreadId}`;
|
||||
runMetadataStorage.setItem(rejoinKey, callbackMeta.run_id);
|
||||
}
|
||||
|
||||
options.onCreated?.(callbackMeta);
|
||||
},
|
||||
}) as AsyncGenerator<
|
||||
EventStreamEvent<StateType, UpdateType, CustomType>
|
||||
>;
|
||||
},
|
||||
{
|
||||
getMessages,
|
||||
setMessages,
|
||||
|
||||
initialValues: historyValues,
|
||||
callbacks: options,
|
||||
|
||||
async onSuccess() {
|
||||
if (rejoinKey) runMetadataStorage?.removeItem(rejoinKey);
|
||||
const shouldRefetch =
|
||||
// We're expecting the whole thread state in onFinish
|
||||
options.onFinish != null ||
|
||||
// We're fetching history, thus we need the latest checkpoint
|
||||
// to ensure we're not accidentally submitting to a wrong branch
|
||||
includeImplicitBranch;
|
||||
|
||||
if (shouldRefetch) {
|
||||
const newHistory = await history.mutate(usableThreadId!);
|
||||
const lastHead = newHistory.at(0);
|
||||
if (lastHead) {
|
||||
// We now have the latest update from /history
|
||||
// Thus we can clear the local stream state
|
||||
options.onFinish?.(lastHead, callbackMeta);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
},
|
||||
onError(error) {
|
||||
options.onError?.(error, callbackMeta);
|
||||
},
|
||||
}
|
||||
);
|
||||
};
|
||||
|
||||
const joinStream = async (
|
||||
runId: string,
|
||||
lastEventId?: string,
|
||||
joinOptions?: { streamMode?: StreamMode | StreamMode[] }
|
||||
) => {
|
||||
// eslint-disable-next-line no-param-reassign
|
||||
lastEventId ??= "-1";
|
||||
if (!threadId) return;
|
||||
|
||||
const callbackMeta: RunCallbackMeta = {
|
||||
thread_id: threadId,
|
||||
run_id: runId,
|
||||
};
|
||||
|
||||
await stream.start(
|
||||
async (signal: AbortSignal) => {
|
||||
return client.runs.joinStream(threadId, runId, {
|
||||
signal,
|
||||
lastEventId,
|
||||
streamMode: joinOptions?.streamMode,
|
||||
}) as AsyncGenerator<
|
||||
EventStreamEvent<StateType, UpdateType, CustomType>
|
||||
>;
|
||||
},
|
||||
{
|
||||
getMessages,
|
||||
setMessages,
|
||||
|
||||
initialValues: historyValues,
|
||||
callbacks: options,
|
||||
async onSuccess() {
|
||||
runMetadataStorage?.removeItem(`lg:stream:${threadId}`);
|
||||
const newHistory = await history.mutate(threadId);
|
||||
const lastHead = newHistory.at(0);
|
||||
if (lastHead) options.onFinish?.(lastHead, callbackMeta);
|
||||
},
|
||||
onError(error) {
|
||||
options.onError?.(error, callbackMeta);
|
||||
},
|
||||
}
|
||||
);
|
||||
};
|
||||
|
||||
const reconnectKey = useMemo(() => {
|
||||
if (!runMetadataStorage || stream.isLoading) return undefined;
|
||||
if (typeof window === "undefined") return undefined;
|
||||
const runId = runMetadataStorage?.getItem(`lg:stream:${threadId}`);
|
||||
if (!runId) return undefined;
|
||||
return { runId, threadId };
|
||||
}, [runMetadataStorage, stream.isLoading, threadId]);
|
||||
|
||||
const shouldReconnect = !!runMetadataStorage;
|
||||
const reconnectRef = useRef({ threadId, shouldReconnect });
|
||||
|
||||
const joinStreamRef = useRef<typeof joinStream>(joinStream);
|
||||
joinStreamRef.current = joinStream;
|
||||
|
||||
useEffect(() => {
|
||||
// reset shouldReconnect when switching threads
|
||||
if (reconnectRef.current.threadId !== threadId) {
|
||||
reconnectRef.current = { threadId, shouldReconnect };
|
||||
}
|
||||
}, [threadId, shouldReconnect]);
|
||||
|
||||
useEffect(() => {
|
||||
if (reconnectKey && reconnectRef.current.shouldReconnect) {
|
||||
reconnectRef.current.shouldReconnect = false;
|
||||
void joinStreamRef.current?.(reconnectKey.runId);
|
||||
}
|
||||
}, [reconnectKey]);
|
||||
// --- END TRANSPORT ---
|
||||
|
||||
const error = stream.error ?? historyError ?? history.error;
|
||||
const values = stream.values ?? historyValues;
|
||||
|
||||
return {
|
||||
get values() {
|
||||
trackStreamMode("values");
|
||||
return values;
|
||||
},
|
||||
|
||||
client,
|
||||
assistantId: options.assistantId,
|
||||
|
||||
error,
|
||||
isLoading: stream.isLoading,
|
||||
|
||||
stop,
|
||||
submit,
|
||||
|
||||
joinStream,
|
||||
|
||||
branch,
|
||||
setBranch,
|
||||
|
||||
get history() {
|
||||
if (historyLimit === false) {
|
||||
throw new Error(
|
||||
"`fetchStateHistory` must be set to `true` to use `history`"
|
||||
);
|
||||
}
|
||||
|
||||
return branchContext.flatHistory;
|
||||
},
|
||||
|
||||
isThreadLoading: history.isLoading && history.data == null,
|
||||
|
||||
get experimental_branchTree() {
|
||||
if (historyLimit === false) {
|
||||
throw new Error(
|
||||
"`fetchStateHistory` must be set to `true` to use `experimental_branchTree`"
|
||||
);
|
||||
}
|
||||
|
||||
return branchContext.branchTree;
|
||||
},
|
||||
|
||||
get interrupt() {
|
||||
if (
|
||||
values != null &&
|
||||
"__interrupt__" in values &&
|
||||
Array.isArray(values.__interrupt__)
|
||||
) {
|
||||
const valueInterrupts = values.__interrupt__;
|
||||
if (valueInterrupts.length === 0) return { when: "breakpoint" };
|
||||
if (valueInterrupts.length === 1) return valueInterrupts[0];
|
||||
|
||||
// TODO: fix the typing of interrupts if multiple interrupts are returned
|
||||
return valueInterrupts;
|
||||
}
|
||||
|
||||
// If we're deferring to old interrupt detection logic, don't show the interrupt if the stream is loading
|
||||
if (stream.isLoading) return undefined;
|
||||
|
||||
const interrupts = branchContext.threadHead?.tasks?.at(-1)?.interrupts;
|
||||
if (interrupts == null || interrupts.length === 0) {
|
||||
// check if there's a next task present
|
||||
const next = branchContext.threadHead?.next ?? [];
|
||||
if (!next.length || error != null) return undefined;
|
||||
return { when: "breakpoint" };
|
||||
}
|
||||
|
||||
// Return only the current interrupt
|
||||
return interrupts.at(-1) as Interrupt<InterruptType> | undefined;
|
||||
},
|
||||
|
||||
get messages() {
|
||||
trackStreamMode("messages-tuple", "values");
|
||||
return getMessages(values);
|
||||
},
|
||||
|
||||
getMessagesMetadata(
|
||||
message: Message,
|
||||
index?: number
|
||||
): MessageMetadata<StateType> | undefined {
|
||||
trackStreamMode("values");
|
||||
|
||||
const streamMetadata = messageManager.get(message.id)?.metadata;
|
||||
const historyMetadata = messageMetadata?.find(
|
||||
(m) => m.messageId === (message.id ?? index)
|
||||
);
|
||||
|
||||
if (streamMetadata != null || historyMetadata != null) {
|
||||
return {
|
||||
...historyMetadata,
|
||||
streamMetadata,
|
||||
} as MessageMetadata<StateType>;
|
||||
}
|
||||
|
||||
return undefined;
|
||||
},
|
||||
};
|
||||
>(
|
||||
options: UseStreamCustomOptions<StateType, Bag>
|
||||
): UseStreamCustom<StateType, Bag>;
|
||||
|
||||
export function useStream<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
ConfigurableType?: Record<string, unknown>;
|
||||
InterruptType?: unknown;
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(options: UseStreamOptions<StateType, Bag>): UseStream<StateType, Bag>;
|
||||
|
||||
export function useStream<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends {
|
||||
ConfigurableType?: Record<string, unknown>;
|
||||
InterruptType?: unknown;
|
||||
CustomEventType?: unknown;
|
||||
UpdateType?: unknown;
|
||||
} = BagTemplate
|
||||
>(
|
||||
options:
|
||||
| UseStreamOptions<StateType, Bag>
|
||||
| UseStreamCustomOptions<StateType, Bag>
|
||||
): UseStream<StateType, Bag> | UseStreamCustom<StateType, Bag> {
|
||||
// Store this in useState to make sure we're not changing the implementation in re-renders
|
||||
const [isCustom] = useState(isCustomOptions(options));
|
||||
|
||||
if (isCustom) {
|
||||
// eslint-disable-next-line react-hooks/rules-of-hooks
|
||||
return useStreamCustom(options as UseStreamCustomOptions<StateType, Bag>);
|
||||
}
|
||||
|
||||
// eslint-disable-next-line react-hooks/rules-of-hooks
|
||||
return useStreamLGP(options as UseStreamOptions<StateType, Bag>);
|
||||
}
|
||||
|
||||
@@ -436,3 +436,59 @@ export interface SubmitOptions<
|
||||
*/
|
||||
threadId?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Transport used to stream the thread.
|
||||
* Only applicable for custom endpoints using `toLangGraphEventStream` or `toLangGraphEventStreamResponse`.
|
||||
*/
|
||||
export interface UseStreamTransport<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends BagTemplate = BagTemplate
|
||||
> {
|
||||
stream: (payload: {
|
||||
input: GetUpdateType<Bag, StateType> | null | undefined;
|
||||
context: GetConfigurableType<Bag> | undefined;
|
||||
command: Command | undefined;
|
||||
signal: AbortSignal;
|
||||
}) => Promise<AsyncGenerator<{ id?: string; event: string; data: unknown }>>;
|
||||
}
|
||||
|
||||
export type UseStreamCustomOptions<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends BagTemplate = BagTemplate
|
||||
> = Pick<
|
||||
UseStreamOptions<StateType, Bag>,
|
||||
| "messagesKey"
|
||||
| "onError"
|
||||
| "onCreated"
|
||||
| "onUpdateEvent"
|
||||
| "onCustomEvent"
|
||||
| "onMetadataEvent"
|
||||
| "onLangChainEvent"
|
||||
| "onDebugEvent"
|
||||
| "onCheckpointEvent"
|
||||
| "onTaskEvent"
|
||||
| "onStop"
|
||||
| "initialValues"
|
||||
> & { transport: UseStreamTransport<StateType, Bag> };
|
||||
|
||||
export type UseStreamCustom<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
Bag extends BagTemplate = BagTemplate
|
||||
> = Pick<
|
||||
UseStream<StateType, Bag>,
|
||||
"values" | "error" | "isLoading" | "stop" | "interrupt" | "messages"
|
||||
> & {
|
||||
submit: (
|
||||
values: GetUpdateType<Bag, StateType> | null | undefined,
|
||||
options?: CustomSubmitOptions<StateType, GetConfigurableType<Bag>>
|
||||
) => Promise<void>;
|
||||
};
|
||||
|
||||
export type CustomSubmitOptions<
|
||||
StateType extends Record<string, unknown> = Record<string, unknown>,
|
||||
ConfigurableType extends Record<string, unknown> = Record<string, unknown>
|
||||
> = Pick<
|
||||
SubmitOptions<StateType, ConfigurableType>,
|
||||
"optimisticValues" | "context" | "command"
|
||||
>;
|
||||
|
||||
Reference in New Issue
Block a user