1
0
Fork 0

feat(askai): remove Ask AI transport layer (#2889)

* feat(askai): add Agent Studio memory support

* refactor(askai): remove Ask AI transport abstraction
This commit is contained in:
Paul Jankowski 2026-06-01 17:00:38 -04:00 committed by GitHub
parent 0101c92a1e
commit 144250a62c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 86 additions and 419 deletions

View file

@ -222,7 +222,6 @@ function DocSearch({ externalUrlRegex, ...props }: AlgoliaSearchBarProps) {
sidePanelEnabled,
showSidepanelButton,
sidePanelOptions,
sidePanelAgentStudio,
sidepanelPortalContainer,
isSidepanelOpen,
sidepanelInitialMessage,
@ -391,7 +390,6 @@ function DocSearch({ externalUrlRegex, ...props }: AlgoliaSearchBarProps) {
apiKey={askAi.apiKey}
appId={askAi.appId}
indexName={askAi.indexName}
agentStudio={sidePanelAgentStudio}
suggestedQuestions={sidePanelOptions?.suggestedQuestions ?? askAi.suggestedQuestions}
isOpen={isSidepanelOpen}
initialMessage={sidepanelInitialMessage}

View file

@ -7,8 +7,6 @@ import './App.css';
import '@docsearch/css/dist/style.css';
import '@docsearch/css/dist/sidepanel.css';
import AgentStudio from './examples/agent-studio';
import AgentStudioSidepanel from './examples/agent-studio-sidepanel';
import Basic from './examples/basic';
import BasicAskAI from './examples/basic-askai';
import Composable from './examples/composable';
@ -99,20 +97,6 @@ function App(): JSX.Element {
<BasicHybrid />
</div>
</section>
<section className="demo-section">
<p className="section-description">Agent Studio</p>
<div className="search-wrapper column">
<AgentStudio />
</div>
</section>
<section className="demo-section">
<p className="section-description">Agent Studio sidepanel</p>
<div className="search-wrapper column">
<AgentStudioSidepanel />
</div>
</section>
</main>
</div>

View file

@ -1,42 +0,0 @@
/* eslint-disable react/react-in-jsx-scope */
import { DocSearch } from '@docsearch/core';
import { SidepanelButton } from '@docsearch/sidepanel/button';
import { Sidepanel } from '@docsearch/sidepanel/sidepanel';
import type { JSX } from 'react';
export default function AgentStudioSidepanel(): JSX.Element {
return (
<DocSearch>
<SidepanelButton
variant="inline"
translations={{
buttonText: 'Agent Studio',
}}
/>
<Sidepanel
indexName="docsearch-markdown"
appId="PMZUYBQDAK"
apiKey="a00716d83c64f6c61905c078b7d5ab66"
assistantId="ccdec697-e3fe-465b-a1c3-657e7bf18aef"
agentStudio={true}
tools={{
printConsoleMessage: {
render({ message: { output } }) {
if (!output) return '';
return output as string;
},
async onToolCall({ input, addToolOutput }) {
// eslint-disable-next-line no-console
console.log((input as any).message);
await addToolOutput({
output: 'Check your console for a nice message :)',
});
},
},
}}
/>
</DocSearch>
);
}

View file

@ -1,42 +0,0 @@
/* eslint-disable react/react-in-jsx-scope */
import { DocSearch } from '@docsearch/core';
import { DocSearchButton, DocSearchAskAiModal } from '@docsearch/modal';
import type { JSX } from 'react';
export default function AgentStudio(): JSX.Element {
return (
<DocSearch>
<DocSearchButton
translations={{
buttonText: 'Ask AI with Agent Studio',
}}
/>
<DocSearchAskAiModal
indexName="docsearch"
appId="PMZUYBQDAK"
apiKey="a00716d83c64f6c61905c078b7d5ab66"
askAi={{
assistantId: 'ccdec697-e3fe-465b-a1c3-657e7bf18aef',
agentStudio: true,
}}
tools={{
printConsoleMessage: {
render({ message: { output } }) {
if (!output) return '';
return output as string;
},
async onToolCall({ input, addToolOutput }) {
// eslint-disable-next-line no-console
console.log((input as any).message);
await addToolOutput({
output: 'Check your console for a nice message :)',
});
},
},
}}
/>
</DocSearch>
);
}

View file

@ -9,14 +9,28 @@ export default function BasicAskAI(): JSX.Element {
appId="PMZUYBQDAK"
apiKey="24b09689d5b4223813d9b8e48563c8f6"
askAi={{
assistantId: 'askAIDemo',
searchParameters: {
facetFilters: ['language:en'],
},
assistantId: 'ccdec697-e3fe-465b-a1c3-657e7bf18aef',
suggestedQuestions: true,
}}
insights={true}
translations={{ button: { buttonText: 'Search with Ask AI' } }}
tools={{
printConsoleMessage: {
render({ message: { output } }) {
if (!output) return '';
return output as string;
},
async onToolCall({ input, addToolOutput }) {
// eslint-disable-next-line no-console
console.log((input as any).message);
await addToolOutput({
output: 'Check your console for a nice message :)',
});
},
},
}}
/>
);
}

View file

@ -11,8 +11,25 @@ export default function SidepanelExample(): JSX.Element {
indexName="docsearch"
appId="PMZUYBQDAK"
apiKey="24b09689d5b4223813d9b8e48563c8f6"
assistantId="askAIDemo"
assistantId="ccdec697-e3fe-465b-a1c3-657e7bf18aef"
variant="floating"
tools={{
printConsoleMessage: {
render({ message: { output } }) {
if (!output) return '';
return output as string;
},
async onToolCall({ input, addToolOutput }) {
// eslint-disable-next-line no-console
console.log((input as any).message);
await addToolOutput({
output: 'Check your console for a nice message :)',
});
},
},
}}
/>
</DocSearch>
);

View file

@ -102,7 +102,6 @@ type AskAiScreenProps = Omit<AskAiScreenStateProps<InternalDocSearchHit>, 'trans
askAiError?: Error;
translations?: AskAiScreenTranslations;
onNewConversation: () => void;
agentStudio?: boolean;
memoryEnabled?: boolean;
};
@ -130,7 +129,6 @@ interface AskAiExchangeCardProps {
tools: ToolCalls;
conversations: StoredSearchPlugin<StoredAskAiState>;
onFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
agentStudio?: boolean;
memoryEnabled?: boolean;
}
@ -144,7 +142,6 @@ function AskAiExchangeCard({
tools,
conversations,
onFeedback,
agentStudio,
memoryEnabled,
}: AskAiExchangeCardProps): JSX.Element {
const { userMessage, assistantMessage } = exchange;
@ -174,7 +171,7 @@ function AskAiExchangeCard({
isLastExchange &&
!displayParts.some((part) => part.type !== 'step-start');
const messageId = agentStudio ? assistantMessage?.id || exchange.id : userMessage?.id || exchange.id;
const messageId = assistantMessage?.id || exchange.id;
return (
<div className="DocSearch-AskAiScreen-Response-Container">
@ -405,7 +402,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
startNewConversationButtonText = 'Start a new conversation',
} = translations;
const { messages, tools, askAiError, status, agentStudio, memoryEnabled } = props;
const { messages, tools, askAiError, status, memoryEnabled } = props;
// Check if there's a thread depth error
const hasThreadDepthError = useMemo(() => {
@ -481,7 +478,6 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
translations={translations}
tools={tools}
conversations={props.conversations}
agentStudio={agentStudio}
memoryEnabled={memoryEnabled}
onSearchQueryClick={handleSearchQueryClick}
onFeedback={props.onFeedback}

View file

@ -57,7 +57,6 @@ export interface AskAiScreenStateProps<TItem extends BaseItem>
suggestedQuestions: SuggestedQuestionHit[];
selectSuggestedQuestion: (question: SuggestedQuestionHit) => void;
onNewConversation: () => void;
agentStudio?: boolean;
memoryEnabled?: boolean;
}
@ -86,7 +85,6 @@ export const AskAiScreenState = React.memo(
status={props.status}
askAiError={props.askAiError}
translations={translations?.askAiScreen}
agentStudio={props.agentStudio}
memoryEnabled={props.memoryEnabled}
/>
);

View file

@ -63,45 +63,17 @@ export type DocSearchAskAi = {
* @default false
*/
suggestedQuestions?: boolean;
// HACK: This is a hack for testing staging, remove before releasing
useStagingEnv?: boolean;
} & (
| {
/**
* **Experimental:** Whether to use Agent Studio as the chat backend.
*
* This is an experimental feature and its API may change without notice in future releases.
* Use with caution in production environments.
*
* @default false
*/
agentStudio?: never;
/**
* The search parameters to use for the ask AI feature.
*
* **NOTE**: If using `agentStudio = true`, the `searchParameters` object is
* keyed by the index name.
*/
searchParameters?: AskAiSearchParameters;
}
| {
agentStudio: false;
searchParameters?: AskAiSearchParameters;
}
| {
agentStudio: true;
/**
* The search parameters to use for the ask AI feature.
* Keyed by the index name.
*
* @example
* {
* "INDEX_NAME": { distinct: false }
* }
*/
searchParameters?: AgentStudioSearchParameters;
}
);
/**
* The search parameters to use for the ask AI feature.
* Keyed by the index name.
*
* @example
* {
* "INDEX_NAME": { distinct: false }
* }
*/
searchParameters?: AgentStudioSearchParameters;
};
export interface DocSearchIndex {
name: string;

View file

@ -109,14 +109,12 @@ export function DocSearchAskAiModal({
const askAiConfig = typeof askAi === 'object' ? askAi : null;
const askAiConfigurationId = typeof askAi === 'string' ? askAi : askAiConfig?.assistantId || null;
const askAiSearchParameters = askAiConfig?.searchParameters;
const askAiUseStagingEnv = askAiConfig?.useStagingEnv || false;
const [askAiState, setAskAiState] = React.useState<AskAiState>('initial');
const suggestedQuestions = useSuggestedQuestions({
assistantId: askAiConfigurationId,
searchClient,
suggestedQuestionsEnabled: askAiConfig?.suggestedQuestions,
});
const agentStudio = askAiConfig?.agentStudio ?? false;
const memoryEnabled = props.memory?.enabled ?? false;
const indexes = normalizeDocSearchIndexes({
@ -141,8 +139,6 @@ export function DocSearchAskAiModal({
appId: askAiConfig?.appId || appId,
indexName: askAiConfig?.indexName || defaultIndexName,
searchParameters: askAiSearchParameters,
useStagingEnv: askAiUseStagingEnv,
agentStudio,
tools,
memory: props.memory,
});
@ -465,7 +461,6 @@ export function DocSearchAskAiModal({
selectAskAiQuestion={handleSelectAskAiQuestion}
suggestedQuestions={suggestedQuestions}
selectSuggestedQuestion={selectSuggestedQuestion}
agentStudio={agentStudio}
memoryEnabled={memoryEnabled}
onAskAiToggle={onAskAiToggle}
onNewConversation={handleNewConversation}

View file

@ -4,49 +4,25 @@ import type { JSX } from 'react';
import React from 'react';
import { createPortal } from 'react-dom';
import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory } from './DocSearch';
import type { AgentStudioSearchParameters, Memory } from './DocSearch';
import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index';
import { SidepanelButton, Sidepanel } from './Sidepanel/index';
import type { ToolCalls } from './types/AskiAi';
export type { DocSearchRef, DocSearchCallbacks } from '@docsearch/core';
export type SidepanelSearchParameters =
| {
/**
* **Experimental:** Whether to use Agent Studio as the chat backend.
*
* This is an experimental feature and its API may change without notice in future releases.
* Use with caution in production environments.
*
* @default false
*/
agentStudio?: never;
/**
* The search parameters to use for the ask AI feature.
*
* **NOTE**: If using `agentStudio = true`, the `searchParameters` object is
* keyed by the index name.
*/
searchParameters?: AskAiSearchParameters;
}
| {
agentStudio: false;
searchParameters?: AskAiSearchParameters;
}
| {
agentStudio: true;
/**
* The search parameters to use for the ask AI feature.
* Keyed by the index name.
*
* @example
* {
* "INDEX_NAME": { distinct: false }
* }
*/
searchParameters?: AgentStudioSearchParameters;
};
export type SidepanelSearchParameters = {
/**
* The search parameters to use for the ask AI feature.
* Keyed by the index name.
*
* @example
* {
* "INDEX_NAME": { distinct: false }
* }
*/
searchParameters?: AgentStudioSearchParameters;
};
export type DocSearchSidepanelProps = DocSearchCallbacks & {
/**

View file

@ -71,7 +71,6 @@ export type ConversationScreenProps = {
status: UseChatHelpers<AIMessage>['status'];
handleFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
streamError?: Error;
agentStudio?: boolean;
memoryEnabled?: boolean;
tools?: ToolCalls;
};
@ -84,7 +83,6 @@ type ConversationnExchangeProps = {
translations?: ConversationScreenTranslations;
onFeedback?: ConversationScreenProps['handleFeedback'];
streamError?: ConversationScreenProps['streamError'];
agentStudio?: boolean;
memoryEnabled?: boolean;
tools: ToolCalls;
};
@ -99,7 +97,6 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
onFeedback,
status,
streamError,
agentStudio,
memoryEnabled,
tools,
},
@ -137,7 +134,7 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
const showActions =
!wasStopped && (!isLastExchange || (isLastExchange && status === 'ready' && Boolean(assistantMessage)));
const messageId = agentStudio ? assistantMessage?.id || exchange.id : userMessage?.id || exchange.id;
const messageId = assistantMessage?.id || exchange.id;
return (
<div className="DocSearch-AskAiScreen-Response-Container" ref={conversationRef}>

View file

@ -118,8 +118,6 @@ export type SidepanelProps = {
* @default `{ 'Ctrl/Cmd+I': true }`
*/
keyboardShortcuts?: SidepanelShortcuts;
// HACK: This is a hack for testing staging, remove before releasing
useStagingEnv?: boolean;
};
type Props = Omit<DocSearchSidepanelProps, 'button' | 'panel'> &
@ -150,8 +148,6 @@ function SidepanelInner(
keyboardShortcuts,
side = 'right',
initialMessage,
useStagingEnv = false,
agentStudio = false,
tools = EMPTY_TOOLS,
memory,
}: Props,
@ -195,8 +191,6 @@ function SidepanelInner(
assistantId,
apiKey,
searchParameters,
useStagingEnv,
agentStudio,
tools,
memory,
});
@ -401,7 +395,6 @@ function SidepanelInner(
handleFeedback={sendFeedback}
translations={translations.conversationScreen}
streamError={askAiError}
agentStudio={agentStudio}
memoryEnabled={memory?.enabled ?? false}
tools={tools}
/>

View file

@ -84,7 +84,6 @@ describe('useAskAi', () => {
renderHook(() =>
useAskAi({
agentStudio: false,
apiKey: 'api-key',
appId: 'app-id',
assistantId: 'assistant-id',
@ -133,7 +132,6 @@ describe('useAskAi', () => {
renderHook(() =>
useAskAi({
agentStudio: false,
apiKey: 'api-key',
appId: 'app-id',
assistantId: 'assistant-id',
@ -159,10 +157,9 @@ describe('useAskAi', () => {
expect(result).toBeUndefined();
});
it('sends the secure user token header when memory.userToken is provided in Agent Studio mode', () => {
it('sends the secure user token header when memory.userToken is provided', () => {
renderHook(() =>
useAskAi({
agentStudio: true,
apiKey: 'api-key',
appId: 'app-id',
assistantId: 'assistant-id',
@ -180,7 +177,6 @@ describe('useAskAi', () => {
it('omits the secure user token header when no memory token is provided', () => {
renderHook(() =>
useAskAi({
agentStudio: true,
apiKey: 'api-key',
appId: 'app-id',
assistantId: 'assistant-id',

View file

@ -1,110 +1,5 @@
import { ASK_AI_API_URL, BETA_ASK_AI_API_URL } from './constants';
// ... existing imports ...
const TOKEN_KEY = 'askai_token';
export const agentStudioBaseUrl = (appId: string): string => `https://${appId}.algolia.net/agent-studio/1`;
type TokenPayload = { exp: number };
const decode = (token: string): TokenPayload => {
const [b64] = token.split('.');
return JSON.parse(atob(b64));
};
const isExpired = (token?: string | null): boolean => {
if (!token) return true;
try {
const { exp } = decode(token);
// refresh 30 s before the backend rejects it
return Date.now() / 1000 > exp - 30;
} catch {
return true;
}
};
let inflight: Promise<string> | null = null;
// call /token once, cache the promise while its running
export const getValidToken = async ({
assistantId,
abortSignal,
useStagingEnv = false,
}: {
assistantId: string;
abortSignal: AbortSignal;
useStagingEnv?: boolean;
// eslint-disable-next-line require-await
}): Promise<string | null> => {
const cached = sessionStorage.getItem(TOKEN_KEY);
if (!isExpired(cached)) return cached!;
const baseUrl = useStagingEnv ? BETA_ASK_AI_API_URL : ASK_AI_API_URL;
if (!inflight) {
inflight = fetch(`${baseUrl}/token`, {
method: 'POST',
headers: {
'x-algolia-assistant-id': assistantId,
'content-type': 'application/json',
},
signal: abortSignal,
})
.then((r) => r.json())
.then(({ token, success, message }) => {
// If request was unsuccessful, throw an error to prevent calling `/chat` without a token
if (!success && message) {
throw new Error(message);
}
sessionStorage.setItem(TOKEN_KEY, token);
return token;
})
.finally(() => (inflight = null));
}
return inflight;
};
export const postFeedback = async ({
assistantId,
thumbs,
messageId,
appId,
abortSignal,
useStagingEnv = false,
}: {
assistantId: string;
thumbs: 0 | 1;
messageId: string;
appId: string;
abortSignal: AbortSignal;
useStagingEnv?: boolean;
}): Promise<Response> => {
const headers = new Headers();
headers.set('x-algolia-assistant-id', assistantId);
headers.set('content-type', 'application/json');
const token = await getValidToken({
assistantId,
abortSignal,
useStagingEnv,
});
headers.set('authorization', `TOKEN ${token}`);
const baseUrl = useStagingEnv ? BETA_ASK_AI_API_URL : ASK_AI_API_URL;
return fetch(`${baseUrl}/feedback`, {
method: 'POST',
body: JSON.stringify({
appId,
messageId,
thumbs,
}),
headers,
});
};
interface AgentStudioValidationError extends Error {
name: 'ValidationError';
detail?: Array<{ type: string; loc: string[]; msg: string }>;

View file

@ -1,4 +1,2 @@
export const MAX_QUERY_SIZE = 512;
export const ASK_AI_API_URL = 'https://askai.algolia.com/chat';
export const BETA_ASK_AI_API_URL = 'https://beta-chat-askai.algolia.com';
export const SUGGESTED_QUETIONS_INDEX_NAME = 'algolia_ask_ai_suggested_questions';

View file

@ -4,21 +4,14 @@ import type { ChatOnToolCallCallback } from 'ai';
import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
import { useCallback, useEffect, useMemo, useRef } from 'react';
import {
agentStudioBaseUrl,
getAgentStudioErrorMessage,
getValidToken,
postAgentStudioFeedback,
postFeedback,
} from './askai';
import { agentStudioBaseUrl, getAgentStudioErrorMessage, postAgentStudioFeedback } from './askai';
import type { Exchange } from './AskAiScreen';
import { ASK_AI_API_URL, BETA_ASK_AI_API_URL } from './constants';
import type { StoredSearchPlugin } from './stored-searches';
import { createStoredConversations } from './stored-searches';
import { type AIMessage, type ToolCalls } from './types/AskiAi';
import { EMPTY_TOOLS } from './utils/ai';
import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory, StoredAskAiState } from '.';
import type { AgentStudioSearchParameters, Memory, StoredAskAiState } from '.';
type UseChat = UseChatHelpers<AIMessage>;
@ -27,21 +20,10 @@ type UseAskAiParams = {
apiKey: string;
appId: string;
indexName: string;
useStagingEnv?: boolean;
searchParameters?: AskAiSearchParameters;
searchParameters?: AgentStudioSearchParameters;
tools: ToolCalls;
agentStudio: boolean;
memory?: Memory;
} & (
| {
agentStudio: false;
searchParameters?: AskAiSearchParameters;
}
| {
agentStudio: true;
searchParameters?: AgentStudioSearchParameters;
}
);
};
type UseAskAiReturn = {
messages: AIMessage[];
@ -81,51 +63,12 @@ const getAgentStudioTransport = ({
});
};
const getAskAiTransport = ({
assistantId,
apiKey,
indexName,
searchParameters,
appId,
abortController,
useStagingEnv,
}: Pick<UseAskAiParams, 'apiKey' | 'appId' | 'assistantId' | 'indexName' | 'searchParameters' | 'useStagingEnv'> & {
abortController: AbortController;
}): DefaultChatTransport<AIMessage> => {
return new DefaultChatTransport({
api: useStagingEnv ? BETA_ASK_AI_API_URL : ASK_AI_API_URL,
headers: async (): Promise<Record<string, string>> => {
if (!assistantId) {
throw new Error('Ask AI assistant ID is required');
}
const token = await getValidToken({
assistantId,
abortSignal: abortController.signal,
useStagingEnv,
});
return {
...(token ? { authorization: `TOKEN ${token}` } : {}),
'X-Algolia-API-Key': apiKey,
'X-Algolia-Application-Id': appId,
'X-Algolia-Index-Name': indexName,
'X-Algolia-Assistant-Id': assistantId || '',
'X-AI-SDK-Version': 'v5',
};
},
body: searchParameters ? { searchParameters } : {},
});
};
export const useAskAi: UseAskAi = ({
assistantId,
apiKey,
appId,
indexName,
useStagingEnv = false,
tools = EMPTY_TOOLS,
agentStudio,
searchParameters,
memory,
}) => {
@ -133,24 +76,14 @@ export const useAskAi: UseAskAi = ({
const askAiTransport = useMemo(
() =>
agentStudio
? getAgentStudioTransport({
apiKey,
appId,
assistantId: assistantId ?? '',
searchParameters,
userToken: memory?.userToken,
})
: getAskAiTransport({
assistantId: assistantId ?? '',
apiKey,
appId,
indexName,
searchParameters,
abortController: abortControllerRef.current,
useStagingEnv,
}),
[apiKey, appId, assistantId, indexName, useStagingEnv, agentStudio, searchParameters, memory?.userToken],
getAgentStudioTransport({
apiKey,
appId,
assistantId: assistantId ?? '',
searchParameters,
userToken: memory?.userToken,
}),
[apiKey, appId, assistantId, searchParameters, memory?.userToken],
);
// Sync ref during render so the stable `handleToolCall` (registered once
@ -202,28 +135,19 @@ export const useAskAi: UseAskAi = ({
async (messageId: string, thumbs: 0 | 1): Promise<void> => {
if (!assistantId) return;
const res = await (agentStudio
? postAgentStudioFeedback({
agentId: assistantId,
vote: thumbs,
messageId,
appId,
apiKey,
abortSignal: abortControllerRef.current.signal,
})
: postFeedback({
assistantId,
thumbs,
messageId,
appId,
abortSignal: abortControllerRef.current.signal,
useStagingEnv,
}));
const res = await postAgentStudioFeedback({
agentId: assistantId,
vote: thumbs,
messageId,
appId,
apiKey,
abortSignal: abortControllerRef.current.signal,
});
if (res.status >= 300) throw new Error('Failed, try again later.');
conversations.addFeedback?.(messageId, thumbs === 1 ? 'like' : 'dislike');
},
[assistantId, agentStudio, appId, apiKey, useStagingEnv, conversations],
[assistantId, appId, apiKey, conversations],
);
const onStopStreaming = async (): Promise<void> => {
@ -253,10 +177,8 @@ export const useAskAi: UseAskAi = ({
const askAiError = useMemo((): Error | undefined => {
if (!error) return undefined;
if (!agentStudio) return error;
return getAgentStudioErrorMessage(error);
}, [error, agentStudio]);
}, [error]);
return {
messages,