mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-08-28 05:25:33 +00:00
Draft: Add chat source annotations (#1165)
This commit is contained in:
Dhravya Shah
GitHub
parent
d169dc078e
commit
f1ff7beb0f
7 changed files with 1103 additions and 22 deletions
|
|
@ -48,7 +48,10 @@ import {
|
|||
import { SpaceSelector } from "@/components/space-selector"
|
||||
import { SuperLoader } from "../superloader"
|
||||
import { UserMessage } from "./message/user-message"
|
||||
import { AgentMessage } from "./message/agent-message"
|
||||
import {
|
||||
AgentMessage,
|
||||
isChatToolDisplayPartType,
|
||||
} from "./message/agent-message"
|
||||
import { ChatGraphContextRail } from "./chat-graph-context-rail"
|
||||
import { ChainOfThought } from "./input/chain-of-thought"
|
||||
import { useIsMobile } from "@hooks/use-mobile"
|
||||
|
|
@ -1147,14 +1150,14 @@ export function ChatSidebar({
|
|||
}) => ({
|
||||
id: m.id,
|
||||
role: m.role,
|
||||
// Strip tool parts (they break convertToModelMessages with tool_use/tool_result
|
||||
// mismatches); keep text/reasoning + source parts so citations survive reload.
|
||||
// Keep chat tool outputs that are meaningful to render after thread reload.
|
||||
parts: (m.parts || []).filter(
|
||||
(p) =>
|
||||
p.type === "text" ||
|
||||
p.type === "reasoning" ||
|
||||
p.type === "source-url" ||
|
||||
p.type === "source-document",
|
||||
p.type === "source-document" ||
|
||||
isChatToolDisplayPartType(p.type),
|
||||
),
|
||||
metadata: m.metadata,
|
||||
createdAt: new Date(m.createdAt),
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import { type ReactNode, useEffect, useMemo, useRef, useState } from "react"
|
||||
import type { UIMessage } from "@ai-sdk/react"
|
||||
import { useQuery } from "@tanstack/react-query"
|
||||
import { Streamdown } from "streamdown"
|
||||
import {
|
||||
BookOpenIcon,
|
||||
|
|
@ -19,18 +20,37 @@ import {
|
|||
} from "lucide-react"
|
||||
import { cn } from "@lib/utils"
|
||||
import { isWebSearchToolName } from "@/lib/chat-web-search-tools"
|
||||
import {
|
||||
buildCitationIndex,
|
||||
fetchDocumentsByIds,
|
||||
getDocumentSourceUrl,
|
||||
isMemoryToolOutputReady,
|
||||
mapDocumentsByKnownIds,
|
||||
type CitationTarget,
|
||||
type DocumentWithMemories,
|
||||
extractMemoryToolOutputs,
|
||||
} from "@/lib/chat-memory-tools"
|
||||
import {
|
||||
parseSourceAnnotatedMarkdown,
|
||||
stripSourceMarkup,
|
||||
} from "@/lib/source-annotations"
|
||||
import { modelNames, type ModelId } from "@/lib/models"
|
||||
import { RelatedMemories } from "./related-memories"
|
||||
import { MessageActions } from "./message-actions"
|
||||
|
||||
const TOOL_META: Record<string, { label: string; icon: typeof SearchIcon }> = {
|
||||
bash: { label: "Memory", icon: TerminalIcon },
|
||||
recallContext: { label: "Recall Memories", icon: BookOpenIcon },
|
||||
discoverSpaces: { label: "Discover Spaces", icon: SearchIcon },
|
||||
web_search: { label: "Web search", icon: GlobeIcon },
|
||||
google_search: { label: "Google search", icon: GlobeIcon },
|
||||
// legacy tool names kept for existing persisted messages
|
||||
searchMemories: { label: "Search Memories", icon: SearchIcon },
|
||||
addMemory: { label: "Add Memory", icon: PlusIcon },
|
||||
fetchMemory: { label: "Fetch Memory", icon: BookOpenIcon },
|
||||
forgetMemory: { label: "Forget Memory", icon: XCircleIcon },
|
||||
updateMemory: { label: "Update Memory", icon: BookOpenIcon },
|
||||
forgetDocument: { label: "Forget Document", icon: XCircleIcon },
|
||||
scheduleTask: { label: "Schedule Task", icon: ClockIcon },
|
||||
listSchedules: { label: "List Schedules", icon: ListIcon },
|
||||
cancelSchedule: { label: "Cancel Schedule", icon: XCircleIcon },
|
||||
|
|
@ -38,7 +58,7 @@ const TOOL_META: Record<string, { label: string; icon: typeof SearchIcon }> = {
|
|||
|
||||
type ToolCallDisplayPart = {
|
||||
type: string
|
||||
state: string
|
||||
state?: string
|
||||
input?: unknown
|
||||
output?: unknown
|
||||
toolCallId?: string
|
||||
|
|
@ -64,6 +84,20 @@ function faviconUrl(host: string): string {
|
|||
return `https://www.google.com/s2/favicons?sz=64&domain=${host}`
|
||||
}
|
||||
|
||||
function safeExternalUrl(url: string | null | undefined): string | null {
|
||||
if (!url) return null
|
||||
if (url.startsWith("/") && !url.startsWith("//")) return url
|
||||
if (url.startsWith("#") && !url.startsWith("#sm-source:")) return url
|
||||
try {
|
||||
const parsed = new URL(url)
|
||||
return parsed.protocol === "http:" || parsed.protocol === "https:"
|
||||
? url
|
||||
: null
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
function isWebSearchPart(part: { type: string; toolName?: string }): boolean {
|
||||
if (part.type === "dynamic-tool") {
|
||||
return isWebSearchToolName(part.toolName ?? "")
|
||||
|
|
@ -74,6 +108,25 @@ function isWebSearchPart(part: { type: string; toolName?: string }): boolean {
|
|||
return false
|
||||
}
|
||||
|
||||
function isMemoryRetrievalToolName(toolName: string): boolean {
|
||||
return (
|
||||
toolName === "searchMemories" ||
|
||||
toolName === "recallContext" ||
|
||||
toolName === "discoverSpaces"
|
||||
)
|
||||
}
|
||||
|
||||
export function isChatToolDisplayPartType(type: string): boolean {
|
||||
return (
|
||||
type === "tool-searchMemories" ||
|
||||
type === "tool-recallContext" ||
|
||||
type === "tool-discoverSpaces" ||
|
||||
type === "tool-forgetMemory" ||
|
||||
type === "tool-updateMemory" ||
|
||||
type === "tool-forgetDocument"
|
||||
)
|
||||
}
|
||||
|
||||
function CitationLink({
|
||||
href,
|
||||
label,
|
||||
|
|
@ -83,7 +136,8 @@ function CitationLink({
|
|||
label: string
|
||||
source?: SourceUrlPart
|
||||
}) {
|
||||
const url = source?.url ?? href
|
||||
const url = safeExternalUrl(source?.url ?? href) ?? ""
|
||||
if (!url) return <>{label}</>
|
||||
const host = sourceHost(url)
|
||||
const rawTitle = source?.title?.trim()
|
||||
const hasTitle =
|
||||
|
|
@ -136,9 +190,138 @@ function CitationLink({
|
|||
)
|
||||
}
|
||||
|
||||
function makeMarkdownComponents(sources: SourceUrlPart[]) {
|
||||
function sourceTitle(
|
||||
target: CitationTarget,
|
||||
document?: DocumentWithMemories,
|
||||
): string {
|
||||
return (
|
||||
document?.title?.trim() ||
|
||||
target.title?.trim() ||
|
||||
document?.customId ||
|
||||
target.customId ||
|
||||
target.documentId ||
|
||||
target.sourceId
|
||||
)
|
||||
}
|
||||
|
||||
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 sourceKind(
|
||||
target: CitationTarget,
|
||||
document?: DocumentWithMemories,
|
||||
): string {
|
||||
return (document?.type || target.type || "memory").replaceAll("_", " ")
|
||||
}
|
||||
|
||||
function SourceCitationLink({
|
||||
sourceId,
|
||||
children,
|
||||
citationIndex,
|
||||
documentByKnownId,
|
||||
}: {
|
||||
sourceId: string
|
||||
children: ReactNode
|
||||
citationIndex: Map<string, CitationTarget>
|
||||
documentByKnownId: Map<string, DocumentWithMemories>
|
||||
}) {
|
||||
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 url = safeExternalUrl(
|
||||
document ? getDocumentSourceUrl(document) : target.url,
|
||||
)
|
||||
const title = sourceTitle(target, document)
|
||||
const summary = sourceSummary(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]">
|
||||
{url ? (
|
||||
<a
|
||||
href={url}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-inherit no-underline outline-none focus-visible:ring-1 focus-visible:ring-white/25"
|
||||
>
|
||||
{children}
|
||||
</a>
|
||||
) : (
|
||||
<button
|
||||
type="button"
|
||||
className="cursor-help border-0 bg-transparent p-0 text-inherit"
|
||||
>
|
||||
{children}
|
||||
</button>
|
||||
)}
|
||||
<span className="ml-1 inline-flex h-3.5 min-w-3.5 translate-y-[-1px] items-center justify-center rounded-full border border-white/10 bg-white/[0.04] px-1 text-[9px] font-medium leading-none text-white/45 transition-colors group-hover/source:text-white/65 group-focus-within/source:text-white/65">
|
||||
{sourceId}
|
||||
</span>
|
||||
<span className="pointer-events-none absolute bottom-full left-1/2 z-[1000] hidden w-72 -translate-x-1/2 pb-2 group-hover/source:block group-focus-within/source:block">
|
||||
<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}
|
||||
</span>
|
||||
<span className="shrink-0 rounded-full bg-white/5 px-2 py-0.5 text-[10px] capitalize text-white/40">
|
||||
{sourceKind(target, document)}
|
||||
</span>
|
||||
</span>
|
||||
{summary ? (
|
||||
<span className="line-clamp-3 text-xs leading-snug text-white/55">
|
||||
{summary}
|
||||
</span>
|
||||
) : null}
|
||||
{url ? (
|
||||
<span className="mt-2 block text-xs font-medium text-blue-300">
|
||||
Open source
|
||||
</span>
|
||||
) : null}
|
||||
</span>
|
||||
</span>
|
||||
</span>
|
||||
)
|
||||
}
|
||||
|
||||
function makeMarkdownComponents(
|
||||
sources: SourceUrlPart[],
|
||||
citationIndex: Map<string, CitationTarget>,
|
||||
documentByKnownId: Map<string, DocumentWithMemories>,
|
||||
) {
|
||||
return {
|
||||
a: ({ href, children }: { href?: string; children?: ReactNode }) => {
|
||||
if (href?.startsWith("#sm-source:")) {
|
||||
const sourceId = (() => {
|
||||
try {
|
||||
return decodeURIComponent(href.slice("#sm-source:".length))
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
})()
|
||||
if (!sourceId) return <>{children}</>
|
||||
return (
|
||||
<SourceCitationLink
|
||||
sourceId={sourceId}
|
||||
citationIndex={citationIndex}
|
||||
documentByKnownId={documentByKnownId}
|
||||
>
|
||||
{children}
|
||||
</SourceCitationLink>
|
||||
)
|
||||
}
|
||||
const label =
|
||||
typeof children === "string"
|
||||
? children
|
||||
|
|
@ -146,14 +329,16 @@ function makeMarkdownComponents(sources: SourceUrlPart[]) {
|
|||
? children.join("")
|
||||
: ""
|
||||
const match = label.match(/^\[?(\d+)\]?$/)
|
||||
if (match && href) {
|
||||
const safeHref = safeExternalUrl(href)
|
||||
if (match && safeHref) {
|
||||
const n = Number(match[1])
|
||||
const source = sources.find((s) => s.url === href) ?? sources[n - 1]
|
||||
return <CitationLink href={href} label={label} source={source} />
|
||||
const source = sources.find((s) => s.url === safeHref) ?? sources[n - 1]
|
||||
return <CitationLink href={safeHref} label={label} source={source} />
|
||||
}
|
||||
if (!safeHref) return <>{children}</>
|
||||
return (
|
||||
<a
|
||||
href={href}
|
||||
href={safeHref}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-blue-400 hover:underline"
|
||||
|
|
@ -397,6 +582,9 @@ function ToolCallDisplay({ part }: { part: ToolCallDisplayPart }) {
|
|||
const isDone = part.state === "output-available"
|
||||
const isError = part.state === "error" || part.state === "output-error"
|
||||
const errorText = part.errorText
|
||||
if (isMemoryRetrievalToolName(toolName) && isMemoryToolOutputReady(part)) {
|
||||
return null
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="rounded-lg border border-[#1E2128] bg-[#0D121A] text-xs my-1 overflow-hidden">
|
||||
|
|
@ -516,7 +704,38 @@ export function AgentMessage({
|
|||
.filter((part) => part.type === "text")
|
||||
.map((part) => part.text)
|
||||
.join(" ")
|
||||
const webSources = (() => {
|
||||
const copyText = stripSourceMarkup(messageText)
|
||||
const memoryOutputs = useMemo(
|
||||
() => extractMemoryToolOutputs(message),
|
||||
[message],
|
||||
)
|
||||
const citationIndex = useMemo(
|
||||
() => buildCitationIndex(memoryOutputs),
|
||||
[memoryOutputs],
|
||||
)
|
||||
const allowedSourceIds = useMemo(
|
||||
() => new Set(citationIndex.keys()),
|
||||
[citationIndex],
|
||||
)
|
||||
const sourceDocumentIds = useMemo(() => {
|
||||
const ids = new Set<string>()
|
||||
for (const target of citationIndex.values()) {
|
||||
if (target.documentId) ids.add(target.documentId)
|
||||
if (target.customId) ids.add(target.customId)
|
||||
}
|
||||
return [...ids].sort()
|
||||
}, [citationIndex])
|
||||
const { data: sourceDocuments = [] } = useQuery({
|
||||
queryKey: ["chat-source-documents", sourceDocumentIds],
|
||||
queryFn: () => fetchDocumentsByIds(sourceDocumentIds),
|
||||
enabled: sourceDocumentIds.length > 0,
|
||||
staleTime: 5 * 60 * 1000,
|
||||
})
|
||||
const documentByKnownId = useMemo(
|
||||
() => mapDocumentsByKnownIds(sourceDocuments),
|
||||
[sourceDocuments],
|
||||
)
|
||||
const webSources = useMemo(() => {
|
||||
const seen = new Set<string>()
|
||||
const out: SourceUrlPart[] = []
|
||||
for (const part of message.parts) {
|
||||
|
|
@ -527,15 +746,13 @@ export function AgentMessage({
|
|||
out.push(source)
|
||||
}
|
||||
return out
|
||||
})()
|
||||
}, [message.parts])
|
||||
const hasAssistantText = message.parts.some(
|
||||
(p) => p.type === "text" && (p as { text?: string }).text?.trim(),
|
||||
)
|
||||
const sourceKey = webSources.map((s) => s.url).join("|")
|
||||
// biome-ignore lint/correctness/useExhaustiveDependencies: keyed by stable source urls
|
||||
const markdownComponents = useMemo(
|
||||
() => makeMarkdownComponents(webSources),
|
||||
[sourceKey],
|
||||
() => makeMarkdownComponents(webSources, citationIndex, documentByKnownId),
|
||||
[webSources, citationIndex, documentByKnownId],
|
||||
)
|
||||
const responseModelLabel = responseModel
|
||||
? `${modelNames[responseModel].name} ${modelNames[responseModel].version}`
|
||||
|
|
@ -603,7 +820,10 @@ export function AgentMessage({
|
|||
className="text-sm text-white/90 chat-markdown-content"
|
||||
>
|
||||
<Streamdown components={markdownComponents}>
|
||||
{runText}
|
||||
{
|
||||
parseSourceAnnotatedMarkdown(runText, allowedSourceIds)
|
||||
.markdown
|
||||
}
|
||||
</Streamdown>
|
||||
</div>
|
||||
)
|
||||
|
|
@ -613,7 +833,7 @@ export function AgentMessage({
|
|||
type: "dynamic-tool"
|
||||
toolName: string
|
||||
toolCallId: string
|
||||
state: string
|
||||
state?: string
|
||||
input?: unknown
|
||||
output?: unknown
|
||||
errorText?: string
|
||||
|
|
@ -651,7 +871,7 @@ export function AgentMessage({
|
|||
<div className="flex min-h-7 items-center gap-2">
|
||||
<MessageActions
|
||||
messageId={message.id}
|
||||
messageText={messageText}
|
||||
messageText={copyText}
|
||||
isLastMessage={isLastAgentMessage}
|
||||
isHovered={isHovered}
|
||||
copiedMessageId={copiedMessageId}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
import type { UIMessage } from "@ai-sdk/react"
|
||||
import {
|
||||
extractDocumentIdsFromMemoryOutput,
|
||||
extractMemoryToolOutputs,
|
||||
} from "@/lib/chat-memory-tools"
|
||||
import { memoryResultsFromSearchToolOutput } from "@/lib/chat-search-memory-results"
|
||||
|
||||
const UUID_IN_STRING =
|
||||
|
|
@ -42,7 +46,8 @@ export function documentIdsFromBashText(text: string): string[] {
|
|||
const found = new Set<string>()
|
||||
// [doc:<nanoid>] annotations from sgrep --include-doc-ids (highest confidence)
|
||||
for (const m of text.matchAll(DOC_ANNOTATION)) {
|
||||
found.add(m[1])
|
||||
const id = m[1]
|
||||
if (id) found.add(id)
|
||||
}
|
||||
// Standard UUID format
|
||||
for (const m of text.matchAll(UUID_IN_STRING)) {
|
||||
|
|
@ -52,7 +57,8 @@ export function documentIdsFromBashText(text: string): string[] {
|
|||
const quoted = /"documentId"\s*:\s*"([^"]+)"/g
|
||||
let q = quoted.exec(text)
|
||||
while (q !== null) {
|
||||
found.add(q[1])
|
||||
const id = q[1]
|
||||
if (id) found.add(id)
|
||||
q = quoted.exec(text)
|
||||
}
|
||||
return [...found]
|
||||
|
|
@ -136,9 +142,23 @@ export function extractHighlightDocumentIdsFromMessages(
|
|||
if (message.role !== "assistant") continue
|
||||
const parts = message.parts
|
||||
if (!parts) continue
|
||||
for (const memoryOutput of extractMemoryToolOutputs(message)) {
|
||||
for (const id of extractDocumentIdsFromMemoryOutput(
|
||||
memoryOutput.output,
|
||||
)) {
|
||||
ids.add(id)
|
||||
}
|
||||
}
|
||||
|
||||
for (const part of parts) {
|
||||
const p = part as Record<string, unknown>
|
||||
if (
|
||||
p.type === "tool-searchMemories" ||
|
||||
p.type === "tool-recallContext" ||
|
||||
p.type === "tool-discoverSpaces"
|
||||
) {
|
||||
continue
|
||||
}
|
||||
|
||||
if (p.type === "source-document") {
|
||||
const sid = (p as { sourceId?: unknown }).sourceId
|
||||
|
|
|
|||
177
apps/web/lib/chat-memory-tools.test.ts
Normal file
177
apps/web/lib/chat-memory-tools.test.ts
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
import { describe, expect, it } from "bun:test"
|
||||
import { extractHighlightDocumentIdsFromMessages } from "./chat-highlight-documents"
|
||||
import {
|
||||
buildCitationIndex,
|
||||
extractDocumentIdsFromMemoryOutput,
|
||||
extractMemoryToolOutputs,
|
||||
getDocumentSourceUrl,
|
||||
mapDocumentsByKnownIds,
|
||||
} from "./chat-memory-tools"
|
||||
|
||||
const assistantMessage = {
|
||||
id: "m1",
|
||||
role: "assistant",
|
||||
parts: [
|
||||
{
|
||||
type: "tool-recallContext",
|
||||
state: "output-available",
|
||||
output: {
|
||||
sourceIds: ["S1"],
|
||||
documentIds: ["topDoc"],
|
||||
results: [
|
||||
{
|
||||
citationId: "S1",
|
||||
content: "memo",
|
||||
document: {
|
||||
id: "docA",
|
||||
customId: "customA",
|
||||
title: "Doc A",
|
||||
type: "google_doc",
|
||||
summary: "sum",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
{
|
||||
type: "tool-discoverSpaces",
|
||||
state: "input-streaming",
|
||||
output: { sourceIds: ["ignored"], documentIds: ["ignoredDoc"] },
|
||||
},
|
||||
{
|
||||
type: "text",
|
||||
text: 'Answer <response source="S1">from memory</response>',
|
||||
},
|
||||
],
|
||||
} as const
|
||||
|
||||
describe("chat memory tool citation mapping", () => {
|
||||
it("extracts only ready memory tool outputs", () => {
|
||||
const outputs = extractMemoryToolOutputs({
|
||||
parts: [
|
||||
...assistantMessage.parts,
|
||||
{
|
||||
type: "tool-searchMemories",
|
||||
state: "done",
|
||||
output: { sourceIds: ["done"], documentIds: ["doneDoc"] },
|
||||
},
|
||||
{
|
||||
type: "tool-searchMemories",
|
||||
output: { sourceIds: ["stateless"], documentIds: ["statelessDoc"] },
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
expect(outputs).toHaveLength(3)
|
||||
expect(outputs.map((output) => output.output.sourceIds?.[0])).toEqual([
|
||||
"S1",
|
||||
"done",
|
||||
"stateless",
|
||||
])
|
||||
})
|
||||
|
||||
it("maps citation ids to document and custom ids", () => {
|
||||
const [output] = extractMemoryToolOutputs(assistantMessage)
|
||||
const index = buildCitationIndex(output ? [output] : [])
|
||||
|
||||
expect(index.get("S1")?.documentId).toBe("docA")
|
||||
expect(index.get("S1")?.customId).toBe("customA")
|
||||
expect(index.has("ignored")).toBe(false)
|
||||
})
|
||||
|
||||
it("extracts graph highlight document ids from memory outputs", () => {
|
||||
const [output] = extractMemoryToolOutputs(assistantMessage)
|
||||
|
||||
expect(output && extractDocumentIdsFromMemoryOutput(output.output)).toEqual(
|
||||
["topDoc", "docA", "customA"],
|
||||
)
|
||||
expect(
|
||||
extractHighlightDocumentIdsFromMessages([assistantMessage as never]),
|
||||
).toEqual(["topDoc", "docA", "customA"])
|
||||
})
|
||||
|
||||
it("keeps graph highlights for legacy memory tool states and ids", () => {
|
||||
const legacyMessage = {
|
||||
id: "legacy",
|
||||
role: "assistant",
|
||||
parts: [
|
||||
{
|
||||
type: "tool-searchMemories",
|
||||
state: "done",
|
||||
output: { results: [{ id: "legacyDoc" }] },
|
||||
},
|
||||
{
|
||||
type: "tool-recallContext",
|
||||
output: { documentIds: ["statelessDoc"] },
|
||||
},
|
||||
],
|
||||
} as const
|
||||
|
||||
expect(extractMemoryToolOutputs(legacyMessage)).toHaveLength(2)
|
||||
expect(
|
||||
extractHighlightDocumentIdsFromMessages([legacyMessage as never]),
|
||||
).toEqual(["legacyDoc", "statelessDoc"])
|
||||
})
|
||||
|
||||
it("normalizes nested discoverSpaces memory results", () => {
|
||||
const outputs = extractMemoryToolOutputs({
|
||||
parts: [
|
||||
{
|
||||
type: "tool-discoverSpaces",
|
||||
state: "output-available",
|
||||
output: {
|
||||
spaces: [
|
||||
{
|
||||
sourceIds: ["S2"],
|
||||
documentIds: ["spaceDoc"],
|
||||
results: [{ citationId: "S2", documentIds: ["nestedDoc"] }],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const index = buildCitationIndex(outputs)
|
||||
expect(index.get("S2")?.documentId).toBe("nestedDoc")
|
||||
expect(
|
||||
extractDocumentIdsFromMemoryOutput(outputs[0]?.output ?? {}),
|
||||
).toEqual(["spaceDoc", "nestedDoc"])
|
||||
})
|
||||
|
||||
it("builds editable Google source URLs from custom ids and API URLs", () => {
|
||||
expect(
|
||||
getDocumentSourceUrl({
|
||||
type: "google_doc",
|
||||
customId: "docCustom",
|
||||
url: "https://docs.googleapis.com/v1/documents/apiDoc",
|
||||
} as never),
|
||||
).toBe("https://docs.google.com/document/d/docCustom/edit")
|
||||
expect(
|
||||
getDocumentSourceUrl({
|
||||
type: "google_doc",
|
||||
url: "https://docs.googleapis.com/v1/documents/apiDoc",
|
||||
} as never),
|
||||
).toBe("https://docs.google.com/document/d/apiDoc/edit")
|
||||
expect(
|
||||
getDocumentSourceUrl({
|
||||
type: "google_sheet",
|
||||
url: "https://sheets.googleapis.com/v4/spreadsheets/sheetId/values/A1",
|
||||
} as never),
|
||||
).toBe("https://docs.google.com/spreadsheets/d/sheetId/edit")
|
||||
expect(
|
||||
getDocumentSourceUrl({
|
||||
type: "google_slide",
|
||||
url: "https://slides.googleapis.com/v1/presentations/slideId/pages",
|
||||
} as never),
|
||||
).toBe("https://docs.google.com/presentation/d/slideId/edit")
|
||||
})
|
||||
|
||||
it("maps documents by all known ids", () => {
|
||||
const mapped = mapDocumentsByKnownIds([
|
||||
{ id: "docA", customId: "customA", type: "text", url: null } as never,
|
||||
])
|
||||
expect(mapped.get("docA")?.id).toBe("docA")
|
||||
expect(mapped.get("customA")?.id).toBe("docA")
|
||||
})
|
||||
})
|
||||
375
apps/web/lib/chat-memory-tools.ts
Normal file
375
apps/web/lib/chat-memory-tools.ts
Normal file
|
|
@ -0,0 +1,375 @@
|
|||
import { $fetch } from "@lib/api"
|
||||
import type { DocumentsWithMemoriesResponseSchema } from "@repo/validation/api"
|
||||
import type { z } from "zod"
|
||||
import { isSafeSourceId } from "./source-annotations"
|
||||
|
||||
export const MEMORY_TOOL_PART_TYPES = [
|
||||
"tool-searchMemories",
|
||||
"tool-recallContext",
|
||||
"tool-discoverSpaces",
|
||||
] as const
|
||||
export const MAX_INLINE_GRAPH_DOCUMENT_IDS = 20
|
||||
|
||||
export type MemoryToolName =
|
||||
| "searchMemories"
|
||||
| "recallContext"
|
||||
| "discoverSpaces"
|
||||
export type ToolDocumentMetadata = {
|
||||
id?: string | undefined
|
||||
internalDocumentId?: string | undefined
|
||||
customId?: string | null | undefined
|
||||
title?: string | null | undefined
|
||||
type?: string | null | undefined
|
||||
summary?: string | null | undefined
|
||||
url?: string | null | undefined
|
||||
}
|
||||
export type MemoryToolResultItem = {
|
||||
id?: string | undefined
|
||||
citationId?: string | undefined
|
||||
kind?: "memory" | "chunk" | "aggregate" | string | undefined
|
||||
content?: string | undefined
|
||||
score?: number | undefined
|
||||
documentId?: string | undefined
|
||||
documentIds?: string[] | undefined
|
||||
internalDocumentId?: string | undefined
|
||||
customId?: string | undefined
|
||||
documents?: ToolDocumentMetadata[] | undefined
|
||||
document?: ToolDocumentMetadata | undefined
|
||||
}
|
||||
export type MemoryToolResult = {
|
||||
query?: string | undefined
|
||||
count?: number | undefined
|
||||
sourceIds?: string[] | undefined
|
||||
documentIds?: string[] | undefined
|
||||
results?: MemoryToolResultItem[] | undefined
|
||||
spaces?:
|
||||
| Array<{
|
||||
results?: MemoryToolResultItem[] | undefined
|
||||
sourceIds?: string[] | undefined
|
||||
documentIds?: string[] | undefined
|
||||
}>
|
||||
| undefined
|
||||
}
|
||||
export type MemoryToolOutput = {
|
||||
output: MemoryToolResult
|
||||
}
|
||||
export type CitationTarget = {
|
||||
sourceId: string
|
||||
documentId?: string | undefined
|
||||
customId?: string | null | undefined
|
||||
title?: string | null | undefined
|
||||
type?: string | null | undefined
|
||||
summary?: string | null | undefined
|
||||
url?: string | null | undefined
|
||||
}
|
||||
|
||||
export type DocumentWithMemories = z.infer<
|
||||
typeof DocumentsWithMemoriesResponseSchema
|
||||
>["documents"][0]
|
||||
|
||||
function isObject(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null
|
||||
}
|
||||
|
||||
function strings(values: unknown): string[] {
|
||||
if (!Array.isArray(values)) return []
|
||||
return values.filter(
|
||||
(value): value is string => typeof value === "string" && value.length > 0,
|
||||
)
|
||||
}
|
||||
|
||||
function normalizeOutput(output: unknown): MemoryToolResult {
|
||||
if (!isObject(output)) return {}
|
||||
|
||||
const nested = [output]
|
||||
for (const key of [
|
||||
"memory",
|
||||
"search",
|
||||
"searchResult",
|
||||
"memoryResult",
|
||||
"memoryOutput",
|
||||
"hints",
|
||||
]) {
|
||||
const value = output[key]
|
||||
if (isObject(value)) nested.push(value)
|
||||
}
|
||||
|
||||
const sourceIds: string[] = []
|
||||
const documentIds: string[] = []
|
||||
const results: MemoryToolResultItem[] = []
|
||||
|
||||
const merged: MemoryToolResult = {
|
||||
query: typeof output.query === "string" ? output.query : undefined,
|
||||
count: typeof output.count === "number" ? output.count : undefined,
|
||||
sourceIds,
|
||||
documentIds,
|
||||
results,
|
||||
spaces: Array.isArray(output.spaces)
|
||||
? (output.spaces.filter(isObject) as MemoryToolResult["spaces"])
|
||||
: undefined,
|
||||
}
|
||||
|
||||
for (const value of nested) {
|
||||
sourceIds.push(...strings(value.sourceIds))
|
||||
documentIds.push(...strings(value.documentIds))
|
||||
if (Array.isArray(value.results))
|
||||
results.push(
|
||||
...(value.results.filter(isObject) as MemoryToolResultItem[]),
|
||||
)
|
||||
}
|
||||
|
||||
for (const space of merged.spaces ?? []) {
|
||||
sourceIds.push(...strings(space.sourceIds))
|
||||
documentIds.push(...strings(space.documentIds))
|
||||
if (Array.isArray(space.results))
|
||||
results.push(
|
||||
...(space.results.filter(isObject) as MemoryToolResultItem[]),
|
||||
)
|
||||
}
|
||||
|
||||
merged.sourceIds = dedupe(merged.sourceIds ?? [])
|
||||
merged.documentIds = dedupe(merged.documentIds ?? [])
|
||||
return merged
|
||||
}
|
||||
|
||||
function dedupe(values: string[]): string[] {
|
||||
return Array.from(new Set(values.filter(Boolean)))
|
||||
}
|
||||
|
||||
function addTarget(
|
||||
index: Map<string, CitationTarget>,
|
||||
key: unknown,
|
||||
target: CitationTarget,
|
||||
) {
|
||||
if (typeof key !== "string" || !isSafeSourceId(key)) return
|
||||
if (!index.has(key)) index.set(key, { ...target, sourceId: key })
|
||||
}
|
||||
|
||||
function citationTargetForResult(
|
||||
sourceId: string,
|
||||
result: MemoryToolResultItem,
|
||||
): CitationTarget {
|
||||
const doc = firstDocumentForResult(result)
|
||||
const docId =
|
||||
doc?.internalDocumentId ??
|
||||
result.internalDocumentId ??
|
||||
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
|
||||
}
|
||||
|
||||
function documentTarget(
|
||||
sourceId: string,
|
||||
doc?: ToolDocumentMetadata | null,
|
||||
): CitationTarget {
|
||||
return {
|
||||
sourceId,
|
||||
documentId: doc?.internalDocumentId ?? doc?.id,
|
||||
customId:
|
||||
doc?.customId ??
|
||||
(doc?.internalDocumentId && doc.id !== doc.internalDocumentId
|
||||
? doc.id
|
||||
: undefined),
|
||||
title: doc?.title,
|
||||
type: doc?.type,
|
||||
summary: doc?.summary,
|
||||
url: doc?.url,
|
||||
}
|
||||
}
|
||||
|
||||
function firstDocumentForResult(
|
||||
result: MemoryToolResultItem,
|
||||
): ToolDocumentMetadata | null {
|
||||
if (result.document && isObject(result.document)) return result.document
|
||||
if (Array.isArray(result.documents) && result.documents.length > 0)
|
||||
return result.documents.find(isObject) ?? null
|
||||
const firstId =
|
||||
result.documentIds?.find(Boolean) ??
|
||||
result.documentId ??
|
||||
result.internalDocumentId ??
|
||||
result.customId
|
||||
return firstId ? { id: firstId, customId: result.customId } : null
|
||||
}
|
||||
|
||||
export function isMemoryToolOutputReady(
|
||||
part: Record<string, unknown>,
|
||||
): boolean {
|
||||
return (
|
||||
part.state === "output-available" ||
|
||||
part.state === "done" ||
|
||||
(part.state === undefined && part.output !== undefined)
|
||||
)
|
||||
}
|
||||
|
||||
export function extractMemoryToolOutputs(message: {
|
||||
parts?: readonly unknown[]
|
||||
}): MemoryToolOutput[] {
|
||||
const parts = Array.isArray(message.parts) ? message.parts : []
|
||||
const outputs: MemoryToolOutput[] = []
|
||||
|
||||
for (let partIndex = 0; partIndex < parts.length; partIndex++) {
|
||||
const part = parts[partIndex]
|
||||
if (!isObject(part)) continue
|
||||
const type = part.type
|
||||
if (
|
||||
typeof type !== "string" ||
|
||||
!MEMORY_TOOL_PART_TYPES.includes(
|
||||
type as (typeof MEMORY_TOOL_PART_TYPES)[number],
|
||||
)
|
||||
)
|
||||
continue
|
||||
if (!isMemoryToolOutputReady(part)) continue
|
||||
outputs.push({ output: normalizeOutput(part.output) })
|
||||
}
|
||||
|
||||
return outputs
|
||||
}
|
||||
|
||||
export function buildCitationIndex(
|
||||
outputs: MemoryToolOutput[],
|
||||
): Map<string, CitationTarget> {
|
||||
const index = new Map<string, CitationTarget>()
|
||||
|
||||
for (const { output } of outputs) {
|
||||
for (const result of output.results ?? []) {
|
||||
if (result.citationId)
|
||||
addTarget(
|
||||
index,
|
||||
result.citationId,
|
||||
citationTargetForResult(result.citationId, result),
|
||||
)
|
||||
}
|
||||
|
||||
for (const sourceId of output.sourceIds ?? []) {
|
||||
if (index.has(sourceId)) continue
|
||||
const matchingResult = (output.results ?? []).find(
|
||||
(result) => result.citationId === sourceId,
|
||||
)
|
||||
if (matchingResult)
|
||||
addTarget(
|
||||
index,
|
||||
sourceId,
|
||||
citationTargetForResult(sourceId, matchingResult),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return index
|
||||
}
|
||||
|
||||
export function extractDocumentIdsFromMemoryOutput(
|
||||
output: MemoryToolResult,
|
||||
): string[] {
|
||||
const ids: string[] = []
|
||||
ids.push(...(output.documentIds ?? []))
|
||||
for (const result of output.results ?? []) {
|
||||
if (result.id) ids.push(result.id)
|
||||
if (result.documentId) ids.push(result.documentId)
|
||||
if (result.internalDocumentId) ids.push(result.internalDocumentId)
|
||||
ids.push(...(result.documentIds ?? []))
|
||||
if (result.document?.id) ids.push(result.document.id)
|
||||
if (result.document?.customId) ids.push(result.document.customId)
|
||||
for (const doc of result.documents ?? []) {
|
||||
if (doc.internalDocumentId) ids.push(doc.internalDocumentId)
|
||||
if (doc.id) ids.push(doc.id)
|
||||
if (doc.customId) ids.push(doc.customId)
|
||||
}
|
||||
}
|
||||
return dedupe(ids).slice(0, MAX_INLINE_GRAPH_DOCUMENT_IDS)
|
||||
}
|
||||
|
||||
export async function fetchDocumentsByIds(
|
||||
ids: string[],
|
||||
): Promise<DocumentWithMemories[]> {
|
||||
const uniqueIds = dedupe(ids)
|
||||
if (uniqueIds.length === 0) return []
|
||||
|
||||
const fetchBy = async (by: "id" | "customId", requestedIds: string[]) => {
|
||||
const response = await $fetch("@post/documents/documents/by-ids", {
|
||||
body: {
|
||||
ids: requestedIds,
|
||||
by,
|
||||
},
|
||||
disableValidation: true,
|
||||
})
|
||||
const result = response as {
|
||||
error?: { message?: string } | null
|
||||
data?: { documents?: DocumentWithMemories[] } | null
|
||||
}
|
||||
if (result.error) {
|
||||
throw new Error("Failed to fetch source documents", {
|
||||
cause: result.error,
|
||||
})
|
||||
}
|
||||
return result.data?.documents ?? []
|
||||
}
|
||||
|
||||
const byIdDocs = await fetchBy("id", uniqueIds)
|
||||
const seen = new Set<string>()
|
||||
const foundLookup = new Set<string>()
|
||||
for (const doc of byIdDocs) {
|
||||
if (doc.id) {
|
||||
seen.add(doc.id)
|
||||
foundLookup.add(doc.id)
|
||||
}
|
||||
if (doc.customId) foundLookup.add(doc.customId)
|
||||
}
|
||||
|
||||
const unresolved = uniqueIds.filter((id) => !foundLookup.has(id))
|
||||
const byCustomDocs =
|
||||
unresolved.length > 0 ? await fetchBy("customId", unresolved) : []
|
||||
const merged = [...byIdDocs]
|
||||
for (const doc of byCustomDocs) {
|
||||
if (doc.id && !seen.has(doc.id)) {
|
||||
seen.add(doc.id)
|
||||
merged.push(doc)
|
||||
}
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
export function mapDocumentsByKnownIds(
|
||||
documents: DocumentWithMemories[],
|
||||
): Map<string, DocumentWithMemories> {
|
||||
const map = new Map<string, DocumentWithMemories>()
|
||||
for (const doc of documents) {
|
||||
if (doc.id) map.set(doc.id, doc)
|
||||
if (doc.customId) map.set(doc.customId, doc)
|
||||
}
|
||||
return map
|
||||
}
|
||||
|
||||
export function getDocumentSourceUrl(
|
||||
document: Pick<DocumentWithMemories, "type" | "url"> & {
|
||||
customId?: string | null
|
||||
},
|
||||
) {
|
||||
const url = document.url ?? null
|
||||
const googleDocTypes: Record<string, { prefix: string; apiPattern: RegExp }> =
|
||||
{
|
||||
google_doc: {
|
||||
prefix: "https://docs.google.com/document/d/",
|
||||
apiPattern: /docs\.googleapis\.com\/v1\/documents\/([A-Za-z0-9_-]+)/,
|
||||
},
|
||||
google_sheet: {
|
||||
prefix: "https://docs.google.com/spreadsheets/d/",
|
||||
apiPattern:
|
||||
/sheets\.googleapis\.com\/v4\/spreadsheets\/([A-Za-z0-9_-]+)/,
|
||||
},
|
||||
google_slide: {
|
||||
prefix: "https://docs.google.com/presentation/d/",
|
||||
apiPattern:
|
||||
/slides\.googleapis\.com\/v1\/presentations\/([A-Za-z0-9_-]+)/,
|
||||
},
|
||||
}
|
||||
const googleDoc = document.type ? googleDocTypes[document.type] : undefined
|
||||
if (!googleDoc) return url
|
||||
if (document.customId) return `${googleDoc.prefix}${document.customId}/edit`
|
||||
const apiId = url?.match(googleDoc.apiPattern)?.[1]
|
||||
return apiId ? `${googleDoc.prefix}${apiId}/edit` : url
|
||||
}
|
||||
92
apps/web/lib/source-annotations.test.ts
Normal file
92
apps/web/lib/source-annotations.test.ts
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import { describe, expect, it } from "bun:test"
|
||||
import {
|
||||
isSafeSourceId,
|
||||
parseSourceAnnotatedMarkdown,
|
||||
stripSourceMarkup,
|
||||
} from "./source-annotations"
|
||||
|
||||
describe("source annotation parsing", () => {
|
||||
it("turns allowed response source spans into internal citation links", () => {
|
||||
const parsed = parseSourceAnnotatedMarkdown(
|
||||
'Alpha <response source="S1">Beta [x]</response> Gamma',
|
||||
new Set(["S1"]),
|
||||
)
|
||||
|
||||
expect(parsed.markdown).toBe("Alpha [Beta \\[x\\]](#sm-source:S1) Gamma")
|
||||
})
|
||||
|
||||
it("renders repeated allowed citations as separate internal links", () => {
|
||||
const parsed = parseSourceAnnotatedMarkdown(
|
||||
'<response source="S1">First</response> and <response source="S1">Second</response>',
|
||||
new Set(["S1"]),
|
||||
)
|
||||
|
||||
expect(parsed.markdown).toBe(
|
||||
"[First](#sm-source:S1) and [Second](#sm-source:S1)",
|
||||
)
|
||||
})
|
||||
|
||||
it("renders unknown or unsafe source ids as plain text", () => {
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown(
|
||||
'Unknown <response source="missing">plain</response>',
|
||||
new Set(["S1"]),
|
||||
).markdown,
|
||||
).toBe("Unknown plain")
|
||||
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown(
|
||||
'Unsafe <response source="bad/id">plain</response>',
|
||||
new Set(["bad/id"]),
|
||||
).markdown,
|
||||
).toBe("Unsafe plain")
|
||||
})
|
||||
|
||||
it("keeps unclosed, incomplete, nested, and malformed source markup safe", () => {
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown(
|
||||
'Lead <response source="S1">unfinished answer',
|
||||
new Set(["S1"]),
|
||||
).markdown,
|
||||
).toBe("Lead unfinished answer")
|
||||
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown("Lead <response", new Set(["S1"])).markdown,
|
||||
).toBe("Lead ")
|
||||
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown(
|
||||
'Outer <response source="S1">A <response source="S2">B</response> C</response>',
|
||||
new Set(["S1", "S2"]),
|
||||
).markdown,
|
||||
).toBe("Outer A B C")
|
||||
|
||||
expect(
|
||||
parseSourceAnnotatedMarkdown(
|
||||
"Malformed <response source=S1>plain</response>",
|
||||
new Set(["S1"]),
|
||||
).markdown,
|
||||
).toBe("Malformed plain")
|
||||
})
|
||||
|
||||
it("does not mutate inline code or fenced code blocks", () => {
|
||||
const input =
|
||||
'`<response source="S1">code</response>`\n```\n<response source="S1">fenced</response>\n```'
|
||||
|
||||
expect(parseSourceAnnotatedMarkdown(input, new Set(["S1"])).markdown).toBe(
|
||||
input,
|
||||
)
|
||||
})
|
||||
|
||||
it("strips source markup for copy text", () => {
|
||||
expect(
|
||||
stripSourceMarkup('Alpha <response source="S1">Beta</response>'),
|
||||
).toBe("Alpha Beta")
|
||||
})
|
||||
|
||||
it("allows only source ids that are safe in internal fragments", () => {
|
||||
expect(isSafeSourceId("S1._:-")).toBe(true)
|
||||
expect(isSafeSourceId("bad/id")).toBe(false)
|
||||
expect(isSafeSourceId("bad space")).toBe(false)
|
||||
})
|
||||
})
|
||||
194
apps/web/lib/source-annotations.ts
Normal file
194
apps/web/lib/source-annotations.ts
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
export type ParsedSourceAnnotations = {
|
||||
markdown: string
|
||||
}
|
||||
|
||||
const RESPONSE_OPEN_PREFIX = "<response"
|
||||
const RESPONSE_CLOSE_TAG = "</response>"
|
||||
const SOURCE_ATTR_PREFIX = 'source="'
|
||||
const SAFE_SOURCE_ID_RE = /^[A-Za-z0-9_.:-]+$/
|
||||
|
||||
export function isSafeSourceId(id: string): boolean {
|
||||
return id.length > 0 && SAFE_SOURCE_ID_RE.test(id)
|
||||
}
|
||||
|
||||
function escapeMarkdownLinkText(text: string): string {
|
||||
return text.replace(/([\\[\]])/g, "\\$1").replace(/\n/g, " ")
|
||||
}
|
||||
|
||||
function parseOpeningTag(
|
||||
text: string,
|
||||
index: number,
|
||||
): { end: number; sourceId: string } | null | "incomplete" {
|
||||
if (!text.startsWith(RESPONSE_OPEN_PREFIX, index)) return null
|
||||
|
||||
const tagEnd = text.indexOf(">", index + RESPONSE_OPEN_PREFIX.length)
|
||||
if (tagEnd === -1) return "incomplete"
|
||||
|
||||
const rawTag = text.slice(index, tagEnd + 1)
|
||||
const inside = rawTag.slice(1, -1).trim()
|
||||
if (!inside.startsWith("response")) return null
|
||||
|
||||
let cursor = "response".length
|
||||
while (
|
||||
inside[cursor] === " " ||
|
||||
inside[cursor] === "\t" ||
|
||||
inside[cursor] === "\n" ||
|
||||
inside[cursor] === "\r"
|
||||
)
|
||||
cursor++
|
||||
if (!inside.startsWith(SOURCE_ATTR_PREFIX, cursor)) return null
|
||||
cursor += SOURCE_ATTR_PREFIX.length
|
||||
|
||||
const sourceEnd = inside.indexOf('"', cursor)
|
||||
if (sourceEnd === -1) return null
|
||||
const sourceId = inside.slice(cursor, sourceEnd)
|
||||
cursor = sourceEnd + 1
|
||||
while (
|
||||
inside[cursor] === " " ||
|
||||
inside[cursor] === "\t" ||
|
||||
inside[cursor] === "\n" ||
|
||||
inside[cursor] === "\r"
|
||||
)
|
||||
cursor++
|
||||
if (cursor !== inside.length) return null
|
||||
if (!isSafeSourceId(sourceId)) return null
|
||||
|
||||
return { end: tagEnd + 1, sourceId }
|
||||
}
|
||||
|
||||
function advanceCodeState(
|
||||
text: string,
|
||||
index: number,
|
||||
state: { inFence: boolean; inInlineCode: boolean; lineStart: boolean },
|
||||
): boolean {
|
||||
if (state.lineStart && text.startsWith("```", index)) {
|
||||
state.inFence = !state.inFence
|
||||
return true
|
||||
}
|
||||
|
||||
if (!state.inFence && text[index] === "`") {
|
||||
state.inInlineCode = !state.inInlineCode
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
function appendChar(
|
||||
text: string,
|
||||
index: number,
|
||||
output: string[],
|
||||
state: { lineStart: boolean },
|
||||
) {
|
||||
const ch = text[index] ?? ""
|
||||
output.push(ch)
|
||||
state.lineStart = ch === "\n"
|
||||
}
|
||||
|
||||
export function parseSourceAnnotatedMarkdown(
|
||||
text: string,
|
||||
allowedSourceIds: ReadonlySet<string>,
|
||||
): ParsedSourceAnnotations {
|
||||
const output: string[] = []
|
||||
const codeState = { inFence: false, inInlineCode: false, lineStart: true }
|
||||
|
||||
let i = 0
|
||||
while (i < text.length) {
|
||||
if (advanceCodeState(text, i, codeState)) {
|
||||
appendChar(text, i, output, codeState)
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
if (
|
||||
!codeState.inFence &&
|
||||
!codeState.inInlineCode &&
|
||||
text.startsWith(RESPONSE_OPEN_PREFIX, i)
|
||||
) {
|
||||
const opening = parseOpeningTag(text, i)
|
||||
if (opening === "incomplete") {
|
||||
break
|
||||
}
|
||||
|
||||
if (opening) {
|
||||
const closeIndex = text.indexOf(RESPONSE_CLOSE_TAG, opening.end)
|
||||
if (closeIndex === -1) {
|
||||
output.push(stripSourceMarkup(text.slice(opening.end)))
|
||||
break
|
||||
}
|
||||
|
||||
const inner = text.slice(opening.end, closeIndex)
|
||||
const hasNested =
|
||||
inner.includes(RESPONSE_OPEN_PREFIX) ||
|
||||
inner.includes(RESPONSE_CLOSE_TAG)
|
||||
const isAllowed = allowedSourceIds.has(opening.sourceId)
|
||||
|
||||
if (hasNested) {
|
||||
const outerCloseIndex = text.indexOf(
|
||||
RESPONSE_CLOSE_TAG,
|
||||
closeIndex + RESPONSE_CLOSE_TAG.length,
|
||||
)
|
||||
const fallbackEnd =
|
||||
outerCloseIndex === -1 ? closeIndex : outerCloseIndex
|
||||
output.push(stripSourceMarkup(text.slice(opening.end, fallbackEnd)))
|
||||
i = fallbackEnd + RESPONSE_CLOSE_TAG.length
|
||||
continue
|
||||
}
|
||||
|
||||
const plainInner = stripSourceMarkup(inner)
|
||||
if (isAllowed && plainInner.trim().length > 0) {
|
||||
output.push(
|
||||
`[${escapeMarkdownLinkText(plainInner)}](#sm-source:${encodeURIComponent(opening.sourceId)})`,
|
||||
)
|
||||
} else {
|
||||
output.push(plainInner)
|
||||
}
|
||||
|
||||
i = closeIndex + RESPONSE_CLOSE_TAG.length
|
||||
codeState.lineStart =
|
||||
output.length === 0 ||
|
||||
output[output.length - 1]?.endsWith("\n") === true
|
||||
continue
|
||||
}
|
||||
|
||||
const nextClose = text.indexOf(RESPONSE_CLOSE_TAG, i)
|
||||
if (nextClose !== -1) {
|
||||
const tagEnd = text.indexOf(">", i)
|
||||
if (tagEnd !== -1 && tagEnd < nextClose) {
|
||||
output.push(stripSourceMarkup(text.slice(tagEnd + 1, nextClose)))
|
||||
i = nextClose + RESPONSE_CLOSE_TAG.length
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
appendChar(text, i, output, codeState)
|
||||
i++
|
||||
}
|
||||
|
||||
return { markdown: output.join("") }
|
||||
}
|
||||
|
||||
export function stripSourceMarkup(text: string): string {
|
||||
let output = ""
|
||||
let i = 0
|
||||
|
||||
while (i < text.length) {
|
||||
if (text.startsWith(RESPONSE_CLOSE_TAG, i)) {
|
||||
i += RESPONSE_CLOSE_TAG.length
|
||||
continue
|
||||
}
|
||||
|
||||
if (text.startsWith(RESPONSE_OPEN_PREFIX, i)) {
|
||||
const tagEnd = text.indexOf(">", i + RESPONSE_OPEN_PREFIX.length)
|
||||
if (tagEnd === -1) break
|
||||
i = tagEnd + 1
|
||||
continue
|
||||
}
|
||||
|
||||
output += text[i]
|
||||
i++
|
||||
}
|
||||
|
||||
return output
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue