diff --git a/apps/web/components/chat/message/agent-message.tsx b/apps/web/components/chat/message/agent-message.tsx index eac0f7b9..7e0df6bc 100644 --- a/apps/web/components/chat/message/agent-message.tsx +++ b/apps/web/components/chat/message/agent-message.tsx @@ -23,6 +23,7 @@ import { isWebSearchToolName } from "@/lib/chat-web-search-tools" import { buildCitationIndex, fetchDocumentsByIds, + getCitationDisplay, getDocumentSourceUrl, isMemoryToolOutputReady, mapDocumentsByKnownIds, @@ -190,37 +191,64 @@ function CitationLink({ ) } -function sourceTitle( +function documentForCitationTarget( target: CitationTarget, - document?: DocumentWithMemories, -): string { + documentByKnownId: Map, +): DocumentWithMemories | undefined { return ( - document?.title?.trim() || - target.title?.trim() || - document?.customId || - target.customId || - target.documentId || - target.sourceId + (target.documentId + ? documentByKnownId.get(target.documentId) + : undefined) ?? + (target.customId ? documentByKnownId.get(target.customId) : undefined) ) } -function sourceSummary( - target: CitationTarget, - document?: DocumentWithMemories, -): string | null { - const summary = - document?.summary || - target.summary || - (document as { content?: string } | undefined)?.content || - null - return summary ? summary.trim() : null -} +function MemorySourcesFallback({ + citationIndex, + documentByKnownId, +}: { + citationIndex: Map + documentByKnownId: Map +}) { + const sources = Array.from(citationIndex.values()).slice(0, 6) + if (sources.length === 0) return null -function sourceKind( - target: CitationTarget, - document?: DocumentWithMemories, -): string { - return (document?.type || target.type || "memory").replaceAll("_", " ") + return ( +
+ Sources + {sources.map((target) => { + const document = documentForCitationTarget(target, documentByKnownId) + const display = getCitationDisplay(target, document) + const url = safeExternalUrl( + document ? getDocumentSourceUrl(document) : target.url, + ) + const label = `${target.sourceId}: ${display.title}` + const className = + "max-w-56 truncate rounded-full border border-white/10 bg-white/[0.035] px-2 py-1 text-white/55 transition-colors hover:bg-white/[0.06] hover:text-white/75" + + return url ? ( + + {label} + + ) : ( + + {label} + + ) + })} +
+ ) } function SourceCitationLink({ @@ -237,16 +265,11 @@ function SourceCitationLink({ const target = citationIndex.get(sourceId) if (!target) return <>{children} - const document = - (target.documentId - ? documentByKnownId.get(target.documentId) - : undefined) ?? - (target.customId ? documentByKnownId.get(target.customId) : undefined) + const document = documentForCitationTarget(target, documentByKnownId) const url = safeExternalUrl( document ? getDocumentSourceUrl(document) : target.url, ) - const title = sourceTitle(target, document) - const summary = sourceSummary(target, document) + const display = getCitationDisplay(target, document) return ( @@ -274,15 +297,15 @@ function SourceCitationLink({ - {title} + {display.title} - {sourceKind(target, document)} + {display.kind} - {summary ? ( + {display.summary ? ( - {summary} + {display.summary} ) : null} {url ? ( @@ -754,6 +777,16 @@ export function AgentMessage({ () => makeMarkdownComponents(webSources, citationIndex, documentByKnownId), [webSources, citationIndex, documentByKnownId], ) + const hasInlineSourceAnnotations = useMemo( + () => + parseSourceAnnotatedMarkdown( + messageText, + allowedSourceIds, + ).markdown.includes("#sm-source:"), + [messageText, allowedSourceIds], + ) + const showMemorySourcesFallback = + citationIndex.size > 0 && !hasInlineSourceAnnotations const responseModelLabel = responseModel ? `${modelNames[responseModel].name} ${modelNames[responseModel].version}` : null @@ -865,6 +898,12 @@ export function AgentMessage({ } return null })} + {showMemorySourcesFallback && ( + + )} {hasAssistantText && ( diff --git a/apps/web/lib/chat-memory-tools.test.ts b/apps/web/lib/chat-memory-tools.test.ts index 9b6411c7..f7e8a4c9 100644 --- a/apps/web/lib/chat-memory-tools.test.ts +++ b/apps/web/lib/chat-memory-tools.test.ts @@ -4,6 +4,7 @@ import { buildCitationIndex, extractDocumentIdsFromMemoryOutput, extractMemoryToolOutputs, + getCitationDisplay, getDocumentSourceUrl, mapDocumentsByKnownIds, } from "./chat-memory-tools" @@ -76,9 +77,101 @@ describe("chat memory tool citation mapping", () => { expect(index.get("S1")?.documentId).toBe("docA") expect(index.get("S1")?.customId).toBe("customA") + expect(index.get("S1")?.content).toBe("memo") expect(index.has("ignored")).toBe(false) }) + it("maps source ids to matching result ids when citation ids are absent", () => { + const outputs = extractMemoryToolOutputs({ + parts: [ + { + type: "tool-searchMemories", + state: "output-available", + output: { + sourceIds: ["memory_no_citation"], + results: [ + { + id: "memory_no_citation", + kind: "memory", + content: "Fallback source text.", + }, + ], + }, + }, + ], + }) + + expect(buildCitationIndex(outputs).get("memory_no_citation")).toMatchObject( + { + sourceId: "memory_no_citation", + memoryId: "memory_no_citation", + content: "Fallback source text.", + }, + ) + }) + + it("keeps memory text for citations without source documents", () => { + const outputs = extractMemoryToolOutputs({ + parts: [ + { + type: "tool-recallContext", + state: "output-available", + output: { + sourceIds: ["memory_1"], + results: [ + { + id: "memory_1", + citationId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }, + ], + }, + }, + ], + }) + + expect(buildCitationIndex(outputs).get("memory_1")).toMatchObject({ + sourceId: "memory_1", + memoryId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }) + }) + + it("uses memory display only for citations without document metadata", () => { + expect( + getCitationDisplay({ + sourceId: "memory_1", + memoryId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }), + ).toEqual({ + title: "Memory", + kind: "memory", + summary: "User prefers concise answers.", + }) + }) + + it("prefers tool-provided document metadata over generic memory display", () => { + expect( + getCitationDisplay({ + sourceId: "S1", + memoryId: "memory_1", + content: "Memory text", + documentId: "doc_1", + customId: "custom_1", + title: "Project Plan", + type: "google_doc", + }), + ).toMatchObject({ + title: "Project Plan", + kind: "google doc", + summary: "Memory text", + }) + }) + it("extracts graph highlight document ids from memory outputs", () => { const [output] = extractMemoryToolOutputs(assistantMessage) diff --git a/apps/web/lib/chat-memory-tools.ts b/apps/web/lib/chat-memory-tools.ts index 140c2ae6..6a39d2f5 100644 --- a/apps/web/lib/chat-memory-tools.ts +++ b/apps/web/lib/chat-memory-tools.ts @@ -55,6 +55,9 @@ export type MemoryToolOutput = { } export type CitationTarget = { sourceId: string + memoryId?: string | undefined + kind?: string | undefined + content?: string | undefined documentId?: string | undefined customId?: string | null | undefined title?: string | null | undefined @@ -63,6 +66,12 @@ export type CitationTarget = { url?: string | null | undefined } +export type CitationDisplay = { + title: string + summary: string | null + kind: string +} + export type DocumentWithMemories = z.infer< typeof DocumentsWithMemoriesResponseSchema >["documents"][0] @@ -156,10 +165,15 @@ function citationTargetForResult( doc?.id ?? result.documentIds?.find(Boolean) ?? result.documentId - const target = documentTarget(sourceId, doc) - target.documentId = target.documentId ?? docId - target.customId = target.customId ?? result.customId - return target + const docTarget = documentTarget(sourceId, doc) + return { + ...docTarget, + memoryId: result.id, + kind: result.kind, + content: result.content, + documentId: docTarget.documentId ?? docId, + customId: docTarget.customId ?? result.customId, + } } function documentTarget( @@ -247,7 +261,7 @@ export function buildCitationIndex( for (const sourceId of output.sourceIds ?? []) { if (index.has(sourceId)) continue const matchingResult = (output.results ?? []).find( - (result) => result.citationId === sourceId, + (result) => result.citationId === sourceId || result.id === sourceId, ) if (matchingResult) addTarget( @@ -344,6 +358,46 @@ export function mapDocumentsByKnownIds( return map } +function hasDisplayMetadata(target: CitationTarget): boolean { + return !!( + target.title?.trim() || + target.documentId || + target.customId || + target.url || + target.type || + target.summary?.trim() + ) +} + +export function getCitationDisplay( + target: CitationTarget, + document?: DocumentWithMemories, +): CitationDisplay { + const memoryOnly = !document && !hasDisplayMetadata(target) + const summary = + document?.summary || + target.summary || + target.content || + (document as { content?: string } | undefined)?.content || + null + + return { + title: memoryOnly + ? "Memory" + : document?.title?.trim() || + target.title?.trim() || + document?.customId || + target.customId || + target.documentId || + target.sourceId, + summary: summary ? summary.trim() : null, + kind: (document?.type || target.type || target.kind || "memory").replaceAll( + "_", + " ", + ), + } +} + export function getDocumentSourceUrl( document: Pick & { customId?: string | null