Skip to content

Commit 84b1473

Browse files
suguanYangsuguanYang
committed
skip full chunk fetch for source-only citations
1 parent c7e93d2 commit 84b1473

4 files changed

Lines changed: 329 additions & 14 deletions

File tree

‎src/components/workspace-citation-focus.test.ts‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,82 @@ describe("useWorkspaceCitationFocus", () => {
9393
requestId: 1,
9494
});
9595
});
96+
97+
it("opens the source without fetching chunks when the citation has no exact target hint", async () => {
98+
const fetchChunks = vi.fn(async () => [prefetchedChunk]);
99+
const selectSource = vi.fn();
100+
const sourceOnlyCitation: ChatCitationView = {
101+
chunkType: "text",
102+
score: 0.5,
103+
source: {
104+
documentId: "document_1",
105+
sourceFileName: "Contract.pdf",
106+
sectionPath: "Root",
107+
},
108+
};
109+
110+
const { result } = renderHook(() =>
111+
useWorkspaceCitationFocus({
112+
fetchChunks,
113+
initialPrefetchedChunksBySourceId: {
114+
source_1: [prefetchedChunk],
115+
},
116+
onSelectSource: selectSource,
117+
selectedSourceId: null,
118+
sources: [readySource],
119+
}),
120+
{ wrapper: createSWRWrapper },
121+
);
122+
123+
await act(async () => {
124+
await result.current.handleCitationClick(
125+
sourceOnlyCitation,
126+
"message_1:0",
127+
);
128+
});
129+
130+
expect(fetchChunks).not.toHaveBeenCalled();
131+
expect(selectSource).toHaveBeenLastCalledWith("source_1");
132+
expect(result.current.prefetchedChunksBySourceId).toEqual({});
133+
expect(result.current.focusedChunk.chunkId).toBeNull();
134+
expect(result.current.pendingCitationId).toBeNull();
135+
});
136+
137+
it("reuses cached chunks for a different source without refetching", async () => {
138+
const fetchChunks = vi.fn(async () => [prefetchedChunk]);
139+
const selectSource = vi.fn();
140+
const otherSource: SourceView = {
141+
id: "source_2",
142+
title: "Other.pdf",
143+
mimeType: "application/pdf",
144+
status: "ready",
145+
documentId: "document_2",
146+
};
147+
148+
const { result } = renderHook(() =>
149+
useWorkspaceCitationFocus({
150+
fetchChunks,
151+
initialPrefetchedChunksBySourceId: {
152+
source_1: [prefetchedChunk],
153+
},
154+
onSelectSource: selectSource,
155+
selectedSourceId: "source_2",
156+
sources: [readySource, otherSource],
157+
}),
158+
{ wrapper: createSWRWrapper },
159+
);
160+
161+
await act(async () => {
162+
await result.current.handleCitationClick(citation, "message_1:0");
163+
});
164+
165+
expect(fetchChunks).not.toHaveBeenCalled();
166+
expect(selectSource).toHaveBeenLastCalledWith("source_1");
167+
expect(result.current.focusedChunk.chunkId).toBe("chunk_1");
168+
expect(Object.keys(result.current.prefetchedChunksBySourceId)).toContain(
169+
"source_1",
170+
);
171+
});
96172
});
97173

