Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions packages/instantsearch.js/src/lib/ai-lite/abstract-chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import type {
FileUIPart,
IdGenerator,
InferUIMessageMetadata,
InferUIMessageToolCall,
InferUIMessageTools,
UIMessage,
UIMessageChunk,
Expand All @@ -21,9 +22,9 @@ import type {
ChatOnDataCallback,
} from './types';

type ActiveResponse = {
type ActiveResponse<TChunk extends UIMessageChunk = UIMessageChunk> = {
abortController: AbortController;
stream?: ReadableStream<UIMessageChunk>;
stream?: ReadableStream<TChunk>;
};

/**
Expand Down Expand Up @@ -685,7 +686,8 @@ export abstract class AbstractChat<TUIMessage extends UIMessage> {
toolName: chunk.toolName,
toolCallId: chunk.toolCallId,
input: chunk.input,
} as any,
dynamic: 'dynamic' in chunk ? chunk.dynamic : undefined,
} as InferUIMessageToolCall<TUIMessage>,
});
if (result && typeof result.then === 'function') {
pendingToolCall = pendingToolCall.then(() => result);
Expand Down
12 changes: 6 additions & 6 deletions packages/instantsearch.js/src/lib/ai-lite/stream-parser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,13 @@ import type { UIMessageChunk } from './types';
* @param stream - The input stream of raw bytes
* @returns A ReadableStream of parsed UIMessageChunk events
*/
export function parseJsonEventStream(
stream: ReadableStream<Uint8Array>
): ReadableStream<UIMessageChunk> {
export function parseJsonEventStream<
TChunk extends UIMessageChunk = UIMessageChunk
>(stream: ReadableStream<Uint8Array>): ReadableStream<TChunk> {
const decoder = new TextDecoder();
let buffer = '';

return new ReadableStream<UIMessageChunk>({
return new ReadableStream<TChunk>({
start(controller) {
const reader = stream.getReader();

Expand All @@ -30,7 +30,7 @@ export function parseJsonEventStream(
const jsonData = extractJsonFromLine(buffer.trim());
if (jsonData) {
try {
const chunk = JSON.parse(jsonData) as UIMessageChunk;
const chunk = JSON.parse(jsonData) as TChunk;
controller.enqueue(chunk);
} catch {
// Ignore parsing errors for incomplete data at end
Expand Down Expand Up @@ -60,7 +60,7 @@ export function parseJsonEventStream(
if (!jsonData) continue;

try {
const chunk = JSON.parse(jsonData) as UIMessageChunk;
const chunk = JSON.parse(jsonData) as TChunk;
controller.enqueue(chunk);
} catch {
// Skip malformed lines
Expand Down
12 changes: 6 additions & 6 deletions packages/instantsearch.js/src/lib/ai-lite/transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ import { resolveValue } from './utils';
import type {
ChatTransport,
HttpChatTransportInitOptions,
InferUIMessageChunk,
UIMessage,
UIMessageChunk,
FetchFunction,
PrepareSendMessagesRequest,
PrepareReconnectToStreamRequest,
Expand Down Expand Up @@ -57,7 +57,7 @@ export abstract class HttpChatTransport<TUIMessage extends UIMessage>
headers: requestHeaders,
body: requestBody,
}: Parameters<ChatTransport<TUIMessage>['sendMessages']>[0]): Promise<
ReadableStream<UIMessageChunk>
ReadableStream<InferUIMessageChunk<TUIMessage>>
> {
const fetchFn = this.fetch ?? fetch;

Expand Down Expand Up @@ -151,7 +151,7 @@ export abstract class HttpChatTransport<TUIMessage extends UIMessage>
body: requestBody,
}: Parameters<
ChatTransport<TUIMessage>['reconnectToStream']
>[0]): Promise<ReadableStream<UIMessageChunk> | null> {
>[0]): Promise<ReadableStream<InferUIMessageChunk<TUIMessage>> | null> {
const fetchFn = this.fetch ?? fetch;

// Resolve configurable values
Expand Down Expand Up @@ -230,7 +230,7 @@ export abstract class HttpChatTransport<TUIMessage extends UIMessage>

protected abstract processResponseStream(
stream: ReadableStream<Uint8Array>
): ReadableStream<UIMessageChunk>;
): ReadableStream<InferUIMessageChunk<TUIMessage>>;
}

/**
Expand All @@ -245,7 +245,7 @@ export class DefaultChatTransport<

protected processResponseStream(
stream: ReadableStream<Uint8Array>
): ReadableStream<UIMessageChunk> {
return parseJsonEventStream(stream);
): ReadableStream<InferUIMessageChunk<TUIMessage>> {
return parseJsonEventStream<InferUIMessageChunk<TUIMessage>>(stream);
}
}
104 changes: 76 additions & 28 deletions packages/instantsearch.js/src/lib/ai-lite/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -227,47 +227,68 @@ type DataUIMessageChunk<DATA_TYPES extends UIDataTypes> = ValueOf<{
};
}>;

export type UIMessageChunk<
METADATA = unknown,
DATA_TYPES extends UIDataTypes = UIDataTypes
> =
| { type: 'text-start'; id: string; providerMetadata?: ProviderMetadata }
| {
type: 'text-delta';
delta: string;
id: string;
providerMetadata?: ProviderMetadata;
}
| { type: 'text-end'; id: string; providerMetadata?: ProviderMetadata }
| { type: 'reasoning-start'; id: string; providerMetadata?: ProviderMetadata }
type ToolUIMessageChunk<TOOLS extends UITools> =
| ValueOf<{
[NAME in keyof TOOLS & string]: {
type: 'tool-input-available';
toolName: NAME;
toolCallId: string;
input: TOOLS[NAME]['input'];
callProviderMetadata?: ProviderMetadata;
providerExecuted?: boolean;
};
}>
| ValueOf<{
[NAME in keyof TOOLS & string]: {
type: 'tool-input-start';
toolName: NAME;
toolCallId: string;
input?: DeepPartial<TOOLS[NAME]['input']>;
providerExecuted?: boolean;
};
}>
| {
type: 'reasoning-delta';
id: string;
delta: string;
providerMetadata?: ProviderMetadata;
type: 'tool-input-delta';
toolName: string;
toolCallId: string;
inputDelta: string;
}
| { type: 'reasoning-end'; id: string; providerMetadata?: ProviderMetadata }
| { type: 'error'; errorText: string }
| ValueOf<{
[NAME in keyof TOOLS & string]: {
type: 'tool-output-available';
toolName: NAME;
toolCallId: string;
output: TOOLS[NAME]['output'];
callProviderMetadata?: ProviderMetadata;
preliminary?: boolean;
};
}>
| ValueOf<{
[NAME in keyof TOOLS & string]: {
type: 'tool-error';
toolName: NAME;
toolCallId: string;
errorText: string;
input?: TOOLS[NAME]['input'];
callProviderMetadata?: ProviderMetadata;
};
}>
| {
type: 'tool-input-available';
toolName: string;
toolCallId: string;
input: unknown;
callProviderMetadata?: ProviderMetadata;
providerExecuted?: boolean;
dynamic: true;
}
| {
type: 'tool-input-start';
toolName: string;
toolCallId: string;
input?: unknown;
providerExecuted?: boolean;
}
| {
type: 'tool-input-delta';
toolName: string;
toolCallId: string;
inputDelta: string;
dynamic: true;
}
| {
type: 'tool-output-available';
Expand All @@ -276,6 +297,7 @@ export type UIMessageChunk<
output: unknown;
callProviderMetadata?: ProviderMetadata;
preliminary?: boolean;
dynamic: true;
}
| {
type: 'tool-error';
Expand All @@ -284,7 +306,32 @@ export type UIMessageChunk<
errorText: string;
input?: unknown;
callProviderMetadata?: ProviderMetadata;
dynamic: true;
};

export type UIMessageChunk<
METADATA = unknown,
DATA_TYPES extends UIDataTypes = UIDataTypes,
TOOLS extends UITools = UITools
> =
| { type: 'text-start'; id: string; providerMetadata?: ProviderMetadata }
| {
type: 'text-delta';
delta: string;
id: string;
providerMetadata?: ProviderMetadata;
}
| { type: 'text-end'; id: string; providerMetadata?: ProviderMetadata }
| { type: 'reasoning-start'; id: string; providerMetadata?: ProviderMetadata }
| {
type: 'reasoning-delta';
id: string;
delta: string;
providerMetadata?: ProviderMetadata;
}
| { type: 'reasoning-end'; id: string; providerMetadata?: ProviderMetadata }
| { type: 'error'; errorText: string }
| ToolUIMessageChunk<TOOLS>
| { type: 'source-url'; sourceId: string; url: string; title?: string }
| {
type: 'source-document';
Expand All @@ -305,7 +352,8 @@ export type UIMessageChunk<

export type InferUIMessageChunk<T extends UIMessage> = UIMessageChunk<
InferUIMessageMetadata<T>,
InferUIMessageData<T>
InferUIMessageData<T>,
InferUIMessageTools<T>
>;

export interface ChatState<UI_MESSAGE extends UIMessage> {
Expand Down Expand Up @@ -334,13 +382,13 @@ export interface ChatTransport<UI_MESSAGE extends UIMessage> {
trigger: 'submit-message' | 'regenerate-message';
messageId?: string;
} & ChatRequestOptions
) => Promise<ReadableStream<UIMessageChunk>>;
) => Promise<ReadableStream<InferUIMessageChunk<UI_MESSAGE>>>;

reconnectToStream: (
options: {
chatId: string;
} & ChatRequestOptions
) => Promise<ReadableStream<UIMessageChunk> | null>;
) => Promise<ReadableStream<InferUIMessageChunk<UI_MESSAGE>> | null>;
}

export type PrepareSendMessagesRequest<UI_MESSAGE extends UIMessage> = (
Expand Down
Loading