Search in conversation

"use client";

import { useMemo, useState } from "react";
import { useAuiState } from "@assistant-ui/react";
import {
  ConversationSearch,
  type SearchHit,
} from "@/components/assistant-ui/elements/conversation-search";

const CONTEXT = 24;

export function ThreadSearchBar() {
  const messages = useAuiState((s) => s.thread.messages);
  const [query, setQuery] = useState("");
  const [activeIndex, setActiveIndex] = useState(0);

  const hits = useMemo<SearchHit[]>(() => {
    if (query === "") return [];
    const needle = query.toLowerCase();
    const found: SearchHit[] = [];

    messages.forEach((message, messageIndex) => {
      const text = message.parts
        .filter((part) => part.type === "text")
        .map((part) => part.text)
        .join(" ");
      const haystack = text.toLowerCase();

      let from = 0;
      let at: number;
      while ((at = haystack.indexOf(needle, from)) !== -1) {
        found.push({
          id: `${message.id}-${at}`,
          before: text.slice(Math.max(0, at - CONTEXT), at),
          match: text.slice(at, at + query.length),
          after: text.slice(at + query.length, at + query.length + CONTEXT),
          position: (messageIndex / Math.max(1, messages.length - 1)) * 100,
        });
        from = at + query.length;
      }
    });

    return found;
  }, [messages, query]);

  return (
    <ConversationSearch
      query={query}
      hits={hits}
      activeIndex={activeIndex}
      onQueryChange={(next) => {
        setQuery(next);
        setActiveIndex(0);
      }}
      onStep={(delta) => {
        if (hits.length === 0) return;
        const next = (activeIndex + delta + hits.length) % hits.length;
        setActiveIndex(next);
        const [, messageId] = hits[next]!.id.match(/^(.+)-\d+$/) ?? [];
        document
          .querySelector(`[data-message-id="${messageId}"]`)
          ?.scrollIntoView({ block: "center", behavior: "smooth" });
      }}
    />
  );
}