feat(askai): Allow dynamic indices for Agent Studio (#2893)

examples/demo-react/src/examples/hybrid.tsx

@@ -10,20 +10,19 @@ export default function BasicHybrid(): JSX.Element {
       <DocSearchButton />
       <DocSearchAskAiModal
         indexName="docsearch"
-        appId="beta3G7FSQDJR3"
+        appId="PMZUYBQDAK"
-        apiKey="0faad3eae2ba413c16355a0f8670c201"
+        apiKey="24b09689d5b4223813d9b8e48563c8f6"
         askAi={{
-          assistantId: 'e3Kl4lTCBlSA',
+          assistantId: 'ccdec697-e3fe-465b-a1c3-657e7bf18aef',
-          indexName: 'docsearch-markdown',
         }}
       />
 
       <SidepanelButton />
       <Sidepanel
         indexName="docsearch-markdown"
-        appId="beta3G7FSQDJR3"
+        appId="PMZUYBQDAK"
-        apiKey="0faad3eae2ba413c16355a0f8670c201"
+        apiKey="24b09689d5b4223813d9b8e48563c8f6"
-        assistantId="e3Kl4lTCBlSA"
+        assistantId="ccdec697-e3fe-465b-a1c3-657e7bf18aef"
       />
     </DocSearch>
   );

packages/docsearch-react/src/DocSearch.tsx

@@ -1,6 +1,6 @@
 import type { AutocompleteOptions, AutocompleteState } from '@algolia/autocomplete-core';
 import { DocSearch as DocSearchProvider, useDocSearch } from '@docsearch/core';
-import type { DocSearchModalShortcuts, DocSearchRef, InitialAskAiMessage } from '@docsearch/core';
+import type { DocSearchModalShortcuts, DocSearchRef } from '@docsearch/core';
 import type { LiteClient, SearchParamsObject } from 'algoliasearch/lite';
 import React, { type JSX } from 'react';
 import { createPortal } from 'react-dom';
@@ -10,7 +10,6 @@ import type { ButtonTranslations } from './DocSearchButton';
 import { DocSearchModal } from './DocSearchModal';
 import type { ModalTranslations } from './DocSearchModal';
 import type { DocSearchHit, DocSearchTheme, InternalDocSearchHit, StoredDocSearchHit } from './types';
-import type { ToolCalls } from './types/AskiAi';
 
 export type { DocSearchRef } from '@docsearch/core';
 
@@ -26,55 +25,6 @@ export type DocSearchTransformClient = {
   transporter: Pick<LiteClient['transporter'], 'algoliaAgent'>;
 };
 
-// Define the specific search parameters allowed for Ask AI
-export type AskAiSearchParameters = {
-  facetFilters?: string[];
-  filters?: string;
-  attributesToRetrieve?: string[];
-  restrictSearchableAttributes?: string[];
-  distinct?: boolean | number | string;
-};
-
-export type AgentStudioSearchParameters = Record<string, Omit<AskAiSearchParameters, 'facetFilters'>>;
-
-export type DocSearchAskAi = {
-  /**
-   * The index name to use for the ask AI feature. Your assistant will search this index for relevant documents.
-   * If not provided, the index name will be used.
-   */
-  indexName?: string;
-  /**
-   * The API key to use for the ask AI feature. Your assistant will use this API key to search the index.
-   * If not provided, the API key will be used.
-   */
-  apiKey?: string;
-  /**
-   * The app ID to use for the ask AI feature. Your assistant will use this app ID to search the index.
-   * If not provided, the app ID will be used.
-   */
-  appId?: string;
-  /**
-   * The assistant ID to use for the ask AI feature.
-   */
-  assistantId: string;
-  /**
-   * Enables displaying suggested questions on Ask AI's new conversation screen.
-   *
-   * @default false
-   */
-  suggestedQuestions?: boolean;
-  /**
-   * 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;
   searchParameters?: SearchParamsObject;
@@ -206,47 +156,6 @@ export interface DocSearchProps {
   keyboardShortcuts?: DocSearchModalShortcuts;
 }
 
-export interface Memory {
-  /**
-   * Determines whether or not to display the memory based tool calls.
-   *
-   * @default false
-   */
-  enabled?: boolean;
-  /**
-   * The JWT used by the agent to know which user's memory to read.
-   *
-   * @see https://www.algolia.com/doc/guides/algolia-ai/agent-studio/how-to/user-authentication
-   */
-  userToken?: string;
-}
-
-export interface DocSearchAIProps extends DocSearchProps {
-  /**
-   * Configuration or assistant id to enable ask ai mode. Pass a string assistant id or a full config object.
-   */
-  askAi: DocSearchAskAi | string;
-  /**
-   * Intercept Ask AI requests (e.g. Submitting a prompt or selecting a suggested question).
-   *
-   * Return `true` to prevent the default modal Ask AI flow (no toggle, no sendMessage).
-   * Useful to route Ask AI into a different UI (e.g. `@docsearch/sidepanel-js`) without flicker.
-   */
-  interceptAskAiEvent?: (initialMessage: InitialAskAiMessage) => boolean | void;
-  /**
-   * Use custom tools driven by Agent Studio.
-   *
-   * For best performance, memoize this object with `useMemo` or define it
-   * outside the component. Inline object literals will be recreated every
-   * render but will not affect correctness.
-   **/
-  tools?: ToolCalls;
-  /**
-   * Configuration for the Agent Studio memory feature.
-   */
-  memory?: Memory;
-}
-
 function DocSearchComponent(props: DocSearchProps, ref: React.ForwardedRef<DocSearchRef>): JSX.Element {
   return (
     <DocSearchProvider {...props} ref={ref}>

packages/docsearch-react/src/DocSearchAI.tsx

@@ -1,11 +1,214 @@
 import { DocSearch as DocSearchProvider, useDocSearch } from '@docsearch/core';
-import type { DocSearchRef } from '@docsearch/core';
+import type { DocSearchRef, InitialAskAiMessage } from '@docsearch/core';
+import type { SearchParamsObject } from 'algoliasearch/lite';
 import React, { type JSX } from 'react';
 import { createPortal } from 'react-dom';
 
-import type { DocSearchAIProps } from './DocSearch';
+import type { DocSearchProps } from './DocSearch';
 import { DocSearchAskAiModal } from './DocSearchAskAiModal';
 import { DocSearchButton } from './DocSearchButton';
+import type { ToolCalls } from './types/AskiAi';
+
+export interface AskAiSearchParameters {
+  facetFilters?: string[];
+  filters?: string;
+  attributesToRetrieve?: string[];
+  restrictSearchableAttributes?: string[];
+  distinct?: boolean | number | string;
+}
+
+export type AgentStudioSearchParameters = Record<string, Omit<AskAiSearchParameters, 'facetFilters'>>;
+
+interface IndexTextParam {
+  exposed: boolean;
+  default?: string;
+}
+
+interface NumberConstraint {
+  min?: number;
+  max?: number;
+}
+
+interface IndexNumberParam {
+  exposed: boolean;
+  default?: number;
+  constraint?: NumberConstraint;
+}
+
+interface StringArrayConstraints {
+  values?: string[];
+}
+
+interface IndexStringArrayParam {
+  exposed: boolean;
+  default?: string[];
+  constraint?: StringArrayConstraints;
+  merge?: boolean;
+}
+
+interface IndexFacetParam {
+  exposed: false;
+  default?: string[];
+}
+
+export interface AgentStudioSearchControls {
+  /**
+   * Augmented query for the MCP search tool to use.
+   *
+   * @default undefined
+   */
+  query?: IndexTextParam;
+  /**
+   * Number of hits for the MCP to return per page.
+   *
+   * @default { exposed: false, default: 7 }
+   */
+  hits_per_page?: IndexNumberParam;
+  /**
+   * The page number the MCP should pull results from.
+   *
+   * @default { exposed: false, default: 0 }
+   */
+  page?: IndexNumberParam;
+  /**
+   * List of attributes that the MCP can retrieve from the index.
+   *
+   * @default { exposed: false, default: ['*'] }
+   */
+  attributesToRetrieve?: IndexStringArrayParam;
+  /**
+   * List of fields that the MCP will return to the Agent.
+   *
+   * @default { exposed: false, default: ["hits", "nbHits", "page", "nbPages", "hitsPerPage", "facets"] }
+   */
+  responseFields?: IndexStringArrayParam;
+  /**
+   * Defined facets the MCP will use when querying the index.
+   *
+   * @default undefined
+   */
+  facets?: IndexFacetParam;
+  /**
+   * Any other custom properties the MCP should send when querying the index.
+   *
+   * @default undefined
+   */
+  custom?: Record<string, unknown>;
+}
+
+export interface AgentStudioIndices {
+  /**
+   * The name of the index used by the search tool.
+   */
+  index: string;
+  /**
+   * A brief description for the search tool.
+   */
+  description: string;
+  /**
+   * A description used to steer the agent on how/when to use the search tool.
+   *
+   * @default ""
+   */
+  enhancedDescription?: string;
+  /**
+   * Default search parameters for the internal (non-MCP) search tool path.
+   *
+   * @default undefined
+   */
+  searchParameters?: SearchParamsObject;
+  /**
+   * Structured search parameters for the MCP-based search tool path.
+   *
+   * Each parameter controls whether it is exposed to the LLM and it's default value.
+   *
+   * @default undefined
+   */
+  searchControls?: AgentStudioSearchControls;
+}
+
+export interface DocSearchAskAi {
+  /**
+   * The index name to use for the Ask AI feature. Your assistant will search for relevant documents.
+   * If not provided, the root index name will be used.
+   */
+  indexName?: string;
+  /**
+   * The API key to use for the ask AI feature. Your assistant will use this API key to search the index.
+   * If not provided, the API key will be used.
+   */
+  apiKey?: string;
+  /**
+   * The app ID to use for the ask AI feature. Your assistant will use this app ID to search the index.
+   * If not provided, the app ID will be used.
+   */
+  appId?: string;
+  /**
+   * The assistant ID to use for the ask AI feature.
+   */
+  assistantId: string;
+  /**
+   * Enables displaying suggested questions on Ask AI's new conversation screen.
+   *
+   * @default false
+   */
+  suggestedQuestions?: boolean;
+  /**
+   * The search parameters to use for the ask AI feature.
+   * Keyed by the index name.
+   *
+   * @example
+   * {
+   *   "INDEX_NAME": { distinct: false }
+   * }
+   */
+  searchParameters?: AgentStudioSearchParameters;
+  /**
+   * List of dynamic indices for the Agent Studio search tool to use.
+   */
+  indices?: AgentStudioIndices[];
+}
+
+export interface Memory {
+  /**
+   * Determines whether or not to display the memory based tool calls.
+   *
+   * @default false
+   */
+  enabled?: boolean;
+  /**
+   * The JWT used by the agent to know which user's memory to read.
+   *
+   * @see https://www.algolia.com/doc/guides/algolia-ai/agent-studio/how-to/user-authentication
+   */
+  userToken?: string;
+}
+
+export interface DocSearchAIProps extends DocSearchProps {
+  /**
+   * Configuration or assistant id to enable ask ai mode. Pass a string assistant id or a full config object.
+   */
+  askAi: DocSearchAskAi | string;
+  /**
+   * Intercept Ask AI requests (e.g. Submitting a prompt or selecting a suggested question).
+   *
+   * Return `true` to prevent the default modal Ask AI flow (no toggle, no sendMessage).
+   * Useful to route Ask AI into a different UI (e.g. `@docsearch/sidepanel-js`) without flicker.
+   */
+  interceptAskAiEvent?: (initialMessage: InitialAskAiMessage) => boolean | void;
+  /**
+   * Use custom tools driven by Agent Studio.
+   *
+   * For best performance, memoize this object with `useMemo` or define it
+   * outside the component. Inline object literals will be recreated every
+   * render but will not affect correctness.
+   **/
+  tools?: ToolCalls;
+  /**
+   * Configuration for the Agent Studio memory feature.
+   */
+  memory?: Memory;
+}
 
 function DocSearchAIComponent(props: DocSearchAIProps, ref: React.ForwardedRef<DocSearchRef>): JSX.Element {
   return (

packages/docsearch-react/src/DocSearchAskAiModal.tsx

@@ -8,7 +8,7 @@ import { AskAiScreenState } from './AskAiScreenState';
 import type { AskAiSearchBoxTranslations } from './components/AskAiSearchBox';
 import { AskAiSearchBox } from './components/AskAiSearchBox';
 import { ModalShell } from './components/ui/ModalShell';
-import type { DocSearchAIProps } from './DocSearch';
+import type { DocSearchAIProps } from './DocSearchAI';
 import type { FooterTranslations } from './Footer';
 import { Footer } from './Footer';
 import { Hit } from './Hit';
@@ -113,7 +113,7 @@ export function DocSearchAskAiModal({
   const searchClient = useSearchClient(appId, apiKey, transformSearchClient);
 
   const askAiConfig = typeof askAi === 'object' ? askAi : null;
-  const askAiConfigurationId = typeof askAi === 'string' ? askAi : askAiConfig?.assistantId || null;
+  const askAiConfigurationId = askAiConfig ? askAiConfig.assistantId : (askAi as string);
   const askAiSearchParameters = askAiConfig?.searchParameters;
   const [askAiState, setAskAiState] = React.useState<AskAiState>('initial');
   const suggestedQuestions = useSuggestedQuestions({
@@ -147,6 +147,7 @@ export function DocSearchAskAiModal({
       searchParameters: askAiSearchParameters,
       tools,
       memory: props.memory,
+      indices: askAiConfig?.indices,
     });
 
   const prevStatus = React.useRef(status);

packages/docsearch-react/src/Sidepanel.tsx

@@ -4,7 +4,7 @@ import type { JSX } from 'react';
 import React from 'react';
 import { createPortal } from 'react-dom';
 
-import type { AgentStudioSearchParameters, Memory } from './DocSearch';
+import type { AgentStudioIndices, AgentStudioSearchParameters, Memory } from './DocSearchAI';
 import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index';
 import { SidepanelButton, Sidepanel } from './Sidepanel/index';
 import type { ToolCalls } from './types/AskiAi';
@@ -73,6 +73,10 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & {
    * Configuration for the Agent Studio memory feature.
    */
   memory?: Memory;
+  /**
+   * List of dynamic indices for the Agent Studio search tool to use.
+   */
+  indices?: AgentStudioIndices[];
 };
 
 type SidepanelProps = DocSearchSidepanelProps & SidepanelSearchParameters;

packages/docsearch-react/src/Sidepanel/Sidepanel.tsx

@@ -150,6 +150,7 @@ function SidepanelInner(
     initialMessage,
     tools = EMPTY_TOOLS,
     memory,
+    indices,
   }: Props,
   ref: React.ForwardedRef<SidepanelRef>,
 ): JSX.Element {
@@ -193,6 +194,7 @@ function SidepanelInner(
     searchParameters,
     tools,
     memory,
+    indices,
   });
 
   const suggestedQuestions = useSuggestedQuestions({

packages/docsearch-react/src/__tests__/api.test.tsx

@@ -5,8 +5,9 @@ import { describe, it, expect, afterEach } from 'vitest';
 import '@testing-library/jest-dom/vitest';
 
 import { DocSearch as DocSearchComponent } from '../DocSearch';
-import type { DocSearchAIProps, DocSearchProps } from '../DocSearch';
+import type { DocSearchProps } from '../DocSearch';
 import { DocSearchAI as DocSearchAIComponent } from '../DocSearchAI';
+import type { DocSearchAIProps } from '../DocSearchAI';
 
 function DocSearch(props: Partial<DocSearchProps>): JSX.Element {
   return <DocSearchComponent appId="woo" apiKey="foo" indexName="bar" {...props} />;

packages/docsearch-react/src/__tests__/useAskAi.test.tsx

@@ -11,7 +11,12 @@ type ToolCall = {
 
 type ChatOptions = {
   onToolCall: (params: { toolCall: ToolCall }) => unknown;
-  transport: { options: { headers?: Record<string, string> } };
+  transport: {
+    options: {
+      headers?: Record<string, string>;
+      body?: Record<string, unknown>;
+    };
+  };
 };
 
 type CustomOnToolCallParams = ToolCall & {
@@ -57,6 +62,14 @@ describe('useAskAi', () => {
     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();
 
@@ -187,4 +200,106 @@ describe('useAskAi', () => {
 
     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: {} });
+  });
 });

packages/docsearch-react/src/useAskAi.ts

@@ -12,18 +12,19 @@ import { type AIMessage, type ToolCalls } from './types/AskiAi';
 import type { OnAskAiFeedback } from './types/Feedback';
 import { EMPTY_TOOLS } from './utils/ai';
 
-import type { AgentStudioSearchParameters, Memory, StoredAskAiState } from '.';
+import type { AgentStudioIndices, AgentStudioSearchParameters, Memory, StoredAskAiState } from '.';
 
 type UseChat = UseChatHelpers<AIMessage>;
 
 type UseAskAiParams = {
-  assistantId?: string | null;
+  assistantId: string;
   apiKey: string;
   appId: string;
   indexName: string;
   searchParameters?: AgentStudioSearchParameters;
   tools: ToolCalls;
   memory?: Memory;
+  indices?: AgentStudioIndices[];
 };
 
 type UseAskAiReturn = {
@@ -44,6 +45,7 @@ type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn;
 type AgentStudioTransportParams = Pick<UseAskAiParams, 'apiKey' | 'appId' | 'assistantId'> & {
   searchParameters?: AgentStudioSearchParameters;
   userToken?: string;
+  indices?: AgentStudioIndices[];
 };
 
 const getAgentStudioTransport = ({
@@ -52,7 +54,21 @@ const getAgentStudioTransport = ({
   assistantId,
   searchParameters,
   userToken,
+  indices,
 }: AgentStudioTransportParams): DefaultChatTransport<AIMessage> => {
+  const algoliaParams: {
+    searchParameters?: AgentStudioSearchParameters;
+    indices?: AgentStudioIndices[];
+  } = {};
+
+  if (searchParameters) {
+    algoliaParams.searchParameters = searchParameters;
+  }
+
+  if (indices && indices.length > 0) {
+    algoliaParams.indices = indices;
+  }
+
   return new DefaultChatTransport({
     api: `${agentStudioBaseUrl(appId)}/agents/${assistantId}/completions?stream=true&compatibilityMode=ai-sdk-5`,
     headers: {
@@ -60,7 +76,7 @@ const getAgentStudioTransport = ({
       'x-algolia-api-key': apiKey,
       ...(userToken ? { 'x-algolia-secure-user-token': userToken } : {}),
     },
-    body: searchParameters ? { algolia: { searchParameters } } : {},
+    body: { algolia: algoliaParams },
   });
 };
 
@@ -72,6 +88,7 @@ export const useAskAi: UseAskAi = ({
   tools = EMPTY_TOOLS,
   searchParameters,
   memory,
+  indices,
 }) => {
   const abortControllerRef = useRef(new AbortController());
 
@@ -80,11 +97,12 @@ export const useAskAi: UseAskAi = ({
       getAgentStudioTransport({
         apiKey,
         appId,
-        assistantId: assistantId ?? '',
+        assistantId,
         searchParameters,
         userToken: memory?.userToken,
+        indices,
       }),
-    [apiKey, appId, assistantId, searchParameters, memory?.userToken],
+    [apiKey, appId, assistantId, searchParameters, memory?.userToken, indices],
   );
 
   // Sync ref during render so the stable `handleToolCall` (registered once