fix(chat): show memory source fallbacks

This commit is contained in:
Dhravya 2026-06-27 04:35:28 +00:00
parent f1ff7beb0f
commit 6798b565cb
3 changed files with 227 additions and 41 deletions

View file

@ -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<string, DocumentWithMemories>,
): 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<string, CitationTarget>
documentByKnownId: Map<string, DocumentWithMemories>
}) {
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 (
<div className="mt-3 flex flex-wrap gap-1.5 text-xs text-white/45">
<span className="mr-0.5 self-center text-white/35">Sources</span>
{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 ? (
<a
key={target.sourceId}
href={url}
target="_blank"
rel="noopener noreferrer"
className={className}
title={display.summary ?? label}
>
{label}
</a>
) : (
<span
key={target.sourceId}
className={className}
title={display.summary ?? label}
>
{label}
</span>
)
})}
</div>
)
}
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 (
<span className="group/source relative inline rounded-[3px] border-b border-dotted border-white/20 bg-white/[0.025] px-px text-white/90 transition-colors hover:border-white/35 hover:bg-white/[0.045] focus-within:border-white/35 focus-within:bg-white/[0.045]">
@ -274,15 +297,15 @@ function SourceCitationLink({
<span className="pointer-events-auto block rounded-xl border border-white/10 bg-[#0B0F16]/95 p-3 text-left shadow-[0_16px_44px_rgba(0,0,0,0.48)] backdrop-blur-xl">
<span className="mb-1 flex items-center justify-between gap-2">
<span className="truncate text-xs font-medium text-white/85">
{title}
{display.title}
</span>
<span className="shrink-0 rounded-full bg-white/5 px-2 py-0.5 text-[10px] capitalize text-white/40">
{sourceKind(target, document)}
{display.kind}
</span>
</span>
{summary ? (
{display.summary ? (
<span className="line-clamp-3 text-xs leading-snug text-white/55">
{summary}
{display.summary}
</span>
) : 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 && (
<MemorySourcesFallback
citationIndex={citationIndex}
documentByKnownId={documentByKnownId}
/>
)}
</div>
</div>
{hasAssistantText && (

View file

@ -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)

View file

@ -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<DocumentWithMemories, "type" | "url"> & {
customId?: string | null