fix(askai): Fix new conversation causing thread depth errors (#2900)
This commit is contained in:
Paul Jankowski
GitHub
parent
ef2bec0e20
commit
dfc1048417
8 changed files with 730 additions and 345 deletions
|
|
@ -77,7 +77,7 @@ bun run pw:run:webkit
|
|||
|
||||
### Imports
|
||||
|
||||
Imports must be ordered alphabetically with newlines between groups:
|
||||
Imported modules must be ordered alphabetically with newlines between groups:
|
||||
|
||||
1. Built-in modules
|
||||
2. External dependencies
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import type {
|
|||
InternalDocSearchHit,
|
||||
OnAskAiFeedback,
|
||||
StoredAskAiMessage,
|
||||
StoredAskAiState,
|
||||
SuggestedQuestionHit,
|
||||
} from './types';
|
||||
import { type AskAiState } from './types/AskiAi';
|
||||
|
|
@ -172,17 +173,28 @@ export function DocSearchAskAiModal({
|
|||
|
||||
const [stoppedStream, setStoppedStream] = React.useState(false);
|
||||
|
||||
const { messages, status, setMessages, sendMessage, stopAskAiStreaming, askAiError, sendFeedback, conversations } =
|
||||
useAskAi({
|
||||
assistantId: askAiConfigurationId,
|
||||
apiKey: askAiConfig?.apiKey || apiKey,
|
||||
appId: askAiConfig?.appId || appId,
|
||||
indexName: askAiConfig?.indexName || defaultIndexName,
|
||||
searchParameters: askAiSearchParameters,
|
||||
tools,
|
||||
memory: props.memory,
|
||||
indices: askAiConfig?.indices,
|
||||
});
|
||||
const {
|
||||
chatId,
|
||||
messages,
|
||||
status,
|
||||
setMessages,
|
||||
sendMessage,
|
||||
stopAskAiStreaming,
|
||||
askAiError,
|
||||
sendFeedback,
|
||||
conversations,
|
||||
startNewConversation,
|
||||
restoreConversation,
|
||||
} = useAskAi({
|
||||
assistantId: askAiConfigurationId,
|
||||
apiKey: askAiConfig?.apiKey || apiKey,
|
||||
appId: askAiConfig?.appId || appId,
|
||||
indexName: askAiConfig?.indexName || defaultIndexName,
|
||||
searchParameters: askAiSearchParameters,
|
||||
tools,
|
||||
memory: props.memory,
|
||||
indices: askAiConfig?.indices,
|
||||
});
|
||||
|
||||
const prevStatus = React.useRef(status);
|
||||
React.useEffect(() => {
|
||||
|
|
@ -200,12 +212,12 @@ export function DocSearchAskAiModal({
|
|||
|
||||
for (const part of messages[0].parts) {
|
||||
if (part.type === 'text') {
|
||||
conversations.add(buildDummyAskAiHit(part.text, messages));
|
||||
conversations.add(buildDummyAskAiHit(part.text, messages, chatId));
|
||||
}
|
||||
}
|
||||
}
|
||||
prevStatus.current = status;
|
||||
}, [status, messages, conversations, disableUserPersonalization, stoppedStream]);
|
||||
}, [status, messages, conversations, disableUserPersonalization, stoppedStream, chatId]);
|
||||
|
||||
// Check if there's a thread depth error (AI-217)
|
||||
const hasThreadDepthError = React.useMemo(() => {
|
||||
|
|
@ -392,14 +404,19 @@ export function DocSearchAskAiModal({
|
|||
|
||||
useRefreshOnInitialQuery({ initialQuery, inputRef, refresh });
|
||||
|
||||
const hasCurrentMessages = messages.length > 0;
|
||||
|
||||
// Refresh the autocomplete results when ask ai is toggled off
|
||||
// helps return to the previous ac state and start screen
|
||||
React.useEffect(() => {
|
||||
if (!isAskAiActive) {
|
||||
autocomplete.refresh();
|
||||
setMessages([]);
|
||||
|
||||
if (hasCurrentMessages) {
|
||||
startNewConversation();
|
||||
}
|
||||
}
|
||||
}, [isAskAiActive, autocomplete, setMessages]);
|
||||
}, [isAskAiActive, autocomplete, startNewConversation, hasCurrentMessages]);
|
||||
|
||||
// Track external state in order to manage internal askAiState
|
||||
React.useEffect(() => {
|
||||
|
|
@ -413,7 +430,7 @@ export function DocSearchAskAiModal({
|
|||
};
|
||||
|
||||
const handleNewConversation = (): void => {
|
||||
setMessages([]);
|
||||
startNewConversation();
|
||||
setAskAiState('new-conversation');
|
||||
};
|
||||
|
||||
|
|
@ -510,10 +527,11 @@ export function DocSearchAskAiModal({
|
|||
onItemClick={(item, event) => {
|
||||
if (item.type === 'askAI' && item.query) {
|
||||
if (item.anchor === 'stored' && 'messages' in item) {
|
||||
setMessages(item.messages as any);
|
||||
const hitMessages = item.messages as StoredAskAiMessage[];
|
||||
restoreConversation(hitMessages, (item as StoredAskAiState).chatId);
|
||||
const initialMessage: InitialAskAiMessage = {
|
||||
query: item.query,
|
||||
messageId: (item.messages as StoredAskAiMessage[])[0].id,
|
||||
messageId: hitMessages[0].id,
|
||||
};
|
||||
|
||||
if (interceptAskAiEvent?.(initialMessage)) {
|
||||
|
|
|
|||
|
|
@ -176,16 +176,18 @@ function SidepanelInner(
|
|||
const searchClient = useSearchClient(appId, apiKey, setSidepanelSearchClient);
|
||||
|
||||
const {
|
||||
chatId,
|
||||
status,
|
||||
sendMessage,
|
||||
stopAskAiStreaming,
|
||||
isStreaming,
|
||||
exchanges,
|
||||
setMessages,
|
||||
conversations,
|
||||
messages,
|
||||
sendFeedback,
|
||||
askAiError,
|
||||
startNewConversation,
|
||||
restoreConversation,
|
||||
} = useAskAi({
|
||||
appId,
|
||||
indexName,
|
||||
|
|
@ -213,13 +215,13 @@ function SidepanelInner(
|
|||
};
|
||||
|
||||
const handleStartNewConversation = (): void => {
|
||||
setMessages([]);
|
||||
startNewConversation();
|
||||
setSidepanelState('new-conversation');
|
||||
};
|
||||
|
||||
const handleSelectQuestion = (question: SuggestedQuestionHit): void => {
|
||||
setStoppedStreaming(false);
|
||||
setMessages([]);
|
||||
startNewConversation();
|
||||
sendMessage(
|
||||
{ text: question.question },
|
||||
{
|
||||
|
|
@ -239,14 +241,14 @@ function SidepanelInner(
|
|||
const handleSelectConversation = React.useCallback(
|
||||
(conversation: StoredAskAiState): void => {
|
||||
if (conversation.messages) {
|
||||
setMessages(conversation.messages);
|
||||
restoreConversation(conversation.messages, conversation.chatId);
|
||||
} else if (conversation.query) {
|
||||
sendMessage({ text: conversation.query });
|
||||
}
|
||||
|
||||
setSidepanelState('conversation');
|
||||
},
|
||||
[sendMessage, setMessages],
|
||||
[sendMessage, restoreConversation],
|
||||
);
|
||||
|
||||
useManageSidepanelLayout({
|
||||
|
|
@ -290,13 +292,13 @@ function SidepanelInner(
|
|||
|
||||
for (const part of messages[0].parts) {
|
||||
if (part.type === 'text') {
|
||||
conversations.add(buildDummyAskAiHit(part.text, messages));
|
||||
conversations.add(buildDummyAskAiHit(part.text, messages, chatId));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
prevStatus.current = status;
|
||||
}, [conversations, status, messages, stoppedStreaming]);
|
||||
}, [conversations, status, messages, stoppedStreaming, chatId]);
|
||||
|
||||
React.useEffect(() => {
|
||||
function setFullViewportHeight(): void {
|
||||
|
|
@ -327,7 +329,7 @@ function SidepanelInner(
|
|||
if (selectedConversation) {
|
||||
handleSelectConversation(selectedConversation);
|
||||
} else {
|
||||
setMessages([]);
|
||||
startNewConversation();
|
||||
sendMessage(
|
||||
{
|
||||
text: initialMessage.query,
|
||||
|
|
@ -342,7 +344,7 @@ function SidepanelInner(
|
|||
);
|
||||
setSidepanelState('conversation');
|
||||
}
|
||||
}, [initialMessage, sendMessage, conversations, handleSelectConversation, setMessages]);
|
||||
}, [initialMessage, sendMessage, conversations, handleSelectConversation, startNewConversation]);
|
||||
|
||||
// Autofocus the prompt input when the sidepanel opens and blur it when
|
||||
// it closes. Disabled on mobile because focusing the textarea triggers the
|
||||
|
|
|
|||
609
packages/docsearch-react/src/__tests__/useAskAi.test.ts
Normal file
609
packages/docsearch-react/src/__tests__/useAskAi.test.ts
Normal file
|
|
@ -0,0 +1,609 @@
|
|||
/* eslint-disable max-classes-per-file */
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { useAskAi } from '../useAskAi';
|
||||
|
||||
interface ToolCall {
|
||||
input: unknown;
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
}
|
||||
|
||||
interface ChatMessage {
|
||||
id: string;
|
||||
role: 'assistant' | 'user';
|
||||
parts: Array<{ type: 'text'; text: string }>;
|
||||
}
|
||||
|
||||
interface ChatOptions {
|
||||
id?: string;
|
||||
messages?: ChatMessage[];
|
||||
onToolCall: (params: { toolCall: ToolCall }) => unknown;
|
||||
transport: {
|
||||
options: {
|
||||
headers?: Record<string, string>;
|
||||
body?: Record<string, unknown>;
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
interface ChatInstance {
|
||||
id?: string;
|
||||
messages?: ChatMessage[];
|
||||
options: ChatOptions;
|
||||
sendMessage: (message: unknown) => void;
|
||||
}
|
||||
|
||||
interface UseChatOptions {
|
||||
chat: ChatInstance;
|
||||
}
|
||||
|
||||
type CustomOnToolCallParams = ToolCall & {
|
||||
addToolOutput: (props: { output: unknown }) => Promise<void>;
|
||||
};
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
addToolOutput: vi.fn(),
|
||||
setMessages: vi.fn(),
|
||||
sendMessage: vi.fn(),
|
||||
useChat: vi.fn(),
|
||||
generateId: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('@ai-sdk/react', () => ({
|
||||
Chat: class Chat {
|
||||
id?: string;
|
||||
messages?: ChatMessage[];
|
||||
options: ChatOptions;
|
||||
|
||||
constructor(options: ChatOptions) {
|
||||
this.id = options.id;
|
||||
this.messages = options.messages;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
sendMessage(message: unknown): void {
|
||||
mocks.sendMessage({ chatId: this.id, message });
|
||||
}
|
||||
},
|
||||
useChat: mocks.useChat,
|
||||
}));
|
||||
|
||||
vi.mock('ai', () => ({
|
||||
DefaultChatTransport: class DefaultChatTransport {
|
||||
options: unknown;
|
||||
|
||||
constructor(options: unknown) {
|
||||
this.options = options;
|
||||
}
|
||||
},
|
||||
lastAssistantMessageIsCompleteWithToolCalls: vi.fn(() => false),
|
||||
generateId: mocks.generateId,
|
||||
}));
|
||||
|
||||
interface SendMessageCall {
|
||||
chatId?: string;
|
||||
message: unknown;
|
||||
}
|
||||
|
||||
describe('useAskAi', () => {
|
||||
let chatOptions: ChatOptions | undefined;
|
||||
let chatOptionsHistory: ChatOptions[] = [];
|
||||
let sendMessageCalls: SendMessageCall[] = [];
|
||||
|
||||
function getOnToolCall(): ChatOptions['onToolCall'] {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.onToolCall;
|
||||
}
|
||||
|
||||
function getTransportHeaders(): Record<string, string> {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.transport.options.headers ?? {};
|
||||
}
|
||||
|
||||
function getTransportBody(): Record<string, unknown> {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.transport.options.body ?? {};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
let idCounter = 0;
|
||||
mocks.generateId.mockImplementation(() => {
|
||||
idCounter += 1;
|
||||
return `generated-id-${idCounter}`;
|
||||
});
|
||||
|
||||
chatOptions = undefined;
|
||||
chatOptionsHistory = [];
|
||||
sendMessageCalls = [];
|
||||
mocks.sendMessage.mockImplementation((call: SendMessageCall) => {
|
||||
sendMessageCalls.push(call);
|
||||
});
|
||||
mocks.useChat.mockImplementation(({ chat }: UseChatOptions) => {
|
||||
chatOptions = chat.options;
|
||||
chatOptionsHistory.push(chat.options);
|
||||
|
||||
return {
|
||||
addToolOutput: mocks.addToolOutput,
|
||||
error: undefined,
|
||||
messages: chat.messages ?? [],
|
||||
sendMessage: chat.sendMessage,
|
||||
setMessages: mocks.setMessages,
|
||||
status: 'ready',
|
||||
stop: vi.fn(),
|
||||
};
|
||||
});
|
||||
});
|
||||
|
||||
it('forwards custom tool output to useChat addToolOutput', async () => {
|
||||
let addToolOutput: CustomOnToolCallParams['addToolOutput'] | undefined;
|
||||
const onToolCall = vi.fn((params: CustomOnToolCallParams) => {
|
||||
addToolOutput = params.addToolOutput;
|
||||
});
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {
|
||||
customAction: {
|
||||
onToolCall,
|
||||
render: () => 'Custom action complete',
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
getOnToolCall()({
|
||||
toolCall: {
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
},
|
||||
});
|
||||
|
||||
expect(onToolCall).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
}),
|
||||
);
|
||||
|
||||
if (!addToolOutput) {
|
||||
throw new Error('addToolOutput was not provided to the custom tool');
|
||||
}
|
||||
|
||||
await addToolOutput({ output: { result: 'output value' } });
|
||||
|
||||
expect(mocks.addToolOutput).toHaveBeenCalledWith({
|
||||
output: { result: 'output value' },
|
||||
tool: 'customAction',
|
||||
toolCallId: 'tool-call-id',
|
||||
});
|
||||
});
|
||||
|
||||
it('does not wait for custom onToolCall to finish', () => {
|
||||
const pendingToolCall = new Promise<void>(() => {});
|
||||
const onToolCall = vi.fn(() => pendingToolCall);
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {
|
||||
customAction: {
|
||||
onToolCall,
|
||||
render: () => 'Custom action complete',
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
const result = getOnToolCall()({
|
||||
toolCall: {
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
},
|
||||
});
|
||||
|
||||
expect(onToolCall).toHaveBeenCalledTimes(1);
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('sends the secure user token header when memory.userToken is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
memory: { userToken: 'secure-user-token' },
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportHeaders()).toMatchObject({
|
||||
'x-algolia-secure-user-token': 'secure-user-token',
|
||||
});
|
||||
});
|
||||
|
||||
it('omits the secure user token header when no memory token is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token');
|
||||
});
|
||||
|
||||
it('sends an empty transport body when no search parameters or indices are provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({ algolia: {} });
|
||||
});
|
||||
|
||||
it('includes searchParameters under the algolia body when provided', () => {
|
||||
const searchParameters = {
|
||||
'index-name': { distinct: false },
|
||||
};
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
searchParameters,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { searchParameters },
|
||||
});
|
||||
});
|
||||
|
||||
it('includes indices under the algolia body when provided', () => {
|
||||
const indices = [
|
||||
{
|
||||
index: 'docsearch-markdown',
|
||||
description: 'Use this to gather specific results.',
|
||||
},
|
||||
];
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
indices,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { indices },
|
||||
});
|
||||
});
|
||||
|
||||
it('includes both searchParameters and indices under the algolia body when both are provided', () => {
|
||||
const searchParameters = {
|
||||
'index-name': { distinct: false },
|
||||
};
|
||||
const indices = [
|
||||
{
|
||||
index: 'docsearch-markdown',
|
||||
description: 'Use this to gather specific results.',
|
||||
},
|
||||
];
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
searchParameters,
|
||||
indices,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { searchParameters, indices },
|
||||
});
|
||||
});
|
||||
|
||||
it('omits indices from the body when an empty indices array is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
indices: [],
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({ algolia: {} });
|
||||
});
|
||||
|
||||
describe('conversation id rotation', () => {
|
||||
it('passes a generated id to useChat on mount', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(chatOptions?.id).toBe('generated-id-1');
|
||||
});
|
||||
|
||||
it('rotates the chat id and clears messages on startNewConversation', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
const initialId = chatOptions?.id;
|
||||
expect(initialId).toBe('generated-id-1');
|
||||
|
||||
act(() => {
|
||||
result.current.startNewConversation();
|
||||
});
|
||||
|
||||
// Rotating the id recreates the `Chat` instance with no seeded messages,
|
||||
// which is what clears the previous conversation.
|
||||
expect(chatOptions?.messages).toBeUndefined();
|
||||
expect(chatOptions?.id).toBe('generated-id-2');
|
||||
expect(chatOptions?.id).not.toBe(initialId);
|
||||
});
|
||||
|
||||
it('generates a fresh id for each new conversation', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
act(() => {
|
||||
result.current.startNewConversation();
|
||||
});
|
||||
const secondId = chatOptions?.id;
|
||||
|
||||
act(() => {
|
||||
result.current.startNewConversation();
|
||||
});
|
||||
const thirdId = chatOptions?.id;
|
||||
|
||||
expect(secondId).toBe('generated-id-2');
|
||||
expect(thirdId).toBe('generated-id-3');
|
||||
expect(secondId).not.toBe(thirdId);
|
||||
});
|
||||
|
||||
it('restores a stored conversation using its persisted chat id', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
const storedMessages: ChatMessage[] = [
|
||||
{
|
||||
id: 'message-1',
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text: 'stored question' }],
|
||||
},
|
||||
];
|
||||
|
||||
act(() => {
|
||||
result.current.restoreConversation(storedMessages, 'stored-chat-id');
|
||||
});
|
||||
|
||||
// Restored messages are seeded through the `messages` option so they
|
||||
// survive the `Chat` instance being recreated when the id changes.
|
||||
expect(chatOptions?.messages).toEqual(storedMessages);
|
||||
expect(chatOptions?.id).toBe('stored-chat-id');
|
||||
});
|
||||
|
||||
it('restores a stored conversation with a fresh id when no chat id is persisted', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
const storedMessages: ChatMessage[] = [
|
||||
{
|
||||
id: 'message-1',
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text: 'stored question' }],
|
||||
},
|
||||
];
|
||||
|
||||
act(() => {
|
||||
result.current.restoreConversation(storedMessages);
|
||||
});
|
||||
|
||||
// The messages must still be seeded even when a fresh id is generated,
|
||||
// otherwise the recreated `Chat` instance would render an empty thread.
|
||||
expect(chatOptions?.messages).toEqual(storedMessages);
|
||||
expect(chatOptions?.id).toBe('generated-id-2');
|
||||
});
|
||||
|
||||
it('does not seed initial messages on mount', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(chatOptions?.messages).toBeUndefined();
|
||||
});
|
||||
|
||||
it('clears seeded messages and rotates the id when starting a new conversation after a restore', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
const storedMessages: ChatMessage[] = [
|
||||
{
|
||||
id: 'message-1',
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text: 'stored question' }],
|
||||
},
|
||||
];
|
||||
|
||||
act(() => {
|
||||
result.current.restoreConversation(storedMessages, 'stored-chat-id');
|
||||
});
|
||||
|
||||
expect(chatOptions?.messages).toEqual(storedMessages);
|
||||
|
||||
act(() => {
|
||||
result.current.startNewConversation();
|
||||
});
|
||||
|
||||
// Starting a new conversation must drop the seeded messages so a later
|
||||
// recreation does not re-hydrate the previous conversation.
|
||||
expect(chatOptions?.messages).toBeUndefined();
|
||||
expect(chatOptions?.id).toBe('generated-id-2');
|
||||
});
|
||||
|
||||
it('sends the next message with the new chat id after startNewConversation', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
const initialId = chatOptions?.id;
|
||||
expect(initialId).toBe('generated-id-1');
|
||||
|
||||
// Capture `sendMessage` from the current render, mirroring how a consumer
|
||||
// (e.g. the Sidepanel) grabs it before rotating the conversation.
|
||||
const sendMessage = result.current.sendMessage;
|
||||
|
||||
act(() => {
|
||||
result.current.startNewConversation();
|
||||
sendMessage({
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text: 'new question' }],
|
||||
});
|
||||
});
|
||||
|
||||
// The rotated id is what every later request must target.
|
||||
expect(chatOptions?.id).toBe('generated-id-2');
|
||||
expect(chatOptions?.id).not.toBe(initialId);
|
||||
|
||||
// Regression: the request must NOT go out on the stale (pre-rotation) id.
|
||||
expect(sendMessageCalls).toHaveLength(1);
|
||||
expect(sendMessageCalls[0].chatId).toBe('generated-id-2');
|
||||
expect(sendMessageCalls[0].chatId).not.toBe(initialId);
|
||||
});
|
||||
|
||||
it('keeps conversation lifecycle callbacks stable when transport inputs change by reference', () => {
|
||||
const { result, rerender } = renderHook(
|
||||
({ indices, searchParameters }) =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
searchParameters,
|
||||
indices,
|
||||
tools: {},
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
searchParameters: {
|
||||
'index-name': { distinct: false },
|
||||
},
|
||||
indices: [{ index: 'index-name', description: 'Test index' }],
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const startNewConversation = result.current.startNewConversation;
|
||||
const restoreConversation = result.current.restoreConversation;
|
||||
|
||||
rerender({
|
||||
searchParameters: {
|
||||
'index-name': { distinct: false },
|
||||
},
|
||||
indices: [{ index: 'index-name', description: 'Test index' }],
|
||||
});
|
||||
|
||||
expect(result.current.startNewConversation).toBe(startNewConversation);
|
||||
expect(result.current.restoreConversation).toBe(restoreConversation);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,305 +0,0 @@
|
|||
import { renderHook } from '@testing-library/react';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { useAskAi } from '../useAskAi';
|
||||
|
||||
type ToolCall = {
|
||||
input: unknown;
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
};
|
||||
|
||||
type ChatOptions = {
|
||||
onToolCall: (params: { toolCall: ToolCall }) => unknown;
|
||||
transport: {
|
||||
options: {
|
||||
headers?: Record<string, string>;
|
||||
body?: Record<string, unknown>;
|
||||
};
|
||||
};
|
||||
};
|
||||
|
||||
type CustomOnToolCallParams = ToolCall & {
|
||||
addToolOutput: (props: { output: unknown }) => Promise<void>;
|
||||
};
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
addToolOutput: vi.fn(),
|
||||
useChat: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('@ai-sdk/react', () => ({
|
||||
useChat: mocks.useChat,
|
||||
}));
|
||||
|
||||
vi.mock('ai', () => ({
|
||||
DefaultChatTransport: class DefaultChatTransport {
|
||||
options: unknown;
|
||||
|
||||
constructor(options: unknown) {
|
||||
this.options = options;
|
||||
}
|
||||
},
|
||||
lastAssistantMessageIsCompleteWithToolCalls: vi.fn(() => false),
|
||||
}));
|
||||
|
||||
describe('useAskAi', () => {
|
||||
let chatOptions: ChatOptions | undefined;
|
||||
|
||||
function getOnToolCall(): ChatOptions['onToolCall'] {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.onToolCall;
|
||||
}
|
||||
|
||||
function getTransportHeaders(): Record<string, string> {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.transport.options.headers ?? {};
|
||||
}
|
||||
|
||||
function getTransportBody(): Record<string, unknown> {
|
||||
if (!chatOptions) {
|
||||
throw new Error('useChat was not initialized');
|
||||
}
|
||||
|
||||
return chatOptions.transport.options.body ?? {};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
chatOptions = undefined;
|
||||
mocks.useChat.mockImplementation((options: ChatOptions) => {
|
||||
chatOptions = options;
|
||||
|
||||
return {
|
||||
addToolOutput: mocks.addToolOutput,
|
||||
error: undefined,
|
||||
messages: [],
|
||||
sendMessage: vi.fn(),
|
||||
setMessages: vi.fn(),
|
||||
status: 'ready',
|
||||
stop: vi.fn(),
|
||||
};
|
||||
});
|
||||
});
|
||||
|
||||
it('forwards custom tool output to useChat addToolOutput', async () => {
|
||||
let addToolOutput: CustomOnToolCallParams['addToolOutput'] | undefined;
|
||||
const onToolCall = vi.fn((params: CustomOnToolCallParams) => {
|
||||
addToolOutput = params.addToolOutput;
|
||||
});
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {
|
||||
customAction: {
|
||||
onToolCall,
|
||||
render: () => 'Custom action complete',
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
getOnToolCall()({
|
||||
toolCall: {
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
},
|
||||
});
|
||||
|
||||
expect(onToolCall).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
}),
|
||||
);
|
||||
|
||||
if (!addToolOutput) {
|
||||
throw new Error('addToolOutput was not provided to the custom tool');
|
||||
}
|
||||
|
||||
await addToolOutput({ output: { result: 'output value' } });
|
||||
|
||||
expect(mocks.addToolOutput).toHaveBeenCalledWith({
|
||||
output: { result: 'output value' },
|
||||
tool: 'customAction',
|
||||
toolCallId: 'tool-call-id',
|
||||
});
|
||||
});
|
||||
|
||||
it('does not wait for custom onToolCall to finish', () => {
|
||||
const pendingToolCall = new Promise<void>(() => {});
|
||||
const onToolCall = vi.fn(() => pendingToolCall);
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {
|
||||
customAction: {
|
||||
onToolCall,
|
||||
render: () => 'Custom action complete',
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
const result = getOnToolCall()({
|
||||
toolCall: {
|
||||
input: { value: 'input value' },
|
||||
toolCallId: 'tool-call-id',
|
||||
toolName: 'customAction',
|
||||
},
|
||||
});
|
||||
|
||||
expect(onToolCall).toHaveBeenCalledTimes(1);
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('sends the secure user token header when memory.userToken is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
memory: { userToken: 'secure-user-token' },
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportHeaders()).toMatchObject({
|
||||
'x-algolia-secure-user-token': 'secure-user-token',
|
||||
});
|
||||
});
|
||||
|
||||
it('omits the secure user token header when no memory token is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token');
|
||||
});
|
||||
|
||||
it('sends an empty transport body when no search parameters or indices are provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({ algolia: {} });
|
||||
});
|
||||
|
||||
it('includes searchParameters under the algolia body when provided', () => {
|
||||
const searchParameters = {
|
||||
'index-name': { distinct: false },
|
||||
};
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
searchParameters,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { searchParameters },
|
||||
});
|
||||
});
|
||||
|
||||
it('includes indices under the algolia body when provided', () => {
|
||||
const indices = [
|
||||
{
|
||||
index: 'docsearch-markdown',
|
||||
description: 'Use this to gather specific results.',
|
||||
},
|
||||
];
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
indices,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { indices },
|
||||
});
|
||||
});
|
||||
|
||||
it('includes both searchParameters and indices under the algolia body when both are provided', () => {
|
||||
const searchParameters = {
|
||||
'index-name': { distinct: false },
|
||||
};
|
||||
const indices = [
|
||||
{
|
||||
index: 'docsearch-markdown',
|
||||
description: 'Use this to gather specific results.',
|
||||
},
|
||||
];
|
||||
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
searchParameters,
|
||||
indices,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({
|
||||
algolia: { searchParameters, indices },
|
||||
});
|
||||
});
|
||||
|
||||
it('omits indices from the body when an empty indices array is provided', () => {
|
||||
renderHook(() =>
|
||||
useAskAi({
|
||||
apiKey: 'api-key',
|
||||
appId: 'app-id',
|
||||
assistantId: 'assistant-id',
|
||||
indexName: 'index-name',
|
||||
tools: {},
|
||||
indices: [],
|
||||
}),
|
||||
);
|
||||
|
||||
expect(getTransportBody()).toEqual({ algolia: {} });
|
||||
});
|
||||
});
|
||||
|
|
@ -14,6 +14,7 @@ export type StoredAskAiMessage = AIMessage & {
|
|||
};
|
||||
|
||||
export type StoredAskAiState = StoredDocSearchHit & {
|
||||
chatId?: string;
|
||||
stopped?: boolean;
|
||||
messages?: StoredAskAiMessage[];
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
import type { UseChatHelpers } from '@ai-sdk/react';
|
||||
import { useChat } from '@ai-sdk/react';
|
||||
import { Chat, useChat } from '@ai-sdk/react';
|
||||
import type { ChatOnToolCallCallback } from 'ai';
|
||||
import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
|
||||
import { useCallback, useEffect, useMemo, useRef } from 'react';
|
||||
import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, generateId } from 'ai';
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
|
||||
|
||||
import { agentStudioBaseUrl, getAgentStudioErrorMessage, postAgentStudioFeedback } from './askai';
|
||||
import type { Exchange } from './AskAiScreen';
|
||||
|
|
@ -28,6 +28,7 @@ type UseAskAiParams = {
|
|||
};
|
||||
|
||||
type UseAskAiReturn = {
|
||||
chatId: string;
|
||||
messages: AIMessage[];
|
||||
status: UseChat['status'];
|
||||
sendMessage: UseChat['sendMessage'];
|
||||
|
|
@ -38,6 +39,14 @@ type UseAskAiReturn = {
|
|||
exchanges: Exchange[];
|
||||
conversations: StoredSearchPlugin<StoredAskAiState>;
|
||||
sendFeedback: OnAskAiFeedback;
|
||||
/**
|
||||
* Create's a new chat instance, clearing existing messages and generating a new conversation ID.
|
||||
*/
|
||||
startNewConversation: () => void;
|
||||
/**
|
||||
* Create's a new chat instance, seeded with an existing conversation's ID and its messages.
|
||||
*/
|
||||
restoreConversation: (restored: AIMessage[], existingConversationId?: string) => void;
|
||||
};
|
||||
|
||||
type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn;
|
||||
|
|
@ -105,6 +114,12 @@ export const useAskAi: UseAskAi = ({
|
|||
[apiKey, appId, assistantId, searchParameters, memory?.userToken, indices],
|
||||
);
|
||||
|
||||
// Store transport in a ref since it is dependent on unstable dependencies:
|
||||
// - searchParameters, an object whose changed values trigger a new transport
|
||||
// - indices, an array whose changed values trigger a new transport
|
||||
const askAiTransportRef = useRef(askAiTransport);
|
||||
askAiTransportRef.current = askAiTransport;
|
||||
|
||||
// Sync ref during render so the stable `handleToolCall` (registered once
|
||||
// by useChat) always sees the latest `tools` without re-creating itself.
|
||||
// Safe because tool calls only fire after a commit, and writes are idempotent.
|
||||
|
|
@ -133,10 +148,25 @@ export const useAskAi: UseAskAi = ({
|
|||
});
|
||||
}, []);
|
||||
|
||||
const { messages, sendMessage, status, setMessages, error, stop, addToolOutput } = useChat({
|
||||
sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
|
||||
transport: askAiTransport,
|
||||
onToolCall: handleToolCall,
|
||||
const createChatInstance = useCallback(
|
||||
(messages?: AIMessage[], id = generateId()): Chat<AIMessage> =>
|
||||
new Chat<AIMessage>({
|
||||
id,
|
||||
messages,
|
||||
sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
|
||||
transport: askAiTransportRef.current,
|
||||
onToolCall: handleToolCall,
|
||||
}),
|
||||
[handleToolCall],
|
||||
);
|
||||
|
||||
const [chatInstance, setChatInstance] = useState((): Chat<AIMessage> => createChatInstance());
|
||||
// Keep a stable reference to the chat instance so reading messages are render safe
|
||||
const chatInstanceRef = useRef<Chat<AIMessage>>(chatInstance);
|
||||
chatInstanceRef.current = chatInstance;
|
||||
|
||||
const { messages, status, setMessages, error, stop, addToolOutput } = useChat({
|
||||
chat: chatInstance,
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -171,12 +201,12 @@ export const useAskAi: UseAskAi = ({
|
|||
[assistantId, appId, apiKey, conversations],
|
||||
);
|
||||
|
||||
const onStopStreaming = async (): Promise<void> => {
|
||||
const onStopStreaming = useCallback(async (): Promise<void> => {
|
||||
abortControllerRef.current.abort();
|
||||
await stop();
|
||||
};
|
||||
}, [stop]);
|
||||
|
||||
const exchanges = useMemo(() => {
|
||||
const exchanges = useMemo((): Exchange[] => {
|
||||
const grouped: Exchange[] = [];
|
||||
|
||||
for (let i = 0; i < messages.length; i++) {
|
||||
|
|
@ -201,9 +231,36 @@ export const useAskAi: UseAskAi = ({
|
|||
return getAgentStudioErrorMessage(error);
|
||||
}, [error]);
|
||||
|
||||
const updateChatInstance = useCallback(
|
||||
(restored?: AIMessage[], existingConversationId?: string): void => {
|
||||
const newChatInstance = createChatInstance(restored, existingConversationId);
|
||||
chatInstanceRef.current = newChatInstance;
|
||||
setChatInstance(newChatInstance);
|
||||
},
|
||||
[createChatInstance],
|
||||
);
|
||||
|
||||
const startNewConversation = useCallback((): void => {
|
||||
updateChatInstance();
|
||||
}, [updateChatInstance]);
|
||||
|
||||
const restoreConversation = useCallback(
|
||||
(restored: AIMessage[], existingConversationId?: string): void => {
|
||||
updateChatInstance(restored, existingConversationId);
|
||||
},
|
||||
[updateChatInstance],
|
||||
);
|
||||
|
||||
// This is so that the public `sendMessage` is always pointed to a stable reference of the chat instance
|
||||
const sendMessageSafe = useCallback<UseChat['sendMessage']>(
|
||||
(...args): Promise<void> => chatInstanceRef.current!.sendMessage(...args),
|
||||
[],
|
||||
);
|
||||
|
||||
return {
|
||||
chatId: chatInstance.id,
|
||||
messages,
|
||||
sendMessage,
|
||||
sendMessage: sendMessageSafe,
|
||||
status,
|
||||
setMessages,
|
||||
askAiError,
|
||||
|
|
@ -212,5 +269,7 @@ export const useAskAi: UseAskAi = ({
|
|||
exchanges,
|
||||
conversations,
|
||||
sendFeedback,
|
||||
startNewConversation,
|
||||
restoreConversation,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -79,13 +79,14 @@ export function extractLinksFromMessage(message: AIMessage | null): ExtractedLin
|
|||
return links;
|
||||
}
|
||||
|
||||
export const buildDummyAskAiHit = (query: string, messages: AIMessage[]): StoredAskAiState => {
|
||||
export const buildDummyAskAiHit = (query: string, messages: AIMessage[], chatId?: string): StoredAskAiState => {
|
||||
const textPart = messages[0].parts.find((part) => part.type === 'text');
|
||||
const sanitizedText = textPart?.text ? sanitizeUserInput(textPart.text) : '';
|
||||
|
||||
return {
|
||||
query,
|
||||
objectID: sanitizedText,
|
||||
chatId,
|
||||
messages,
|
||||
type: 'askAI',
|
||||
anchor: 'stored',
|
||||
|
|
|
|||
Loading…
Reference in a new issue