ask-ai-search.tsx

ask-ai-search.tsxapps/web/src/components/search/ask-ai-search.tsxraw ← search14,904 B← Back to the app
1"use client";
2 
3import { useEffect, useRef, useState } from "react";
4import Link from "next/link";
5import ReactMarkdown from "react-markdown";
6import { AlertTriangle, Loader2, MessageSquareText, SendHorizonal, X } from "lucide-react";
7import { askRag, AiPreStreamError, type RagSource } from "@/lib/ai";
8import { cn } from "@/lib/utils";
9 
10/**
11 * TASK-0125 — RAG workspace search UI (shared by the search-page card and the
12 * `/ask` full-page route).
13 *
14 * One question input → the cited, streaming answer:
15 *   - the FIRST SSE frame is `sources` (the server-side citations) — the
16 *     `[n]` chips in the answer + the "Sources" list render from it, so they
17 *     are deterministic and independent of the model's output;
18 *   - the answer streams as markdown; the model's inline `[n]` markers are
19 *     replaced with the matching citation chip (a real link);
20 *   - each chip links to the source page/database (deep link by the workspace
21 *     slug + the source's `pageId`/`databaseId`/`recordId`) and shows the
22 *     source snippet on hover (title tooltip + the Sources row below).
23 *
24 * States: loading skeleton, empty ("No results — try different wording"),
25 * error banner (TASK-0122 semantics: budget / provider / retryable), and the
26 * streamed answer. A11y: the answer is an `aria-live` region; chips are real
27 * links; the form submits on Enter. Mobile: the input + answer stack in the
28 * same column (no fixed widths).
29 */
30 
31interface AskAiSearchProps {
32  /** The workspace ID (the RAG endpoint's path param). */
33  workspaceId: string;
34  /** The workspace slug (citation deep links navigate by slug). */
35  workspaceSlug: string;
36  /**
37   * The search page's committed query — seeded into the input on first mount
38   * (the card inherits the search context) and kept in sync until the user
39   * types or asks.
40   */
41  initialQuery?: string;
42  /** Compact (card) vs full (route) layout. */
43  variant?: "card" | "page";
44  /** Called with the question as soon as the user asks (URL shareability). */
45  onAsk?: (q: string) => void;
46}
47 
48type Status = "idle" | "loading" | "done" | "error";
49 
50/** Build the citation deep link for a source (slug-based, type-aware). */
51function sourceHref(slug: string, s: RagSource): string | null {
52  const enc = (id: string) => encodeURIComponent(id);
53  if (s.type === "page") return s.pageId ? `/w/${enc(slug)}/page/${enc(s.pageId)}` : null;
54  if (s.type === "database") return s.id ? `/w/${enc(slug)}/databases/${enc(s.id)}` : null;
55  if (s.type === "record") return s.databaseId ? `/w/${enc(slug)}/databases/${enc(s.databaseId)}` : null;
56  if (s.type === "block") return s.pageId ? `/w/${enc(slug)}/page/${enc(s.pageId)}` : null;
57  return null;
58}
59 
60/** Map a pre-stream / in-band error to a human banner (TASK-0122 semantics). */
61function errorText(e: unknown): { title: string; hint?: string; retryable?: boolean } {
62  if (e instanceof AiPreStreamError) {
63    if (e.code === "AI_BUDGET_EXCEEDED" || e.status === 402) {
64      return { title: "AI budget exhausted", hint: "Ask your workspace admin to raise the monthly AI limit." };
65    }
66    if (e.code === "RATE_LIMITED" || e.status === 429) {
67      return {
68        title: "AI call limit reached",
69        hint: e.retryAfterSeconds ? `Try again in ${e.retryAfterSeconds}s.` : "Slow down a little and retry.",
70        retryable: true
71      };
72    }
73    if (e.code === "SEARCH_UNAVAILABLE" || e.status === 503) {
74      return {
75        title: "Search is temporarily unavailable",
76        hint: "Your content may not be indexed yet — try again shortly.",
77        retryable: true
78      };
79    }
80    if (e.status === 404) return { title: "Workspace not found" };
81    if (e.status === 403) return { title: "You do not have access to this workspace" };
82  }
83  const msg = e instanceof Error && e.message ? e.message : "Something went wrong.";
84  return { title: "The AI request failed", hint: msg, retryable: true };
85}
86 
87export function AskAiSearch({ workspaceId, workspaceSlug, initialQuery = "", variant = "card", onAsk }: AskAiSearchProps) {
88  const [input, setInput] = useState(initialQuery);
89  const [question, setQuestion] = useState("");
90  const [status, setStatus] = useState<Status>("idle");
91  const [answer, setAnswer] = useState("");
92  const [sources, setSources] = useState<RagSource[]>([]);
93  const [err, setErr] = useState<{ title: string; hint?: string; retryable?: boolean } | null>(null);
94  const abortRef = useRef<AbortController | null>(null);
95 
96  // Keep the input seeded from the search query until the user types/asks.
97  useEffect(() => {
98    if (question === "") setInput(initialQuery);
99  }, [initialQuery, question]);
100 
101  // Abort any in-flight stream on unmount (card removed / route left).
102  useEffect(() => () => abortRef.current?.abort(), []);
103 
104  const ask = (q: string) => {
105    const query = q.trim();
106    if (!query || status === "loading" || !workspaceId) return;
107    abortRef.current?.abort();
108    const ac = new AbortController();
109    abortRef.current = ac;
110    setQuestion(query);
111    setAnswer("");
112    setSources([]);
113    setErr(null);
114    setStatus("loading");
115    onAsk?.(query);
116 
117    askRag({
118      workspaceId,
119      query,
120      signal: ac.signal,
121      onFrame: (frame) => {
122        if (frame.type === "sources") setSources(frame.sources ?? []);
123        else if (frame.type === "delta") setAnswer((prev) => prev + frame.text);
124      }
125    })
126      .then((res) => {
127        setSources((prev) => (res.sources.length > 0 ? res.sources : prev));
128        setAnswer((prev) => (prev || res.answer));
129        if (res.reason === "error") {
130          setErr({ title: res.error?.message ?? "The AI request failed", hint: res.error?.code, retryable: true });
131        }
132        setStatus(res.reason === "error" && !res.answer ? "error" : "done");
133      })
134      .catch((e: unknown) => {
135        if (e instanceof Error && e.name === "AbortError") {
136          setStatus((s) => (s === "loading" ? "idle" : s));
137          return;
138        }
139        setErr(errorText(e));
140        setStatus("error");
141      });
142  };
143 
144  const stop = () => abortRef.current?.abort();
145 
146  const hasAnswer = answer.trim().length > 0;
147 
148  // Render the streamed answer as markdown. The model's inline `[n]`
149  // citation markers (the convention the mock + the RAG system prompt both
150  // use) are rewritten into markdown links to the source's deep link and
151  // rendered as citation chips; the `a` override chip-ifies ONLY the
152  // hrefs we injected (a model-supplied link is rendered normally). The
153  // Sources list below carries the same chips, so a model that skips inline
154  // markers still yields links.
155  const renderAnswer = () => {
156    const injected = new Map<string, RagSource>();
157    const md = answer.replace(/\[(\d+)\]/g, (tok, nStr: string) => {
158      const n = Number(nStr);
159      const s = sources.find((x) => x.n === n);
160      if (!s) return tok;
161      const href = sourceHref(workspaceSlug, s);
162      if (!href) return tok;
163      injected.set(href, s);
164      return `[${n}](${href})`;
165    });
166    return (
167      <div className="text-sm leading-relaxed">
168        <ReactMarkdown
169          components={{
170            p: ({ children }) => <p className="mb-2 last:mb-0">{children}</p>,
171            ul: ({ children }) => <ul className="mb-2 list-disc pl-5 last:mb-0">{children}</ul>,
172            ol: ({ children }) => <ol className="mb-2 list-decimal pl-5 last:mb-0">{children}</ol>,
173            li: ({ children }) => <li className="mb-0.5">{children}</li>,
174            h1: ({ children }) => <h3 className="mb-1 mt-2 font-semibold">{children}</h3>,
175            h2: ({ children }) => <h3 className="mb-1 mt-2 font-semibold">{children}</h3>,
176            h3: ({ children }) => <h4 className="mb-1 mt-2 font-medium">{children}</h4>,
177            blockquote: ({ children }) => (
178              <blockquote className="mb-2 border-l-2 border-border pl-3 text-muted-foreground">{children}</blockquote>
179            ),
180            pre: ({ children }) => <pre className="mb-2 overflow-x-auto rounded bg-muted p-2 text-xs">{children}</pre>,
181            code: ({ children }) => <code className="rounded bg-muted px-1 py-0.5 text-xs">{children}</code>,
182            a: ({ href, children }) => {
183              if (href) {
184                const s = injected.get(href);
185                if (s) {
186                  return (
187                    <Link
188                      href={href}
189                      data-testid={`ask-ai-chip-${s.n}`}
190                      title={`${s.title} — ${s.snippet}`}
191                      aria-label={`Citation ${s.n}: ${s.title}`}
192                      className="mx-0.5 inline-block rounded bg-primary/10 px-1 text-[10px] font-medium text-primary no-underline hover:bg-primary/20"
193                    >
194                      [{s.n}]
195                    </Link>
196                  );
197                }
198                return (
199                  <a href={href} target="_blank" rel="noreferrer" className="text-primary underline">
200                    {children}
201                  </a>
202                );
203              }
204              return <>{children}</>;
205            }
206          }}
207        >
208          {md}
209        </ReactMarkdown>
210      </div>
211    );
212  };
213 
214  return (
215    <div
216      data-testid="ask-ai"
217      className={cn(
218        "w-full",
219        variant === "page" ? "mx-auto max-w-3xl px-6 py-8" : "rounded-xl border border-border bg-card p-4 shadow-sm"
220      )}
221    >
222      <div className={cn("mb-3 flex items-center gap-2", variant === "page" && "mb-5")}>
223        <MessageSquareText className="size-4 text-primary" aria-hidden />
224        <h2 className={cn("font-semibold", variant === "page" ? "text-xl" : "text-sm")}>Ask AI about this workspace</h2>
225      </div>
226 
227      <div className={cn("relative", variant === "page" && "max-w-3xl")}>
228        <input
229          data-testid="ask-ai-input"
230          type="text"
231          value={input}
232          onChange={(e) => setInput(e.target.value)}
233          onKeyDown={(e) => {
234            if (e.key === "Enter" && !e.shiftKey) {
235              e.preventDefault();
236              ask(input);
237            }
238          }}
239          placeholder='Ask AI about this workspace… e.g. "What is the launch code word?"'
240          aria-label="Ask AI about this workspace"
241          disabled={status === "loading"}
242          className="w-full rounded-lg border border-input bg-background py-2.5 pl-3 pr-24 text-sm outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:opacity-60"
243        />
244        <div className="absolute right-1.5 top-1/2 flex -translate-y-1/2 items-center gap-1">
245          {status === "loading" ? (
246            <button
247              type="button"
248              data-testid="ask-ai-stop"
249              onClick={stop}
250              className="inline-flex items-center gap-1 rounded-md border border-border px-2 py-1.5 text-xs hover:bg-accent"
251            >
252              <X className="size-3.5" aria-hidden /> Stop
253            </button>
254          ) : (
255            <button
256              type="button"
257              data-testid="ask-ai-submit"
258              onClick={() => ask(input)}
259              disabled={!input.trim() || !workspaceId}
260              className="inline-flex items-center gap-1 rounded-md bg-primary px-2.5 py-1.5 text-xs font-medium text-primary-foreground hover:opacity-90 disabled:opacity-50"
261            >
262              <SendHorizonal className="size-3.5" aria-hidden /> Ask
263            </button>
264          )}
265        </div>
266      </div>
267 
268      {status === "loading" && (
269        <div
270          className="mt-4 flex items-center gap-2 text-sm text-muted-foreground"
271          role="status"
272          data-testid="ask-ai-loading"
273        >
274          <Loader2 className="size-4 animate-spin" aria-hidden /> Retrieving workspace content and generating an answer…
275        </div>
276      )}
277 
278      {err !== null && status !== "loading" && (
279        <div
280          data-testid="ask-ai-error"
281          role="alert"
282          className="mt-4 flex items-start gap-2 rounded-md border border-destructive/40 bg-destructive/5 px-3 py-2.5 text-sm"
283        >
284          <AlertTriangle className="mt-0.5 size-4 shrink-0 text-destructive" aria-hidden />
285          <div>
286            <p className="font-medium text-foreground">{err.title}</p>
287            {err.hint ? <p className="mt-0.5 text-xs text-muted-foreground">{err.hint}</p> : null}
288            {err.retryable ? <p className="mt-1 text-xs text-muted-foreground">Press Ask to retry.</p> : null}
289          </div>
290        </div>
291      )}
292 
293      {status !== "loading" && err === null && !question && (
294        <p className="mt-4 text-sm text-muted-foreground" data-testid="ask-ai-empty">
295          Ask a question about anything in this workspace — pages, databases, records and comments.
296        </p>
297      )}
298 
299      {status !== "loading" && err === null && question && sources.length === 0 && !hasAnswer && (
300        <div
301          className="mt-4 rounded-md border border-border px-4 py-6 text-center text-sm text-muted-foreground"
302          data-testid="ask-ai-empty"
303        >
304          No results — try different wording.
305        </div>
306      )}
307 
308      {hasAnswer && (
309        <div
310          data-testid="ask-ai-answer"
311          aria-live="polite"
312          className="mt-4 rounded-lg border border-border bg-card p-4"
313        >
314          {renderAnswer()}
315          {status === "loading" ? (
316            <span className="ml-0.5 inline-block h-3 w-1.5 animate-pulse bg-foreground/60 align-middle" aria-hidden />
317          ) : null}
318        </div>
319      )}
320 
321      {sources.length > 0 && (
322        <div className="mt-4" data-testid="ask-ai-sources">
323          <h3 className="mb-1.5 flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
324            <MessageSquareText className="size-3.5" aria-hidden /> Sources
325          </h3>
326          <ul className="space-y-1.5">
327            {sources.map((s) => {
328              const href = sourceHref(workspaceSlug, s);
329              const row = (
330                <>
331                  <span className="rounded bg-primary/10 px-1 text-[10px] font-medium text-primary">[{s.n}]</span>
332                  <span className="font-medium text-foreground">{s.title || "(untitled)"}</span>
333                  <span className="text-xs text-muted-foreground">· {s.type}</span>
334                  {s.snippet ? <span className="ml-1 truncate text-xs text-muted-foreground">— {s.snippet}</span> : null}
335                </>
336              );
337              return (
338                <li key={`${s.docId}-${s.n}`} data-testid={`ask-ai-source-${s.n}`} className="text-sm">
339                  {href ? (
340                    <Link href={href} className="hover:underline">
341                      {row}
342                    </Link>
343                  ) : (
344                    <span>{row}</span>
345                  )}
346                </li>
347              );
348            })}
349          </ul>
350        </div>
351      )}
352    </div>
353  );
354}
355