Skip to content

Render image references from retrieved evidence #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
suguanYang wants to merge 1 commit into from
Closed

Render image references from retrieved evidence #94

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions src/components/chat-message-list.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,52 @@ describe("ChatMessageList", () => {
).toBeNull();
});

it("does not hide image cards when source links dedupe the same section", () => {
render(
React.createElement(ChatMessageList, {
messages: [
{
id: "assistant_1",
role: "assistant",
content: "这里是相关身份证明图片。",
citations: [
{
chunkType: "text",
score: 0.9,
source: {
documentId: "doc_1",
sourceFileName: "商务标文件.pdf",
sectionPath: "二、法定代表人身份证明",
},
},
{
chunkType: "image",
score: 0.9,
assetUrl: "https://blob.example/images/image-6-id-front.jpg",
source: {
documentId: "doc_1",
sourceFileName: "商务标文件.pdf",
sectionPath: "二、法定代表人身份证明",
},
},
],
},
],
}),
);

expect(
screen.getByRole("img", {
name: "商务标文件.pdf · 二、法定代表人身份证明",
}),
).toBeTruthy();
expect(
screen.getAllByRole("button", {
name: "Open source 商务标文件.pdf · 二、法定代表人身份证明",
}),
).toHaveLength(1);
});

it("shows thinking progress after existing messages while sending", () => {
render(
React.createElement(ChatMessageList, {
Expand Down
21 changes: 15 additions & 6 deletions src/components/chat-message-list.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,10 @@ function MessageBubble({
message,
sourceTitlesByDocumentId,
);
const displayImageCitations = getDisplayImageCitations(displayCitations);
const displayImageCitations = getDisplayImageCitations(
message,
sourceTitlesByDocumentId,
);

return (
<div className="flex min-w-0 flex-col items-start">
Expand Down Expand Up @@ -336,18 +339,24 @@ function getDisplayCitations(
}

function getDisplayImageCitations(
citations: readonly DisplayCitation[],
message: ChatMessageView,
sourceTitlesByDocumentId: Readonly<Record<string, string>>,
): readonly DisplayImageCitation[] {
const seenAssetUrls = new Set<string>();
const imageCitations: DisplayImageCitation[] = [];

for (const citation of citations) {
const assetUrl = getTrimmedCitationField(citation.citation.assetUrl);
if (!assetUrl || !isImageCitation(citation.citation, assetUrl)) continue;
for (const [index, citation] of (message.citations ?? []).entries()) {
const assetUrl = getTrimmedCitationField(citation.assetUrl);
if (!assetUrl || !isImageCitation(citation, assetUrl)) continue;
if (seenAssetUrls.has(assetUrl)) continue;

seenAssetUrls.add(assetUrl);
imageCitations.push({ ...citation, assetUrl });
imageCitations.push({
citation,
citationId: chatPanelModel.getCitationId(message.id, index),
label: chatPanelModel.getCitationLabel(citation, sourceTitlesByDocumentId),
assetUrl,
});
}

return imageCitations;
Expand Down
71 changes: 71 additions & 0 deletions src/domains/chat/index.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,77 @@ describe("answerQuestionWithRetrieval", () => {
]);
});

it("turns retrieved evidence image filenames into image citations", async () => {
const result = makeRetrievalResult({
content: "This section contains identity proof attachments.",
source: {
documentId: "doc_identity",
sourceFileName: "document-generated.pdf",
sectionPath: "二、法定代表人身份证明",
},
});
const retrieval = {
query: vi.fn().mockResolvedValue({
results: [result],
evidenceText:
"[image-6-中华人民共和国居民身份证.jpg]\n[image-7-中国居民身份证.jpg]",
referencedChunks: [],
namespace: "notebook-workspace",
query: "公民身份证明 图片",
routerUsed: "workflow_single_step",
answerText: null,
}),
};
const generateAnswer = vi.fn().mockResolvedValue("这里是相关身份证明图片。");
const generateRetrievalQuery = vi.fn().mockResolvedValue("公民身份证明 图片");
const loadSourceAssetUrls = vi.fn().mockResolvedValue({
"images/image-6-中华人民共和国居民身份证.jpg":
"https://blob.example/images/image-6-id-front.jpg",
"images/image-7-中国居民身份证.jpg":
"https://blob.example/images/image-7-id-back.jpg",
});

const answer = await Effect.runPromise(
answerQuestionWithRetrieval({
question: "请发送几张关于公民身份的图片给我",
namespace: "notebook-workspace",
sources: [
makeSource({
id: "source_identity",
title: "商务标文件.pdf",
knowhereDocumentId: "doc_identity",
}),
],
excludedSourceIds: [],
retrieval,
generateRetrievalQuery,
generateAnswer,
loadSourceAssetUrls,
messages: [],
}),
);

expect(generateAnswer).toHaveBeenCalledWith({
question: "请发送几张关于公民身份的图片给我",
retrievalQuery: "公民身份证明 图片",
messages: [],
evidenceText:
"[image-6-中华人民共和国居民身份证.jpg]\n[image-7-中国居民身份证.jpg]",
mediaAssetContext:
"- 商务标文件.pdf / images/image-6-中华人民共和国居民身份证.jpg: https://blob.example/images/image-6-id-front.jpg\n" +
"- 商务标文件.pdf / images/image-7-中国居民身份证.jpg: https://blob.example/images/image-7-id-back.jpg",
});
expect(answer.citations.map((citation) => citation.assetUrl)).toEqual([
undefined,
"https://blob.example/images/image-6-id-front.jpg",
"https://blob.example/images/image-7-id-back.jpg",
]);
expect(answer.citations.slice(1).map((citation) => citation.chunkType)).toEqual([
"image",
"image",
]);
});

it("returns a deterministic no-results answer without calling the model", async () => {
const retrieval = {
query: vi.fn().mockResolvedValue({
Expand Down
1 change: 1 addition & 0 deletions src/domains/chat/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ export const answerQuestionWithRetrieval = (
results: useNotebookSourceTitles(response.results, input.sources),
sources: input.sources,
loadSourceAssetUrls: input.loadSourceAssetUrls,
evidenceText,
}),
)
const mediaAssetContext = formatRetrievedMediaAssetContext(results)
Expand Down
47 changes: 47 additions & 0 deletions src/domains/chat/media-assets.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,53 @@ describe("chat media assets", () => {
)
})

it("adds image citation results for asset filenames that only appear in evidence text", async () => {
const loadSourceAssetUrls = vi.fn().mockResolvedValue({
"images/image-6-中华人民共和国居民身份证.jpg":
"https://blob.example/images/image-6-id-front.jpg",
"images/image-7-中国居民身份证.jpg":
"https://blob.example/images/image-7-id-back.jpg",
})

const results = await enrichRetrievalResultsWithAssetUrls({
results: [
makeRetrievalResult({
content: "The section contains citizen identity proof copies.",
source: {
documentId: "doc_identity",
sourceFileName: "商务标文件.pdf",
sectionPath: "二、法定代表人身份证明",
},
}),
],
sources: [
makeSource({
id: "source_identity",
title: "商务标文件.pdf",
knowhereDocumentId: "doc_identity",
}),
],
loadSourceAssetUrls,
evidenceText:
"[image-6-中华人民共和国居民身份证.jpg]\n[image-7-中国居民身份证.jpg]",
})

expect(results).toHaveLength(3)
expect(results[0]?.assetUrl).toBeUndefined()
expect(results.slice(1).map((result) => result.assetUrl)).toEqual([
"https://blob.example/images/image-6-id-front.jpg",
"https://blob.example/images/image-7-id-back.jpg",
])
expect(results.slice(1).map((result) => result.chunkType)).toEqual([
"image",
"image",
])
expect(results.slice(1).map((result) => result.source.sectionPath)).toEqual([
"images/image-6-中华人民共和国居民身份证.jpg",
"images/image-7-中国居民身份证.jpg",
])
})

it("formats a bounded media asset context for the grounded prompt", () => {
const context = formatRetrievedMediaAssetContext([
makeRetrievalResult({
Expand Down
Loading
Loading