feat: enhance web search tool integration with citation management and UI enhancements

This commit is contained in:
Anish Sarkar 2026-03-30 01:38:36 +05:30
parent 9eab427b56
commit 74826b3714
13 changed files with 133 additions and 209 deletions

View file

@ -12,6 +12,7 @@ import type { FC } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { commentsEnabledAtom, targetCommentIdAtom } from "@/atoms/chat/current-thread.atom";
import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
import { CitationMetadataProvider } from "@/components/assistant-ui/citation-metadata-context";
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
@ -78,7 +79,7 @@ export const MessageError: FC = () => {
const AssistantMessageInner: FC = () => {
return (
<>
<CitationMetadataProvider>
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
<MessagePrimitive.Parts
components={{
@ -116,9 +117,10 @@ const AssistantMessageInner: FC = () => {
create_confluence_page: CreateConfluencePageToolUI,
update_confluence_page: UpdateConfluencePageToolUI,
delete_confluence_page: DeleteConfluencePageToolUI,
link_preview: () => null,
multi_link_preview: () => null,
scrape_webpage: () => null,
web_search: () => null,
link_preview: () => null,
multi_link_preview: () => null,
scrape_webpage: () => null,
},
Fallback: ToolFallback,
},
@ -130,7 +132,7 @@ const AssistantMessageInner: FC = () => {
<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
<AssistantActionBar />
</div>
</>
</CitationMetadataProvider>
);
};

View file

@ -0,0 +1,61 @@
"use client";
import { useAuiState } from "@assistant-ui/react";
import { createContext, type FC, type ReactNode, useContext, useMemo } from "react";
export interface CitationMeta {
title: string;
snippet?: string;
}
type CitationMetadataMap = ReadonlyMap<string, CitationMeta>;
const CitationMetadataContext = createContext<CitationMetadataMap>(new Map());
interface ToolCallResult {
status?: string;
citations?: Record<string, { title: string; snippet?: string }>;
}
interface MessageContent {
type: string;
toolName?: string;
result?: unknown;
}
export const CitationMetadataProvider: FC<{ children: ReactNode }> = ({ children }) => {
const content = useAuiState(({ message }) => (message as { content?: MessageContent[] })?.content);
const metadataMap = useMemo<CitationMetadataMap>(() => {
if (!content || !Array.isArray(content)) return new Map();
const merged = new Map<string, CitationMeta>();
for (const part of content) {
if (part.type !== "tool-call" || part.toolName !== "web_search" || !part.result) {
continue;
}
const result = part.result as ToolCallResult;
const citations = result.citations;
if (!citations || typeof citations !== "object") continue;
for (const [url, meta] of Object.entries(citations)) {
if (url.startsWith("http") && meta.title && !merged.has(url)) {
merged.set(url, { title: meta.title, snippet: meta.snippet });
}
}
}
return merged;
}, [content]);
return (
<CitationMetadataContext.Provider value={metadataMap}>{children}</CitationMetadataContext.Provider>
);
};
export function useCitationMetadata(url: string): CitationMeta | undefined {
const map = useContext(CitationMetadataContext);
return map.get(url);
}

View file

@ -1,9 +1,10 @@
"use client";
import { ExternalLink } from "lucide-react";
import type { FC } from "react";
import { useState } from "react";
import { useCitationMetadata } from "@/components/assistant-ui/citation-metadata-context";
import { SourceDetailPanel } from "@/components/new-chat/source-detail-panel";
import { Citation } from "@/components/tool-ui/citation";
interface InlineCitationProps {
chunkId: number;
@ -55,21 +56,23 @@ interface UrlCitationProps {
/**
* Inline citation for live web search results (URL-based chunk IDs).
* Renders a clickable badge showing the source domain that opens the URL in a new tab.
* Renders a compact chip with favicon + domain and a hover popover showing the
* page title and snippet (extracted deterministically from web_search tool results).
*/
export const UrlCitation: FC<UrlCitationProps> = ({ url }) => {
const domain = extractDomain(url);
const meta = useCitationMetadata(url);
return (
<a
<Citation
id={`url-cite-${url}`}
href={url}
target="_blank"
rel="noopener noreferrer"
className="text-[10px] font-bold bg-primary/80 hover:bg-primary text-primary-foreground rounded-full h-4 px-1.5 inline-flex items-center gap-0.5 align-super cursor-pointer transition-colors ml-0.5 no-underline"
title={url}
>
<ExternalLink className="size-2.5 shrink-0" />
{domain}
</a>
title={meta?.title || domain}
snippet={meta?.snippet}
domain={domain}
favicon={`https://www.google.com/s2/favicons?domain=${domain}&sz=32`}
variant="inline"
type="webpage"
/>
);
};

View file

@ -10,6 +10,7 @@ import {
import { CheckIcon, CopyIcon } from "lucide-react";
import Image from "next/image";
import { type FC, type ReactNode, useState } from "react";
import { CitationMetadataProvider } from "@/components/assistant-ui/citation-metadata-context";
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
@ -142,30 +143,33 @@ const PublicAssistantMessage: FC = () => {
className="aui-assistant-message-root group fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
data-role="assistant"
>
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
<MessagePrimitive.Parts
components={{
Text: MarkdownText,
tools: {
by_name: {
generate_podcast: GeneratePodcastToolUI,
generate_report: GenerateReportToolUI,
generate_video_presentation: GenerateVideoPresentationToolUI,
display_image: GenerateImageToolUI,
generate_image: GenerateImageToolUI,
<CitationMetadataProvider>
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
<MessagePrimitive.Parts
components={{
Text: MarkdownText,
tools: {
by_name: {
generate_podcast: GeneratePodcastToolUI,
generate_report: GenerateReportToolUI,
generate_video_presentation: GenerateVideoPresentationToolUI,
display_image: GenerateImageToolUI,
generate_image: GenerateImageToolUI,
web_search: () => null,
link_preview: () => null,
multi_link_preview: () => null,
scrape_webpage: () => null,
},
Fallback: ToolFallback,
},
Fallback: ToolFallback,
},
}}
/>
</div>
}}
/>
</div>
<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
<PublicAssistantActionBar />
</div>
<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
<PublicAssistantActionBar />
</div>
</CitationMetadataProvider>
</MessagePrimitive.Root>
);
};

View file

@ -1,5 +1,4 @@
import { z } from "zod";
import { defineToolUiContract } from "../shared/contract";
import {
ToolUIIdSchema,
ToolUIReceiptSchema,
@ -37,16 +36,3 @@ export const SerializableCitationSchema = z.object({
});
export type SerializableCitation = z.infer<typeof SerializableCitationSchema>;
const SerializableCitationSchemaContract = defineToolUiContract(
"Citation",
SerializableCitationSchema,
);
export const parseSerializableCitation: (
input: unknown,
) => SerializableCitation = SerializableCitationSchemaContract.parse;
export const safeParseSerializableCitation: (
input: unknown,
) => SerializableCitation | null = SerializableCitationSchemaContract.safeParse;

View file

@ -1,19 +0,0 @@
import { z } from "zod";
import { parseWithSchema, safeParseWithSchema } from "./parse";
export interface ToolUiContract<T> {
schema: z.ZodType<T>;
parse: (input: unknown) => T;
safeParse: (input: unknown) => T | null;
}
export function defineToolUiContract<T>(
componentName: string,
schema: z.ZodType<T>,
): ToolUiContract<T> {
return {
schema,
parse: (input: unknown) => parseWithSchema(schema, input, componentName),
safeParse: (input: unknown) => safeParseWithSchema(schema, input),
};
}

View file

@ -1,27 +0,0 @@
import { z } from "zod";
export const AspectRatioSchema = z
.enum(["auto", "1:1", "4:3", "16:9", "9:16"])
.default("auto");
export type AspectRatio = z.infer<typeof AspectRatioSchema>;
export const MediaFitSchema = z.enum(["cover", "contain"]).default("cover");
export type MediaFit = z.infer<typeof MediaFitSchema>;
export const RATIO_CLASS_MAP: Record<AspectRatio, string> = {
auto: "",
"1:1": "aspect-square",
"4:3": "aspect-[4/3]",
"16:9": "aspect-video",
"9:16": "aspect-[9/16]",
};
export function getRatioClass(ratio: AspectRatio): string {
return RATIO_CLASS_MAP[ratio];
}
export function getFitClass(fit: MediaFit): string {
return fit === "cover" ? "object-cover" : "object-contain";
}

View file

@ -1,30 +0,0 @@
export function formatDuration(durationMs: number): string {
const totalSeconds = Math.round(durationMs / 1000);
const hours = Math.floor(totalSeconds / 3600);
const minutes = Math.floor((totalSeconds % 3600) / 60);
const seconds = totalSeconds % 60;
if (hours > 0) {
return `${hours}:${minutes.toString().padStart(2, "0")}:${seconds
.toString()
.padStart(2, "0")}`;
}
return `${minutes}:${seconds.toString().padStart(2, "0")}`;
}
/**
* Format file size in bytes to human-readable string.
* @example formatFileSize(1024) => "1 KB"
* @example formatFileSize(1536000) => "1.5 MB"
*/
export function formatFileSize(bytes: number): string {
if (bytes < 1024) return `${bytes} B`;
const units = ["KB", "MB", "GB"];
let size = bytes / 1024;
let unit = 0;
while (size >= 1024 && unit < units.length - 1) {
size /= 1024;
unit += 1;
}
return `${size.toFixed(size >= 10 ? 0 : 1)} ${units[unit]}`;
}

View file

@ -1,17 +1,3 @@
export {
AspectRatioSchema,
MediaFitSchema,
RATIO_CLASS_MAP,
getRatioClass,
getFitClass,
type AspectRatio,
type MediaFit,
} from "./aspect-ratio";
export { OVERLAY_GRADIENT } from "./overlay-gradient";
export { formatDuration, formatFileSize } from "./format-utils";
export { sanitizeHref } from "./sanitize-href";
export {
resolveSafeNavigationHref,

View file

@ -1,19 +0,0 @@
export const OVERLAY_GRADIENT = `linear-gradient(
to bottom,
hsl(0, 0%, 0%) 0%,
hsla(0, 0%, 0%, 0.987) 8.3%,
hsla(0, 0%, 0%, 0.951) 16.6%,
hsla(0, 0%, 0%, 0.896) 24.6%,
hsla(0, 0%, 0%, 0.825) 32.5%,
hsla(0, 0%, 0%, 0.741) 40.1%,
hsla(0, 0%, 0%, 0.648) 47.6%,
hsla(0, 0%, 0%, 0.55) 54.8%,
hsla(0, 0%, 0%, 0.45) 61.7%,
hsla(0, 0%, 0%, 0.352) 68.3%,
hsla(0, 0%, 0%, 0.259) 74.5%,
hsla(0, 0%, 0%, 0.175) 80.4%,
hsla(0, 0%, 0%, 0.104) 86%,
hsla(0, 0%, 0%, 0.049) 91.1%,
hsla(0, 0%, 0%, 0.013) 95.8%,
hsla(0, 0%, 0%, 0) 100%
)` as const;

View file

@ -1,51 +0,0 @@
import { z } from "zod";
function formatZodPath(path: Array<string | number | symbol>): string {
if (path.length === 0) return "root";
return path
.map((segment) =>
typeof segment === "number" ? `[${segment}]` : String(segment),
)
.join(".");
}
/**
* Format Zod errors into a compact `path: message` string.
*/
export function formatZodError(error: z.ZodError): string {
const parts = error.issues.map((issue) => {
const path = formatZodPath(issue.path);
return `${path}: ${issue.message}`;
});
return Array.from(new Set(parts)).join("; ");
}
/**
* Parse unknown input and throw a readable error.
*/
export function parseWithSchema<T>(
schema: z.ZodType<T>,
input: unknown,
name: string,
): T {
const res = schema.safeParse(input);
if (!res.success) {
throw new Error(`Invalid ${name} payload: ${formatZodError(res.error)}`);
}
return res.data;
}
/**
* Parse unknown input, returning `null` instead of throwing on failure.
*
* Use this in assistant-ui `render` functions where `args` stream in
* incrementally and may be incomplete until the tool call finishes.
*/
export function safeParseWithSchema<T>(
schema: z.ZodType<T>,
input: unknown,
): T | null {
const res = schema.safeParse(input);
return res.success ? res.data : null;
}