diff --git a/packages/docsearch-react/src/AskAiScreen.tsx b/packages/docsearch-react/src/AskAiScreen.tsx index 9d519116..83e1831f 100644 --- a/packages/docsearch-react/src/AskAiScreen.tsx +++ b/packages/docsearch-react/src/AskAiScreen.tsx @@ -3,7 +3,7 @@ import React, { type JSX, useMemo, useState, useEffect } from 'react'; import { AggregatedSearchBlock } from './AggregatedSearchBlock'; import type { AskAiScreenStateProps } from './AskAiScreenState'; -import { ToolCall } from './components/ui/ToolCall'; +import { ToolCall, type ToolCallTranslations } from './components/ToolCall'; import { AlertIcon, LoadingIcon } from './icons'; import { MemoizedMarkdown } from './MemoizedMarkdown'; import type { StoredSearchPlugin } from './stored-searches'; @@ -12,56 +12,88 @@ import { type AIMessage, type ToolCalls } from './types/AskiAi'; import { extractLinksFromMessage, getMessageContent, isThreadDepthError, isAIToolPart } from './utils/ai'; import { groupConsecutiveToolResults } from './utils/groupConsecutiveToolResults'; -export type AskAiScreenTranslations = Partial<{ - // Misc texts - disclaimerText: string; - relatedSourcesText: string; - thinkingText: string; - copyButtonText: string; - copyButtonCopiedText: string; - // Feedback buttons - copyButtonTitle: string; - likeButtonTitle: string; - dislikeButtonTitle: string; - thanksForFeedbackText: string; - // Tool call texts - preToolCallText: string; - duringToolCallText: string; - afterToolCallText: string; - /** - * Build the full jsx element for the aggregated search block. - * If provided, completely overrides the default english renderer. - */ - aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode; +export type AskAiScreenTranslations = Partial< + // Inherit the shared tool-call translations, but expose the search-related + // keys under AskAiScreen's own public names (see mapping below). + Omit & { + // Misc texts + disclaimerText: string; + relatedSourcesText: string; + thinkingText: string; + copyButtonText: string; + copyButtonCopiedText: string; + // Feedback buttons + copyButtonTitle: string; + likeButtonTitle: string; + dislikeButtonTitle: string; + thanksForFeedbackText: string; + // Tool call texts + /** + * Text shown while assistant is performing search tool call. + * Maps to `ToolCallTranslations.searchingText`. + */ + duringToolCallText: string; + /** + * Text shown while assistant is finished performing tool call. + * Maps to `ToolCallTranslations.toolCallResultText`. + */ + afterToolCallText: string; + /** + * Build the full jsx element for the aggregated search block. + * If provided, completely overrides the default english renderer. + */ + aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode; + /** + * Generate the list connective parts only (backwards compatibility). + * Receives full list of queries and should return translation parts for before/after/separators. + * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }). + */ + aggregatedToolCallText?: (queries: string[]) => { + before?: string; + separator?: string; + lastSeparator?: string; + after?: string; + }; + /** + * Message that's shown when user has stopped the streaming of a message. + */ + stoppedStreamingText: string; + /** + * Error title shown if there is an error while chatting. + */ + errorTitleText: string; + /** + * Message shown when thread depth limit is exceeded (AI-217 error). + */ + threadDepthExceededMessage: string; + /** + * Button text for starting a new conversation after thread depth error. + */ + startNewConversationButtonText: string; + } +>; - /** - * Generate the list connective parts only (backwards compatibility). - * Receives full list of queries and should return translation parts for before/after/separators. - * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }). - */ - aggregatedToolCallText?: (queries: string[]) => { - before?: string; - separator?: string; - lastSeparator?: string; - after?: string; +/** + * Maps AskAiScreen's public translation keys to the shared `ToolCallTranslations` + * shape consumed by the `ToolCall` component, applying default English values. + */ +function toToolCallTranslations(translations: AskAiScreenTranslations): ToolCallTranslations { + const { + preToolCallText = 'Searching...', + duringToolCallText = 'Searching...', + afterToolCallText = 'Searched for', + savedMemoryToolResultText = 'Saved to memory', + memoryToolResultText = 'Used memory to enhance results', + } = translations; + + return { + preToolCallText, + searchingText: duringToolCallText, + toolCallResultText: afterToolCallText, + savedMemoryToolResultText, + memoryToolResultText, }; - /** - * Message that's shown when user has stopped the streaming of a message. - */ - stoppedStreamingText: string; - /** - * Error title shown if there is an error while chatting. - */ - errorTitleText: string; - /** - * Message shown when thread depth limit is exceeded (AI-217 error). - */ - threadDepthExceededMessage: string; - /** - * Button text for starting a new conversation after thread depth error. - */ - startNewConversationButtonText: string; -}>; +} type AskAiScreenProps = Omit, 'translations'> & { messages: AIMessage[]; @@ -71,6 +103,7 @@ type AskAiScreenProps = Omit, 'trans translations?: AskAiScreenTranslations; onNewConversation: () => void; agentStudio?: boolean; + memoryEnabled?: boolean; }; interface AskAiScreenHeaderProps { @@ -98,6 +131,7 @@ interface AskAiExchangeCardProps { conversations: StoredSearchPlugin; onFeedback?: (messageId: string, thumbs: 0 | 1) => Promise; agentStudio?: boolean; + memoryEnabled?: boolean; } function AskAiExchangeCard({ @@ -111,16 +145,13 @@ function AskAiExchangeCard({ conversations, onFeedback, agentStudio, + memoryEnabled, }: AskAiExchangeCardProps): JSX.Element { const { userMessage, assistantMessage } = exchange; - const { - stoppedStreamingText = 'You stopped this response', - errorTitleText = 'Chat error', - preToolCallText = 'Searching...', - afterToolCallText = 'Searched for', - duringToolCallText = 'Searching...', - } = translations; + const { stoppedStreamingText = 'You stopped this response', errorTitleText = 'Chat error' } = translations; + + const toolCallTranslations = useMemo(() => toToolCallTranslations(translations), [translations]); const isThreadDepth = isThreadDepthError(askAiError); @@ -191,13 +222,10 @@ function AskAiExchangeCard({ return ( ); @@ -377,7 +405,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps): startNewConversationButtonText = 'Start a new conversation', } = translations; - const { messages, tools, askAiError, status, agentStudio } = props; + const { messages, tools, askAiError, status, agentStudio, memoryEnabled } = props; // Check if there's a thread depth error const hasThreadDepthError = useMemo(() => { @@ -454,6 +482,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps): tools={tools} conversations={props.conversations} agentStudio={agentStudio} + memoryEnabled={memoryEnabled} onSearchQueryClick={handleSearchQueryClick} onFeedback={props.onFeedback} /> diff --git a/packages/docsearch-react/src/AskAiScreenState.tsx b/packages/docsearch-react/src/AskAiScreenState.tsx index eac3ff79..0d07fe6f 100644 --- a/packages/docsearch-react/src/AskAiScreenState.tsx +++ b/packages/docsearch-react/src/AskAiScreenState.tsx @@ -58,6 +58,7 @@ export interface AskAiScreenStateProps selectSuggestedQuestion: (question: SuggestedQuestionHit) => void; onNewConversation: () => void; agentStudio?: boolean; + memoryEnabled?: boolean; } export const AskAiScreenState = React.memo( @@ -86,6 +87,7 @@ export const AskAiScreenState = React.memo( askAiError={props.askAiError} translations={translations?.askAiScreen} agentStudio={props.agentStudio} + memoryEnabled={props.memoryEnabled} /> ); } diff --git a/packages/docsearch-react/src/DocSearch.tsx b/packages/docsearch-react/src/DocSearch.tsx index 97b62530..8482eff9 100644 --- a/packages/docsearch-react/src/DocSearch.tsx +++ b/packages/docsearch-react/src/DocSearch.tsx @@ -234,6 +234,21 @@ 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. @@ -254,6 +269,10 @@ export interface DocSearchAIProps extends DocSearchProps { * render but will not affect correctness. **/ tools?: ToolCalls; + /** + * Configuration for the Agent Studio memory feature. + */ + memory?: Memory; } function DocSearchComponent(props: DocSearchProps, ref: React.ForwardedRef): JSX.Element { diff --git a/packages/docsearch-react/src/DocSearchAskAiModal.tsx b/packages/docsearch-react/src/DocSearchAskAiModal.tsx index 70792bca..771876d4 100644 --- a/packages/docsearch-react/src/DocSearchAskAiModal.tsx +++ b/packages/docsearch-react/src/DocSearchAskAiModal.tsx @@ -117,6 +117,7 @@ export function DocSearchAskAiModal({ suggestedQuestionsEnabled: askAiConfig?.suggestedQuestions, }); const agentStudio = askAiConfig?.agentStudio ?? false; + const memoryEnabled = props.memory?.enabled ?? false; const indexes = normalizeDocSearchIndexes({ indexName, @@ -143,6 +144,7 @@ export function DocSearchAskAiModal({ useStagingEnv: askAiUseStagingEnv, agentStudio, tools, + memory: props.memory, }); const prevStatus = React.useRef(status); @@ -464,6 +466,7 @@ export function DocSearchAskAiModal({ suggestedQuestions={suggestedQuestions} selectSuggestedQuestion={selectSuggestedQuestion} agentStudio={agentStudio} + memoryEnabled={memoryEnabled} onAskAiToggle={onAskAiToggle} onNewConversation={handleNewConversation} onItemClick={(item, event) => { diff --git a/packages/docsearch-react/src/Sidepanel.tsx b/packages/docsearch-react/src/Sidepanel.tsx index 247be8f9..34c0cb0d 100644 --- a/packages/docsearch-react/src/Sidepanel.tsx +++ b/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, AskAiSearchParameters } from './DocSearch'; +import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory } from './DocSearch'; import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index'; import { SidepanelButton, Sidepanel } from './Sidepanel/index'; import type { ToolCalls } from './types/AskiAi'; @@ -93,6 +93,10 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & { * render but will not affect correctness. **/ tools?: ToolCalls; + /** + * Configuration for the Agent Studio memory feature. + */ + memory?: Memory; }; type SidepanelProps = DocSearchSidepanelProps & SidepanelSearchParameters; diff --git a/packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx b/packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx index af54bb0f..bff280be 100644 --- a/packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx +++ b/packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx @@ -3,7 +3,7 @@ import type { JSX } from 'react'; import React, { memo, useMemo } from 'react'; import { AskAiSourcesPanel, type Exchange } from '../AskAiScreen'; -import { ToolCall, type ToolCallTranslations } from '../components/ui/ToolCall'; +import { ToolCall, type ToolCallTranslations } from '../components/ToolCall'; import { AlertIcon, LoadingIcon } from '../icons'; import { MemoizedMarkdown } from '../MemoizedMarkdown'; import type { StoredSearchPlugin } from '../stored-searches'; @@ -72,6 +72,7 @@ export type ConversationScreenProps = { handleFeedback?: (messageId: string, thumbs: 0 | 1) => Promise; streamError?: Error; agentStudio?: boolean; + memoryEnabled?: boolean; tools?: ToolCalls; }; @@ -84,12 +85,24 @@ type ConversationnExchangeProps = { onFeedback?: ConversationScreenProps['handleFeedback']; streamError?: ConversationScreenProps['streamError']; agentStudio?: boolean; + memoryEnabled?: boolean; tools: ToolCalls; }; const ConversationExchange = React.forwardRef( ( - { exchange, translations = {}, isLastExchange, conversations, onFeedback, status, streamError, agentStudio, tools }, + { + exchange, + translations = {}, + isLastExchange, + conversations, + onFeedback, + status, + streamError, + agentStudio, + memoryEnabled, + tools, + }, conversationRef, ): JSX.Element => { const { userMessage, assistantMessage } = exchange; @@ -105,6 +118,8 @@ const ConversationExchange = React.forwardRef getMessageContent(assistantMessage), [assistantMessage]); @@ -172,8 +187,11 @@ const ConversationExchange = React.forwardRef ); } diff --git a/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx b/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx index dad30e74..3d3694bd 100644 --- a/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx +++ b/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx @@ -153,6 +153,7 @@ function SidepanelInner( useStagingEnv = false, agentStudio = false, tools = EMPTY_TOOLS, + memory, }: Props, ref: React.ForwardedRef, ): JSX.Element { @@ -197,6 +198,7 @@ function SidepanelInner( useStagingEnv, agentStudio, tools, + memory, }); const suggestedQuestions = useSuggestedQuestions({ @@ -400,6 +402,7 @@ function SidepanelInner( translations={translations.conversationScreen} streamError={askAiError} agentStudio={agentStudio} + memoryEnabled={memory?.enabled ?? false} tools={tools} /> )} diff --git a/packages/docsearch-react/src/__tests__/ToolCall.test.tsx b/packages/docsearch-react/src/__tests__/ToolCall.test.tsx index 83ff2fb9..7cfea76d 100644 --- a/packages/docsearch-react/src/__tests__/ToolCall.test.tsx +++ b/packages/docsearch-react/src/__tests__/ToolCall.test.tsx @@ -1,15 +1,17 @@ -import { render, screen } from '@testing-library/react'; +import { render, screen, within } from '@testing-library/react'; import '@testing-library/jest-dom/vitest'; import React from 'react'; import { describe, it, expect, vi } from 'vitest'; -import { ToolCall, type ToolCallTranslations } from '../components/ui/ToolCall'; -import type { AIToolPart, ToolCalls } from '../types/AskiAi'; +import { ToolCall, type ToolCallTranslations } from '../components/ToolCall'; +import type { AIToolPart, MemoryToolPart, ToolCalls } from '../types/AskiAi'; const TRANSLATIONS: ToolCallTranslations = { preToolCallText: 'Searching for', searchingText: 'Searching...', toolCallResultText: 'Searched', + savedMemoryToolResultText: 'Saved to memory', + memoryToolResultText: 'Used memory to enhance results', }; describe('ToolCall', () => { @@ -75,6 +77,172 @@ describe('ToolCall', () => { }); }); + describe('memory tools', () => { + it.each([ + { + description: 'tool-algolia_ponder', + part: { + type: 'tool-algolia_ponder', + toolCallId: 'memory-ponder', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + expectedMessage: 'Used memory to enhance results', + }, + { + description: 'tool-algolia_memorize', + part: { + type: 'tool-algolia_memorize', + toolCallId: 'memory-memorize', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + expectedMessage: 'Saved to memory', + }, + { + description: 'tool-algolia_memory_search', + part: { + type: 'tool-algolia_memory_search', + toolCallId: 'memory-search', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + expectedMessage: 'Used memory to enhance results', + }, + ] satisfies Array<{ + description: string; + part: MemoryToolPart; + expectedMessage: string; + }>)( + 'renders the memory result text for $description when output is available and memory is enabled', + ({ part, expectedMessage }) => { + const { container } = render( + , + ); + + expect(within(container).getByText(expectedMessage)).toBeInTheDocument(); + }, + ); + + it.each([ + { + description: 'input-streaming', + part: { + type: 'tool-algolia_ponder', + toolCallId: 'memory-input-streaming', + state: 'input-streaming', + input: { value: 'remember this' }, + }, + expectedMessage: 'Used memory to enhance results', + }, + { + description: 'input-available', + part: { + type: 'tool-algolia_memorize', + toolCallId: 'memory-input-available', + state: 'input-available', + input: { value: 'remember this' }, + }, + expectedMessage: 'Saved to memory', + }, + { + description: 'output-available', + part: { + type: 'tool-algolia_memory_search', + toolCallId: 'memory-output-available', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + expectedMessage: 'Used memory to enhance results', + }, + ] satisfies Array<{ + description: string; + part: MemoryToolPart; + expectedMessage: string; + }>)('renders the memory result text in $description state when memory is enabled', ({ part, expectedMessage }) => { + const { container } = render( + , + ); + + expect(within(container).getByText(expectedMessage)).toBeInTheDocument(); + }); + + it('renders nothing for memory tool output errors when memory is enabled', () => { + const part: MemoryToolPart = { + type: 'tool-algolia_memorize', + toolCallId: 'memory-error', + state: 'output-error', + input: { value: 'remember this' }, + errorText: 'Memory tool failed', + }; + + const { container } = render( + , + ); + + expect(container).toBeEmptyDOMElement(); + }); + + it.each([ + { + description: 'memoryEnabled is false', + part: { + type: 'tool-algolia_ponder', + toolCallId: 'memory-disabled', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + memoryEnabled: false, + }, + { + description: 'memoryEnabled is omitted (defaults to false)', + part: { + type: 'tool-algolia_memory_search', + toolCallId: 'memory-default-off', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }, + memoryEnabled: undefined, + }, + ] satisfies Array<{ + description: string; + part: MemoryToolPart; + memoryEnabled: boolean | undefined; + }>)('renders nothing for memory tools when $description', ({ part, memoryEnabled }) => { + const { container } = render( + , + ); + + expect(container).toBeEmptyDOMElement(); + }); + + it('prioritizes a matching custom tool over the memory tool rendering', () => { + const part: MemoryToolPart = { + type: 'tool-algolia_ponder', + toolCallId: 'memory-custom', + state: 'output-available', + input: { value: 'remember this' }, + output: { stored: true }, + }; + const tools: ToolCalls = { + algolia_ponder: { + render: () => 'Custom memory render', + }, + }; + + const { container } = render(); + + expect(within(container).getByText('Custom memory render')).toBeInTheDocument(); + expect(within(container).queryByText('Used memory to enhance results')).not.toBeInTheDocument(); + }); + }); + describe('custom tools', () => { it('renders the custom loading text while output is unavailable', () => { const part: AIToolPart = { diff --git a/packages/docsearch-react/src/__tests__/useAskAi.test.tsx b/packages/docsearch-react/src/__tests__/useAskAi.test.tsx index 67f39784..a3549a69 100644 --- a/packages/docsearch-react/src/__tests__/useAskAi.test.tsx +++ b/packages/docsearch-react/src/__tests__/useAskAi.test.tsx @@ -11,6 +11,7 @@ type ToolCall = { type ChatOptions = { onToolCall: (params: { toolCall: ToolCall }) => unknown; + transport: { options: { headers?: Record } }; }; type CustomOnToolCallParams = ToolCall & { @@ -48,6 +49,14 @@ describe('useAskAi', () => { return chatOptions.onToolCall; } + function getTransportHeaders(): Record { + if (!chatOptions) { + throw new Error('useChat was not initialized'); + } + + return chatOptions.transport.options.headers ?? {}; + } + beforeEach(() => { vi.clearAllMocks(); @@ -149,4 +158,37 @@ describe('useAskAi', () => { expect(onToolCall).toHaveBeenCalledTimes(1); expect(result).toBeUndefined(); }); + + it('sends the secure user token header when memory.userToken is provided in Agent Studio mode', () => { + renderHook(() => + useAskAi({ + agentStudio: true, + 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({ + agentStudio: true, + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token'); + }); }); diff --git a/packages/docsearch-react/src/components/ui/ToolCall.tsx b/packages/docsearch-react/src/components/ToolCall.tsx similarity index 68% rename from packages/docsearch-react/src/components/ui/ToolCall.tsx rename to packages/docsearch-react/src/components/ToolCall.tsx index 57f4e945..4ee082cc 100644 --- a/packages/docsearch-react/src/components/ui/ToolCall.tsx +++ b/packages/docsearch-react/src/components/ToolCall.tsx @@ -1,8 +1,10 @@ import type { JSX } from 'react'; import React, { useMemo } from 'react'; -import { LoadingIcon, SearchIcon, ToolIcon } from '../../icons'; -import type { AIToolPart, SearchToolPart, ToolCalls, ToolDefinition } from '../../types/AskiAi'; +import { LoadingIcon, MemoryIcon, SearchIcon, ToolIcon } from '../icons'; +import type { AIToolPart, MemoryToolPart, SearchToolPart, ToolCalls, ToolDefinition } from '../types/AskiAi'; + +import { ToolState } from './ui/ToolState'; export type ToolCallTranslations = { /** @@ -17,6 +19,14 @@ export type ToolCallTranslations = { * Text shown while assistant is finished performing tool call. */ toolCallResultText: string; + /** + * Text shown when the agent saved related information to memory. + */ + savedMemoryToolResultText: string; + /** + * Text shown when the agent used the memory tool to enhance results. + */ + memoryToolResultText: string; }; interface ToolCallProps { @@ -24,24 +34,7 @@ interface ToolCallProps { translations: ToolCallTranslations; tools: ToolCalls; onSearchQueryClick?: (query: string) => void; -} - -interface ToolStateProps { - variant: 'Call' | 'PartialCall' | 'Result'; - icon: React.ReactNode; - shimmer?: boolean; - children: React.ReactNode; -} - -function ToolState({ icon, shimmer = false, children, variant }: ToolStateProps) { - const className = `DocSearch-AskAiScreen-MessageContent-Tool Tool--${variant}${shimmer ? ' shimmer' : ''}`; - - return ( -
- {icon} - {children} -
- ); + memoryEnabled?: boolean; } interface SearchToolProps { @@ -148,21 +141,59 @@ function CustomTool({ tool, part }: CustomToolProps) { ); } +function MemoryTool({ part, translations }: { part: MemoryToolPart; translations: ToolCallTranslations }) { + const { savedMemoryToolResultText, memoryToolResultText } = translations; + + if (part.state === 'output-error') return null; + + if (part.type === 'tool-algolia_memorize') { + return ( + }> + {savedMemoryToolResultText} + + ); + } + + return ( + }> + {memoryToolResultText} + + ); +} + function isSearchToolPart(part: AIToolPart): part is SearchToolPart { return part.type === 'tool-searchIndex' || part.type.startsWith('tool-algolia_search_index'); } -export function ToolCall({ part, translations, tools, onSearchQueryClick }: ToolCallProps): JSX.Element | null { - const normalizedToolName = part.type.replace('tool-', ''); - const dynamicTool = tools[normalizedToolName]; - - if (dynamicTool) { - return ; - } - - if (!isSearchToolPart(part)) { - return null; - } - - return ; +function isMemoryToolPart(part: AIToolPart): part is MemoryToolPart { + return ( + part.type === 'tool-algolia_ponder' || + part.type === 'tool-algolia_memorize' || + part.type === 'tool-algolia_memory_search' + ); +} + +export function ToolCall({ + part, + translations, + tools, + onSearchQueryClick, + memoryEnabled = false, +}: ToolCallProps): JSX.Element | null { + const normalizedToolName = part.type.replace('tool-', ''); + const customTool = tools[normalizedToolName]; + + if (customTool) { + return ; + } + + if (memoryEnabled && isMemoryToolPart(part)) { + return ; + } + + if (isSearchToolPart(part)) { + return ; + } + + return null; } diff --git a/packages/docsearch-react/src/components/ui/ToolState.tsx b/packages/docsearch-react/src/components/ui/ToolState.tsx new file mode 100644 index 00000000..b7d3f8b1 --- /dev/null +++ b/packages/docsearch-react/src/components/ui/ToolState.tsx @@ -0,0 +1,19 @@ +import React from 'react'; + +export interface ToolStateProps { + variant: 'Call' | 'PartialCall' | 'Result'; + icon: React.ReactNode; + shimmer?: boolean; + children: React.ReactNode; +} + +export function ToolState({ icon, shimmer = false, children, variant }: ToolStateProps): React.JSX.Element { + const className = `DocSearch-AskAiScreen-MessageContent-Tool Tool--${variant}${shimmer ? ' shimmer' : ''}`; + + return ( +
+ {icon} + {children} +
+ ); +} diff --git a/packages/docsearch-react/src/icons/MemoryIcon.tsx b/packages/docsearch-react/src/icons/MemoryIcon.tsx new file mode 100644 index 00000000..b2f36778 --- /dev/null +++ b/packages/docsearch-react/src/icons/MemoryIcon.tsx @@ -0,0 +1,20 @@ +import React, { type JSX } from 'react'; + +export function MemoryIcon(): JSX.Element { + return ( + + + + + ); +} diff --git a/packages/docsearch-react/src/icons/index.ts b/packages/docsearch-react/src/icons/index.ts index 2819a437..675749f1 100644 --- a/packages/docsearch-react/src/icons/index.ts +++ b/packages/docsearch-react/src/icons/index.ts @@ -19,3 +19,4 @@ export * from './ExpandIcon'; export * from './FolderIcon'; export * from './EditIcon'; export * from './ToolIcon'; +export * from './MemoryIcon'; diff --git a/packages/docsearch-react/src/types/AskiAi.ts b/packages/docsearch-react/src/types/AskiAi.ts index c1771d53..8dbfe18c 100644 --- a/packages/docsearch-react/src/types/AskiAi.ts +++ b/packages/docsearch-react/src/types/AskiAi.ts @@ -32,6 +32,11 @@ export interface AlgoliaMCPSearchTool { }; } +export interface MemoryTool { + input: unknown; + output: unknown; +} + export type ToolDefinition = { /** * Use the tool's input and output to build a string output for the tool result. @@ -57,6 +62,12 @@ export type ToolDefinition = { export type ToolCalls = Record; +type AgentStudioMemoryTools = { + algolia_ponder: MemoryTool; + algolia_memorize: MemoryTool; + algolia_memory_search: MemoryTool; +}; + type SearchTools = { [K in `algolia_search_index_${string}`]: AlgoliaMCPSearchTool; } & { @@ -66,13 +77,14 @@ type SearchTools = { type CustomTools = { [K in string as K extends keyof SearchTools | `algolia_search_index_${string}` ? never : K]: CustomTool; }; -type Tools = CustomTools & SearchTools; +type Tools = AgentStudioMemoryTools & CustomTools & SearchTools; export type AIMessage = UIMessage<{ stopped?: boolean }, UIDataTypes, Tools>; export type AIMessagePart = UIMessagePart; export type SearchToolPart = ToolUIPart; +export type MemoryToolPart = ToolUIPart; export type AIToolPart = ToolUIPart; export interface AggregatedToolCallPart { diff --git a/packages/docsearch-react/src/useAskAi.ts b/packages/docsearch-react/src/useAskAi.ts index 3b0d0b87..b581e16b 100644 --- a/packages/docsearch-react/src/useAskAi.ts +++ b/packages/docsearch-react/src/useAskAi.ts @@ -18,7 +18,7 @@ import { createStoredConversations } from './stored-searches'; import { type AIMessage, type ToolCalls } from './types/AskiAi'; import { EMPTY_TOOLS } from './utils/ai'; -import type { AgentStudioSearchParameters, AskAiSearchParameters, StoredAskAiState } from '.'; +import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory, StoredAskAiState } from '.'; type UseChat = UseChatHelpers; @@ -31,6 +31,7 @@ type UseAskAiParams = { searchParameters?: AskAiSearchParameters; tools: ToolCalls; agentStudio: boolean; + memory?: Memory; } & ( | { agentStudio: false; @@ -59,6 +60,7 @@ type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn; type AgentStudioTransportParams = Pick & { searchParameters?: AgentStudioSearchParameters; + userToken?: string; }; const getAgentStudioTransport = ({ @@ -66,12 +68,14 @@ const getAgentStudioTransport = ({ apiKey, assistantId, searchParameters, + userToken, }: AgentStudioTransportParams): DefaultChatTransport => { return new DefaultChatTransport({ api: `${agentStudioBaseUrl(appId)}/agents/${assistantId}/completions?stream=true&compatibilityMode=ai-sdk-5`, headers: { 'x-algolia-application-id': appId, 'x-algolia-api-key': apiKey, + ...(userToken ? { 'x-algolia-secure-user-token': userToken } : {}), }, body: searchParameters ? { algolia: { searchParameters } } : {}, }); @@ -123,6 +127,7 @@ export const useAskAi: UseAskAi = ({ tools = EMPTY_TOOLS, agentStudio, searchParameters, + memory, }) => { const abortControllerRef = useRef(new AbortController()); @@ -134,6 +139,7 @@ export const useAskAi: UseAskAi = ({ appId, assistantId: assistantId ?? '', searchParameters, + userToken: memory?.userToken, }) : getAskAiTransport({ assistantId: assistantId ?? '', @@ -144,7 +150,7 @@ export const useAskAi: UseAskAi = ({ abortController: abortControllerRef.current, useStagingEnv, }), - [apiKey, appId, assistantId, indexName, useStagingEnv, agentStudio, searchParameters], + [apiKey, appId, assistantId, indexName, useStagingEnv, agentStudio, searchParameters, memory?.userToken], ); // Sync ref during render so the stable `handleToolCall` (registered once