diff --git a/apps/web/components/document-modal/graph-list-memories.tsx b/apps/web/components/document-modal/graph-list-memories.tsx index dc62cfba..c8f4ae8d 100644 --- a/apps/web/components/document-modal/graph-list-memories.tsx +++ b/apps/web/components/document-modal/graph-list-memories.tsx @@ -2,6 +2,7 @@ import { useState } from "react" import { cn } from "@lib/utils" import { Tabs, TabsList, TabsTrigger, TabsContent } from "@ui/components/tabs" import { MemoryGraph } from "../memory-graph" +import { useProject } from "@/stores" export interface MemoryEntry { id: string @@ -172,6 +173,7 @@ export function GraphListMemories({ memoryEntries: MemoryEntry[] documentId?: string }) { + const { effectiveContainerTags } = useProject() const [expandedMemories, setExpandedMemories] = useState>( new Set(), ) @@ -263,7 +265,10 @@ export function GraphListMemories({
diff --git a/apps/web/components/memory-graph/hooks/use-graph-api.ts b/apps/web/components/memory-graph/hooks/use-graph-api.ts index 94c08e33..7991f600 100644 --- a/apps/web/components/memory-graph/hooks/use-graph-api.ts +++ b/apps/web/components/memory-graph/hooks/use-graph-api.ts @@ -13,6 +13,7 @@ const PAGE_SIZE = 100 interface UseGraphApiOptions { containerTags?: string[] + documentIds?: string[] enabled?: boolean } @@ -81,7 +82,20 @@ function toGraphMemory(mem: ApiMemoryEntry): GraphApiMemory { } } -function toGraphDocument(doc: ApiDocument): GraphApiDocument { +function toGraphDocument( + doc: ApiDocument, + containerTags?: string[], +): GraphApiDocument { + const allowedContainerTags = new Set(containerTags?.filter(Boolean) ?? []) + const memoryEntries = + allowedContainerTags.size > 0 + ? doc.memoryEntries.filter( + (mem) => + mem.spaceContainerTag != null && + allowedContainerTags.has(mem.spaceContainerTag), + ) + : doc.memoryEntries + return { id: doc.id, title: doc.title, @@ -89,12 +103,15 @@ function toGraphDocument(doc: ApiDocument): GraphApiDocument { documentType: doc.type, createdAt: doc.createdAt, updatedAt: doc.updatedAt, - memories: doc.memoryEntries.map(toGraphMemory), + memories: memoryEntries.map(toGraphMemory), } } export function useGraphApi(options: UseGraphApiOptions = {}) { - const { containerTags, enabled = true } = options + const { containerTags, documentIds, enabled = true } = options + const filteredDocumentIds = documentIds?.filter(Boolean) + const hasDocumentIds = + filteredDocumentIds != null && filteredDocumentIds.length > 0 const { data, @@ -104,19 +121,33 @@ export function useGraphApi(options: UseGraphApiOptions = {}) { hasNextPage, fetchNextPage, } = useInfiniteQuery({ - queryKey: ["documents-with-memories", containerTags, []], + queryKey: [ + "documents-with-memories", + containerTags, + [], + filteredDocumentIds, + ], initialPageParam: 1, queryFn: async ({ pageParam }) => { - const response = await $fetch("@post/documents/documents", { - body: { - page: pageParam as number, - limit: PAGE_SIZE, - sort: "createdAt", - order: "desc", - containerTags, - }, - disableValidation: true, - }) + const response = hasDocumentIds + ? await $fetch("@post/documents/documents/by-ids", { + body: { + ids: filteredDocumentIds, + by: "id", + containerTags, + }, + disableValidation: true, + }) + : await $fetch("@post/documents/documents", { + body: { + page: pageParam as number, + limit: PAGE_SIZE, + sort: "createdAt", + order: "desc", + containerTags, + }, + disableValidation: true, + }) if (response.error) { throw new Error(response.error?.message || "Failed to fetch documents") @@ -134,8 +165,10 @@ export function useGraphApi(options: UseGraphApiOptions = {}) { const documents = useMemo(() => { if (!data?.pages) return [] - return data.pages.flatMap((page) => page.documents.map(toGraphDocument)) - }, [data]) + return data.pages.flatMap((page) => + page.documents.map((doc) => toGraphDocument(doc, containerTags)), + ) + }, [data, containerTags]) const totalCount = data?.pages[0]?.pagination.totalItems ?? 0 diff --git a/apps/web/components/memory-graph/memory-graph-wrapper.tsx b/apps/web/components/memory-graph/memory-graph-wrapper.tsx index ae2bf327..0b2ca826 100644 --- a/apps/web/components/memory-graph/memory-graph-wrapper.tsx +++ b/apps/web/components/memory-graph/memory-graph-wrapper.tsx @@ -29,6 +29,7 @@ export function MemoryGraph({ error: externalError = null, variant = "console", containerTags, + documentIds, maxNodes, canvasRef, ...rest @@ -57,6 +58,7 @@ export function MemoryGraph({ totalCount, } = useGraphApi({ containerTags, + documentIds, enabled: containerSize.width > 0 && containerSize.height > 0, })