ToolCall.tsx

import type { JSX } from 'react';
import React, { useMemo } from 'react';

import { LoadingIcon, MemoryIcon, SearchIcon, ToolIcon } from '../icons';
import type {
  AIToolPart,
  MemoryToolPart,
  SearchToolPart,
  ToolCalls,
  ToolDefinition,
} from '../types/AskiAi';
import { getSearchToolQueries, isSearchToolPart } from '../utils/ai';

import { ToolState } from './ui/ToolState';

export type ToolCallTranslations = {
  /** Text shown while assistant is preparing tool call. */
  preToolCallText: string;
  /** Text shown while assistant is performing search tool call. */
  searchingText: string;
  /** 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 {
  part: AIToolPart;
  translations: ToolCallTranslations;
  tools: ToolCalls;
  onSearchQueryClick?: (query: string) => void;
  memoryEnabled?: boolean;
}

interface SearchToolProps {
  part: SearchToolPart;
  translations: ToolCallTranslations;
  onSearchQueryClick?: (query: string) => void;
}

function SearchTool({
  part,
  translations,
  onSearchQueryClick,
}: SearchToolProps) {
  const { searchingText, preToolCallText, toolCallResultText } = translations;

  switch (part.state) {
    case 'input-streaming':
      return (
        <ToolState
          shimmer={true}
          icon={
            <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
          }
          variant="PartialCall"
        >
          <span>{searchingText}</span>
        </ToolState>
      );
    case 'input-available': {
      const queries = getSearchToolQueries(part);

      return (
        <>
          {queries.map((q, index) => (
            <ToolState
              // oxlint-disable-next-line react/no-array-index-key
              key={`${part.toolCallId}-call-${index}`}
              shimmer={true}
              icon={
                <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
              }
              variant="Call"
            >
              <span>
                {preToolCallText} {`"${q}" ...`}
              </span>
            </ToolState>
          ))}
        </>
      );
    }
    case 'output-available': {
      const queries = getSearchToolQueries(part);

      return (
        <>
          {queries.map((q, index) => (
            <ToolState
              // oxlint-disable-next-line react/no-array-index-key
              key={`${part.toolCallId}-result-${index}`}
              icon={<SearchIcon />}
              variant="Result"
            >
              <span>
                {toolCallResultText}{' '}
                {onSearchQueryClick ? (
                  <span
                    role="button"
                    tabIndex={0}
                    className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
                    onKeyDown={(e) => {
                      if (e.key === 'Enter' || e.key === ' ') {
                        e.preventDefault();
                        onSearchQueryClick(q);
                      }
                    }}
                    onClick={() => onSearchQueryClick(q)}
                  >
                    {' '}
                    &quot;{q}&quot;
                  </span>
                ) : (
                  <span className="DocSearch-AskAiScreen-MessageContent-Tool-Query">
                    {' '}
                    &quot;{q}&quot;
                  </span>
                )}
              </span>
            </ToolState>
          ))}
        </>
      );
    }
    default:
      return null;
  }
}

interface CustomToolProps {
  tool: ToolDefinition;
  part: AIToolPart;
}

function CustomTool({ tool, part }: CustomToolProps) {
  const toolOutput = useMemo(() => {
    if (part.state !== 'output-available') return null;

    return tool.render({ message: { input: part.input, output: part.output } });
  }, [part.input, part.output, tool, part.state]);

  if (part.state === 'output-error') {
    return null;
  }

  if (part.state !== 'output-available') {
    const { callingToolText = 'Loading tool' } = tool.translations || {};

    return (
      <ToolState shimmer={true} variant="PartialCall" icon={<ToolIcon />}>
        <span>{callingToolText}</span>
      </ToolState>
    );
  }

  if (!toolOutput) return null;

  return (
    <ToolState icon={<ToolIcon />} variant="Result">
      <span>{toolOutput}</span>
    </ToolState>
  );
}

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 (
      <ToolState variant="Result" icon={<MemoryIcon />}>
        {savedMemoryToolResultText}
      </ToolState>
    );
  }

  return (
    <ToolState variant="Result" icon={<MemoryIcon />}>
      {memoryToolResultText}
    </ToolState>
  );
}

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 <CustomTool tool={customTool} part={part} />;
  }

  if (memoryEnabled && isMemoryToolPart(part)) {
    return <MemoryTool part={part} translations={translations} />;
  }

  if (isSearchToolPart(part)) {
    return (
      <SearchTool
        part={part}
        translations={translations}
        onSearchQueryClick={onSearchQueryClick}
      />
    );
  }

  return null;
}