feat(sdk-js): strongly-typed messages, state/update-type, stream mode

This commit is contained in:
Tat Dat Duong
2025-02-08 20:05:31 -08:00
parent ce80114cf9
commit da3a1a862d
5 changed files with 361 additions and 58 deletions
+69 -37
View File
@@ -27,7 +27,7 @@ import {
createParser,
} from "./utils/eventsource-parser/index.js";
import { IterableReadableStream } from "./utils/stream.js";
import {
import type {
RunsCreatePayload,
RunsStreamPayload,
RunsWaitPayload,
@@ -38,6 +38,7 @@ import {
import { mergeSignals } from "./utils/signals.js";
import { getEnvironmentVariable } from "./utils/env.js";
import { _getFetchImplementation } from "./singletons/fetch.js";
import type { TypedAsyncGenerator, StreamMode } from "./types.stream.js";
/**
* Get the API key from the environment.
* Precedence:
@@ -457,15 +458,17 @@ export class AssistantsClient extends BaseClient {
}
}
export class ThreadsClient extends BaseClient {
export class ThreadsClient<TStateType = DefaultValues> extends BaseClient {
/**
* Get a thread by ID.
*
* @param threadId ID of the thread.
* @returns The thread.
*/
async get(threadId: string): Promise<Thread> {
return this.fetch<Thread>(`/threads/${threadId}`);
async get<ValuesType = TStateType>(
threadId: string,
): Promise<Thread<ValuesType>> {
return this.fetch<Thread<ValuesType>>(`/threads/${threadId}`);
}
/**
@@ -481,8 +484,8 @@ export class ThreadsClient extends BaseClient {
metadata?: Metadata;
threadId?: string;
ifExists?: OnConflictBehavior;
}): Promise<Thread> {
return this.fetch<Thread>(`/threads`, {
}): Promise<Thread<TStateType>> {
return this.fetch<Thread<TStateType>>(`/threads`, {
method: "POST",
json: {
metadata: payload?.metadata,
@@ -497,8 +500,8 @@ export class ThreadsClient extends BaseClient {
* @param threadId ID of the thread to be copied
* @returns Newly copied thread
*/
async copy(threadId: string): Promise<Thread> {
return this.fetch<Thread>(`/threads/${threadId}/copy`, {
async copy(threadId: string): Promise<Thread<TStateType>> {
return this.fetch<Thread<TStateType>>(`/threads/${threadId}/copy`, {
method: "POST",
});
}
@@ -542,7 +545,7 @@ export class ThreadsClient extends BaseClient {
* @param query Query options
* @returns List of threads
*/
async search(query?: {
async search<ValuesType = TStateType>(query?: {
/**
* Metadata to filter threads by.
*/
@@ -561,8 +564,8 @@ export class ThreadsClient extends BaseClient {
* Must be one of 'idle', 'busy', 'interrupted' or 'error'.
*/
status?: ThreadStatus;
}): Promise<Thread[]> {
return this.fetch<Thread[]>("/threads/search", {
}): Promise<Thread<ValuesType>[]> {
return this.fetch<Thread<ValuesType>[]>("/threads/search", {
method: "POST",
json: {
metadata: query?.metadata ?? undefined,
@@ -579,7 +582,7 @@ export class ThreadsClient extends BaseClient {
* @param threadId ID of the thread.
* @returns Thread state.
*/
async getState<ValuesType = DefaultValues>(
async getState<ValuesType = TStateType>(
threadId: string,
checkpoint?: Checkpoint | string,
options?: { subgraphs?: boolean },
@@ -613,7 +616,7 @@ export class ThreadsClient extends BaseClient {
* @param threadId The ID of the thread.
* @returns
*/
async updateState<ValuesType = DefaultValues>(
async updateState<ValuesType = TStateType>(
threadId: string,
options: {
values: ValuesType;
@@ -672,7 +675,7 @@ export class ThreadsClient extends BaseClient {
* @param options Additional options.
* @returns List of thread states.
*/
async getHistory<ValuesType = DefaultValues>(
async getHistory<ValuesType = TStateType>(
threadId: string,
options?: {
limit?: number;
@@ -696,24 +699,43 @@ export class ThreadsClient extends BaseClient {
}
}
export class RunsClient extends BaseClient {
stream(
export class RunsClient<
TStateType = DefaultValues,
TUpdateType = TStateType,
TCustomEventType = unknown,
> extends BaseClient {
stream<
TStreamMode extends StreamMode | StreamMode[] = [],
TSubgraphs extends boolean = false,
>(
threadId: null,
assistantId: string,
payload?: Omit<RunsStreamPayload, "multitaskStrategy" | "onCompletion">,
): AsyncGenerator<{
event: StreamEvent;
data: any;
}>;
payload?: Omit<
RunsStreamPayload<TStreamMode, TSubgraphs>,
"multitaskStrategy" | "onCompletion"
>,
): TypedAsyncGenerator<
TStreamMode,
TSubgraphs,
TStateType,
TUpdateType,
TCustomEventType
>;
stream(
stream<
TStreamMode extends StreamMode | StreamMode[] = [],
TSubgraphs extends boolean = false,
>(
threadId: string,
assistantId: string,
payload?: RunsStreamPayload,
): AsyncGenerator<{
event: StreamEvent;
data: any;
}>;
payload?: RunsStreamPayload<TStreamMode, TSubgraphs>,
): TypedAsyncGenerator<
TStreamMode,
TSubgraphs,
TStateType,
TUpdateType,
TCustomEventType
>;
/**
* Create a run and stream the results.
@@ -722,14 +744,20 @@ export class RunsClient extends BaseClient {
* @param assistantId Assistant ID to use for this run.
* @param payload Payload for creating a run.
*/
async *stream(
async *stream<
TStreamMode extends StreamMode | StreamMode[] = [],
TSubgraphs extends boolean = false,
>(
threadId: string | null,
assistantId: string,
payload?: RunsStreamPayload,
): AsyncGenerator<{
event: StreamEvent;
data: any;
}> {
payload?: RunsStreamPayload<TStreamMode, TSubgraphs>,
): TypedAsyncGenerator<
TStreamMode,
TSubgraphs,
TStateType,
TUpdateType,
TCustomEventType
> {
const json: Record<string, any> = {
input: payload?.input,
command: payload?.command,
@@ -767,7 +795,7 @@ export class RunsClient extends BaseClient {
let onEndEvent: () => void;
const textDecoder = new TextDecoder();
const stream: ReadableStream<{ event: string; data: any }> = (
const stream: ReadableStream<{ event: any; data: any }> = (
response.body || new ReadableStream({ start: (ctrl) => ctrl.close() })
).pipeThrough(
new TransformStream({
@@ -1284,7 +1312,11 @@ export class StoreClient extends BaseClient {
}
}
export class Client {
export class Client<
TStateType = DefaultValues,
TUpdateType = TStateType,
TCustomEventType = unknown,
> {
/**
* The client for interacting with assistants.
*/
@@ -1293,12 +1325,12 @@ export class Client {
/**
* The client for interacting with threads.
*/
public threads: ThreadsClient;
public threads: ThreadsClient<TStateType>;
/**
* The client for interacting with runs.
*/
public runs: RunsClient;
public runs: RunsClient<TStateType, TUpdateType, TCustomEventType>;
/**
* The client for interacting with cron runs.
+21
View File
@@ -26,3 +26,24 @@ export type {
export { overrideFetchImplementation } from "./singletons/fetch.js";
export type { OnConflictBehavior, Command } from "./types.js";
export type { StreamMode } from "./types.stream.js";
export type {
ValuesPayload,
MessagesTuplePayload,
MetadataPayload,
UpdatesPayload,
CustomPayload,
MessagesPayload,
DebugPayload,
EventsPayload,
} from "./types.stream.js";
export type {
Message,
HumanMessage,
AIMessage,
ToolMessage,
SystemMessage,
FunctionMessage,
RemoveMessage,
} from "./types.messages.js";
+108
View File
@@ -0,0 +1,108 @@
type ImageDetail = "auto" | "low" | "high";
type MessageContentImageUrl = {
type: "image_url";
image_url: string | { url: string; detail?: ImageDetail | undefined };
};
type MessageContentText = { type: "text"; text: string };
type MessageContentComplex = MessageContentText | MessageContentImageUrl;
type MessageAdditionalKwargs = {
[x: string]: unknown;
function_call?: { arguments: string; name: string } | undefined;
tool_calls?:
| {
id: string;
function: { arguments: string; name: string };
type: "function";
index?: number | undefined;
}[]
| undefined;
};
export type HumanMessage = {
type: "human";
id?: string | undefined;
content: string | MessageContentComplex[];
};
export type AIMessage = {
type: "ai";
id?: string | undefined;
content: string | MessageContentComplex[];
tool_calls?:
| {
name: string;
args: { [x: string]: { [x: string]: any } };
id?: string | undefined;
type?: "tool_call" | undefined;
}[]
| undefined;
invalid_tool_calls?:
| {
name?: string | undefined;
args?: string | undefined;
id?: string | undefined;
error?: string | undefined;
type?: "invalid_tool_call" | undefined;
}[]
| undefined;
usage_metadata?:
| {
input_tokens: number;
output_tokens: number;
total_tokens: number;
input_token_details?:
| {
audio?: number | undefined;
cache_read?: number | undefined;
cache_creation?: number | undefined;
}
| undefined;
output_token_details?:
| { audio?: number | undefined; reasoning?: number | undefined }
| undefined;
}
| undefined;
additional_kwargs?: MessageAdditionalKwargs | undefined;
response_metadata?: Record<string, unknown> | undefined;
};
export type ToolMessage = {
type: "tool";
name?: string | undefined;
id?: string | undefined;
content: string | MessageContentComplex[];
status?: "error" | "success" | undefined;
lc_direct_tool_output: boolean;
tool_call_id: string;
additional_kwargs?: MessageAdditionalKwargs | undefined;
response_metadata?: Record<string, unknown> | undefined;
};
export type SystemMessage = {
type: "system";
id?: string | undefined;
content: string | MessageContentComplex[];
};
export type FunctionMessage = {
type: "function";
id?: string | undefined;
content: string | MessageContentComplex[];
};
export type RemoveMessage = {
type: "remove";
id: string;
content: string | MessageContentComplex[];
};
export type Message =
| HumanMessage
| AIMessage
| ToolMessage
| SystemMessage
| FunctionMessage
| RemoveMessage;
+155
View File
@@ -0,0 +1,155 @@
import type { Message } from "./types.messages.js";
/**
* Stream modes
* - "values": Stream only the state values.
* - "messages": Stream complete messages.
* - "messages-tuple": Stream (message chunk, metadata) tuples.
* - "updates": Stream updates to the state.
* - "events": Stream events occurring during execution.
* - "debug": Stream detailed debug information.
* - "custom": Stream custom events.
*/
export type StreamMode =
| "values"
| "messages"
| "updates"
| "events"
| "debug"
| "custom"
| "messages-tuple";
type MessageTupleMetadata = {
tags: string[];
[key: string]: unknown;
};
type AsSubgraph<TEvent extends { event: string; data: unknown }> = {
event: TEvent["event"] | `${TEvent["event"]}|${string}`;
data: TEvent["data"];
};
export type ValuesPayload<StateType> = { event: "values"; data: StateType };
/** @internal */
export type ValuesPayloadSubgraphs<StateType> = AsSubgraph<
ValuesPayload<StateType>
>;
export type MessagesTuplePayload = {
event: "messages";
// TODO: add types for message and config, which do not depend on LangChain
// while making sure it's easy to keep them in sync.
data: [message: Message, config: MessageTupleMetadata];
};
/** @internal */
export type MessagesTuplePayloadSubgraphs = AsSubgraph<MessagesTuplePayload>;
export type MetadataPayload = {
event: "metadata";
data: { run_id: string; thread_id: string };
};
/** @internal */
export type MetadataPayloadSubgraphs = AsSubgraph<MetadataPayload>;
export type UpdatesPayload<UpdateType> = {
event: "updates";
data: { [node: string]: UpdateType };
};
/** @internal */
export type UpdatesPayloadSubgraphs<UpdateType> = AsSubgraph<
UpdatesPayload<UpdateType>
>;
export type CustomPayload<T> = { event: "custom"; data: T };
/** @internal */
export type CustomPayloadSubgraphs<T> = AsSubgraph<CustomPayload<T>>;
type MessagesMetadataPayload = {
event: "messages/metadata";
data: { [messageId: string]: { metadata: unknown } };
};
type MessagesCompletePayload = {
event: "messages/complete";
data: Message[];
};
type MessagesPartialPayload = {
event: "messages/partial";
data: Message[];
};
export type MessagesPayload =
| MessagesMetadataPayload
| MessagesCompletePayload
| MessagesPartialPayload;
/** @internal */
export type MessagesPayloadSubgraphs =
| AsSubgraph<MessagesMetadataPayload>
| AsSubgraph<MessagesCompletePayload>
| AsSubgraph<MessagesPartialPayload>;
export type DebugPayload = { event: "debug"; data: unknown };
/** @internal */
export type DebugPayloadSubgraphs = AsSubgraph<DebugPayload>;
export type EventsPayload = { event: "events"; data: unknown };
/** @internal */
export type EventsPayloadSubgraphs = AsSubgraph<EventsPayload>;
type GetStreamModeMap<
TStreamMode extends StreamMode | StreamMode[],
TStateType = unknown,
TUpdateType = TStateType,
TCustomType = unknown,
> =
| {
values: ValuesPayload<TStateType>;
updates: UpdatesPayload<TUpdateType>;
custom: CustomPayload<TCustomType>;
debug: DebugPayload;
messages: MessagesPayload;
"messages-tuple": MessagesTuplePayload;
events: EventsPayload;
}[TStreamMode extends StreamMode[] ? TStreamMode[number] : TStreamMode]
| MetadataPayload;
type GetSubgraphsStreamModeMap<
TStreamMode extends StreamMode | StreamMode[],
TStateType = unknown,
TUpdateType = TStateType,
TCustomType = unknown,
> =
| {
values: ValuesPayloadSubgraphs<TStateType>;
updates: UpdatesPayloadSubgraphs<TUpdateType>;
custom: CustomPayloadSubgraphs<TCustomType>;
debug: DebugPayloadSubgraphs;
messages: MessagesPayloadSubgraphs;
"messages-tuple": MessagesTuplePayloadSubgraphs;
events: EventsPayloadSubgraphs;
}[TStreamMode extends StreamMode[] ? TStreamMode[number] : TStreamMode]
| MetadataPayloadSubgraphs;
export type TypedAsyncGenerator<
TStreamMode extends StreamMode | StreamMode[] = [],
TSubgraphs extends boolean = false,
TStateType = unknown,
TUpdateType = TStateType,
TCustomType = unknown,
> = AsyncGenerator<
TSubgraphs extends true
? GetSubgraphsStreamModeMap<
TStreamMode,
TStateType,
TUpdateType,
TCustomType
>
: GetStreamModeMap<TStreamMode, TStateType, TUpdateType, TCustomType>
>;
+8 -21
View File
@@ -1,23 +1,6 @@
import { Checkpoint, Config, Metadata } from "./schema.js";
import { StreamMode } from "./types.stream.js";
/**
* Stream modes
* - "values": Stream only the state values.
* - "messages": Stream complete messages.
* - "messages-tuple": Stream (message chunk, metadata) tuples.
* - "updates": Stream updates to the state.
* - "events": Stream events occurring during execution.
* - "debug": Stream detailed debug information.
* - "custom": Stream custom events.
*/
export type StreamMode =
| "values"
| "messages"
| "updates"
| "events"
| "debug"
| "custom"
| "messages-tuple";
export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
export type OnConflictBehavior = "raise" | "do_nothing";
export type OnCompletionBehavior = "complete" | "continue";
@@ -31,6 +14,7 @@ export type StreamEvent =
| "messages/partial"
| "messages/metadata"
| "messages/complete"
| "messages"
| (string & {});
export interface Send {
@@ -148,16 +132,19 @@ interface RunsInvokePayload {
command?: Command;
}
export interface RunsStreamPayload extends RunsInvokePayload {
export interface RunsStreamPayload<
TStreamMode extends StreamMode | StreamMode[] = [],
TSubgraphs extends boolean = false,
> extends RunsInvokePayload {
/**
* One of `"values"`, `"messages"`, `"messages-tuple"`, `"updates"`, `"events"`, `"debug"`, `"custom"`.
*/
streamMode?: StreamMode | Array<StreamMode>;
streamMode?: TStreamMode;
/**
* Stream output from subgraphs. By default, streams only the top graph.
*/
streamSubgraphs?: boolean;
streamSubgraphs?: TSubgraphs;
/**
* Pass one or more feedbackKeys if you want to request short-lived signed URLs