98174
function createSWRWrapper({

‎src/components/workspace-citation-focus.ts‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ export function useWorkspaceCitationFocus({
8585
onSelectSource(sourceId)
8686
if (sourceId) {
8787
setPrefetchedChunksBySourceId((current) =>
88-
removeRecordKey(current, sourceId),
88+
workspaceCitationState.removePrefetchedChunks(current, sourceId),
8989
)
9090
}
9191
requestChunkFocus(null)
@@ -119,6 +119,37 @@ export function useWorkspaceCitationFocus({
119119
return
120120
}
121121

122+
if (!workspaceCitationState.hasExactCitationTargetHint(citation)) {
123+
setPrefetchedChunksBySourceId((current) =>
124+
workspaceCitationState.removePrefetchedChunks(current, source.id),
125+
)
126+
if (selectedSourceId !== source.id) onSelectSource(source.id)
127+
requestChunkFocus(null)
128+
return
129+
}
130+
131+
const cachedChunks = prefetchedChunksBySourceId[source.id]
132+
if (cachedChunks) {
133+
const cachedFocusId =
134+
workspaceCitationState.getLoadedCitationChunkId({
135+
citation,
136+
selectedSourceId: source.id,
137+
sourceId: source.id,
138+
selectedChunks: cachedChunks,
139+
hasMoreSelectedChunks: false,
140+
})
141+
setPrefetchedChunksBySourceId((current) =>
142+
workspaceCitationState.upsertPrefetchedChunks(
143+
current,
144+
source.id,
145+
cachedChunks,
146+
),
147+
)
148+
if (selectedSourceId !== source.id) onSelectSource(source.id)
149+
requestChunkFocus(cachedFocusId)
150+
return
151+
}
152+
122153
requestChunkFocus(null)
123154
const chunks = await fetchChunks(source.id)
124155
setPrefetchedChunksBySourceId((current) =>
@@ -148,6 +179,7 @@ export function useWorkspaceCitationFocus({
148179
fetchChunks,
149180
hasMoreSelectedChunks,
150181
onSelectSource,
182+
prefetchedChunksBySourceId,
151183
requestChunkFocus,
152184
selectedChunks,
153185
selectedSourceId,
@@ -170,12 +202,3 @@ export function useWorkspaceCitationFocus({
170202
selectedSource,
171203
}
172204
}
173-
174-
function removeRecordKey<TValue>(
175-
record: Readonly<Record<string, TValue>>,
176-
key: string,
177-
): Record<string, TValue> {
178-
const nextRecord = { ...record }
179-
delete nextRecord[key]
180-
return nextRecord
181-
}

‎src/components/workspace-citation-state.test.ts‎

Lines changed: 164 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,23 @@
11
import { describe, expect, it } from "vitest"
22

3-
import { workspaceCitationState } from "./workspace-citation-state"
3+
import {
4+
maxPrefetchedChunkSources,
5+
workspaceCitationState,
6+
} from "./workspace-citation-state"
47
import type { ChatCitationView } from "@/domains/chat/types"
58
import type { ParsedChunkView } from "@/domains/chunks/types"
69
import type { SourceView } from "@/domains/sources/types"
710

11+
function makeChunk(chunkId: string, documentId: string): ParsedChunkView {
12+
return {
13+
chunkId,
14+
documentId,
15+
type: "text",
16+
content: `${chunkId} content`,
17+
sourceTitle: `${documentId}.pdf`,
18+
}
19+
}
20+
821
describe("workspaceCitationState", () => {
922
it("finds the Source and loaded Parsed Chunk for a Citation", () => {
1023
const source: SourceView = {
@@ -77,4 +90,154 @@ describe("workspaceCitationState", () => {
7790
}),
7891
).toBeNull()
7992
})
93+
94+
describe("hasExactCitationTargetHint", () => {
95+
it("returns true when the citation has content text", () => {
96+
const citation: ChatCitationView = {
97+
chunkType: "text",
98+
score: 0.5,
99+
content: "Revenue grew this quarter.",
100+
source: { documentId: "document_1", sourceFileName: "Contract.pdf" },
101+
}
102+
103+
expect(
104+
workspaceCitationState.hasExactCitationTargetHint(citation),
105+
).toBe(true)
106+
})
107+
108+
it("returns true when the citation has a meaningful section path", () => {
109+
const citation: ChatCitationView = {
110+
chunkType: "text",
111+
score: 0.5,
112+
source: {
113+
documentId: "document_1",
114+
sectionPath: "Revenue",
115+
},
116+
}
117+
118+
expect(
119+
workspaceCitationState.hasExactCitationTargetHint(citation),
120+
).toBe(true)
121+
})
122+
123+
it("returns false for source-only citations with no useful target hint", () => {
124+
const citation: ChatCitationView = {
125+
chunkType: "text",
126+
score: 0.5,
127+
source: {
128+
documentId: "document_1",
129+
sourceFileName: "Contract.pdf",
130+
sectionPath: "Root",
131+
},
132+
}
133+
134+
expect(
135+
workspaceCitationState.hasExactCitationTargetHint(citation),
136+
).toBe(false)
137+
})
138+
139+
it("returns false for empty content and missing section path", () => {
140+
const citation: ChatCitationView = {
141+
chunkType: "text",
142+
score: 0.5,
143+
content: " ",
144+
source: { documentId: "document_1" },
145+
}
146+
147+
expect(
148+
workspaceCitationState.hasExactCitationTargetHint(citation),
149+
).toBe(false)
150+
})
151+
})
152+
153+
describe("upsertPrefetchedChunks bounded LRU", () => {
154+
it("keeps insertion order so the newest source appears last", () => {
155+
const next = workspaceCitationState.upsertPrefetchedChunks(
156+
workspaceCitationState.upsertPrefetchedChunks({}, "source_a", [
157+
makeChunk("a", "doc_a"),
158+
]),
159+
"source_b",
160+
[makeChunk("b", "doc_b")],
161+
)
162+
163+
expect(Object.keys(next)).toEqual(["source_a", "source_b"])
164+
})
165+
166+
it("refreshes recency for an existing source by moving it to the end", () => {
167+
const seed = workspaceCitationState.upsertPrefetchedChunks(
168+
workspaceCitationState.upsertPrefetchedChunks({}, "source_a", [
169+
makeChunk("a", "doc_a"),
170+
]),
171+
"source_b",
172+
[makeChunk("b", "doc_b")],
173+
)
174+
175+
const refreshed = workspaceCitationState.upsertPrefetchedChunks(
176+
seed,
177+
"source_a",
178+
[makeChunk("a", "doc_a"), makeChunk("a2", "doc_a")],
179+
)
180+
181+
expect(Object.keys(refreshed)).toEqual(["source_b", "source_a"])
182+
expect(refreshed["source_a"]?.length).toBe(2)
183+
})
184+
185+
it(`evicts the oldest entries when more than ${maxPrefetchedChunkSources} sources are stored`, () => {
186+
const seeded = Array.from({ length: maxPrefetchedChunkSources }).reduce<
187+
Readonly<Record<string, ParsedChunkView[]>>
188+
>(
189+
(acc, _value, index) =>
190+
workspaceCitationState.upsertPrefetchedChunks(
191+
acc,
192+
`source_${index}`,
193+
[makeChunk(`chunk_${index}`, `doc_${index}`)],
194+
),
195+
{},
196+
)
197+
198+
const next = workspaceCitationState.upsertPrefetchedChunks(
199+
seeded,
200+
"source_overflow",
201+
[makeChunk("chunk_overflow", "doc_overflow")],
202+
)
203+
204+
const keys = Object.keys(next)
205+
expect(keys.length).toBe(maxPrefetchedChunkSources)
206+
expect(keys).not.toContain("source_0")
207+
expect(keys).toContain("source_overflow")
208+
expect(keys[keys.length - 1]).toBe("source_overflow")
209+
})
210+
})
211+
212+
describe("removePrefetchedChunks", () => {
213+
it("returns the same record when the source is not cached", () => {
214+
const seed = workspaceCitationState.upsertPrefetchedChunks({}, "source_a", [
215+
makeChunk("a", "doc_a"),
216+
])
217+
218+
const next = workspaceCitationState.removePrefetchedChunks(
219+
seed,
220+
"missing",
221+
)
222+
223+
expect(next).toBe(seed)
224+
})
225+
226+
it("removes only the requested source", () => {
227+
const seed = workspaceCitationState.upsertPrefetchedChunks(
228+
workspaceCitationState.upsertPrefetchedChunks({}, "source_a", [
229+
makeChunk("a", "doc_a"),
230+
]),
231+
"source_b",
232+
[makeChunk("b", "doc_b")],
233+
)
234+
235+
const next = workspaceCitationState.removePrefetchedChunks(
236+
seed,
237+
"source_a",
238+
)
239+
240+
expect(Object.keys(next)).toEqual(["source_b"])
241+
})
242+
})
80243
})

0 commit comments

Comments
 (0)