mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-08-10 17:14:14 +00:00
375 lines
10 KiB
TypeScript
375 lines
10 KiB
TypeScript
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
|
|
}
|