Add image asset citations to chat

src/components/chat-message-list.test.ts

@@ -59,6 +59,47 @@ describe("ChatMessageList", () => {
     ).toBeTruthy();
   });
 
+  it("renders image citations as viewable image attachments", () => {
+    render(
+      React.createElement(ChatMessageList, {
+        messages: [
+          {
+            id: "assistant_1",
+            role: "assistant",
+            content: "Here is the launch image.",
+            citations: [
+              {
+                chunkType: "image",
+                score: 0.9,
+                assetUrl: "https://blob.example/images/launch.jpg",
+                source: {
+                  documentId: "doc_1",
+                  sourceFileName: "spacex-s1.pdf",
+                  sectionPath: "Assets / images / launch.jpg",
+                },
+              },
+            ],
+          },
+        ],
+      }),
+    );
+
+    const image = screen.getByRole("img", {
+      name: "spacex-s1.pdf · Assets / images / launch.jpg",
+    });
+    expect(image.getAttribute("src")).toBe(
+      "https://blob.example/images/launch.jpg",
+    );
+    expect(
+      screen.queryByRole("link", {
+        name: "https://blob.example/images/launch.jpg",
+      }),
+    ).toBeNull();
+    expect(
+      screen.queryByText("https://blob.example/images/launch.jpg"),
+    ).toBeNull();
+  });
+
   it("shows thinking progress after existing messages while sending", () => {
     render(
       React.createElement(ChatMessageList, {

src/components/chat-message-list.tsx

@@ -2,7 +2,7 @@
 
 import { type CSSProperties, type ReactElement } from "react";
 import { type VirtualItem } from "@tanstack/react-virtual";
-import { MessageCircle } from "lucide-react";
+import { ImageIcon, MessageCircle } from "lucide-react";
 
 import { useChatMessageListWorkflow } from "@/components/chat-message-list-workflow";
 import { chatPanelModel } from "@/components/chat-panel-model";
@@ -19,6 +19,10 @@ type DisplayCitation = {
   readonly label: string;
 };
 
+type DisplayImageCitation = DisplayCitation & {
+  readonly assetUrl: string;
+};
+
 export type ChatMessageListProps = {
   readonly isDisabled?: boolean;
   readonly isSending?: boolean;
@@ -240,11 +244,40 @@ function MessageBubble({
     message,
     sourceTitlesByDocumentId,
   );
+  const displayImageCitations = getDisplayImageCitations(displayCitations);
 
   return (
     <div className="flex min-w-0 flex-col items-start">
       <div className="max-w-[92%] overflow-hidden rounded-2xl rounded-tl-sm border border-border/70 bg-card px-3 py-2.5 text-sm leading-relaxed text-foreground shadow-xs sm:max-w-[90%] sm:px-4 sm:py-3">
         <p className="whitespace-pre-wrap break-words">{message.content}</p>
+        {displayImageCitations.length > 0 && (
+          <div className="mt-3 border-t border-border/70 pt-2.5">
+            <p className="mb-2 flex items-center gap-1.5 text-[10px] font-bold uppercase tracking-wider text-muted-foreground">
+              <ImageIcon className="size-3" />
+              Images
+            </p>
+            <div className="grid gap-2">
+              {displayImageCitations.map(({ assetUrl, citationId, label }) => (
+                <figure
+                  key={`${citationId}-image`}
+                  className="overflow-hidden rounded-lg border border-border bg-muted/25"
+                >
+                  {/* eslint-disable-next-line @next/next/no-img-element -- Chat image citation dimensions are not known before render. */}
+                  <img
+                    src={assetUrl}
+                    alt={label}
+                    className="max-h-64 w-full object-contain"
+                  />
+                  <figcaption className="border-t border-border/70 bg-background/80 px-2.5 py-2">
+                    <span className="block break-words text-[11px] font-semibold text-foreground">
+                      {label}
+                    </span>
+                  </figcaption>
+                </figure>
+              ))}
+            </div>
+          </div>
+        )}
         {displayCitations.length > 0 && (
           <div className="mt-3 border-t border-border/70 pt-2.5">
             <p className="mb-1.5 text-[10px] font-bold uppercase tracking-wider text-muted-foreground">
@@ -302,6 +335,49 @@ function getDisplayCitations(
   return displayCitations;
 }
 
+function getDisplayImageCitations(
+  citations: readonly DisplayCitation[],
+): 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;
+    if (seenAssetUrls.has(assetUrl)) continue;
+
+    seenAssetUrls.add(assetUrl);
+    imageCitations.push({ ...citation, assetUrl });
+  }
+
+  return imageCitations;
+}
+
+function isImageCitation(
+  citation: ChatCitationView,
+  assetUrl: string,
+): boolean {
+  return (
+    citation.chunkType.toLowerCase() === "image" ||
+    hasImageFileExtension(assetUrl)
+  );
+}
+
+function hasImageFileExtension(assetUrl: string): boolean {
+  const pathname = getUrlPathname(assetUrl).toLowerCase();
+  return [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"].some(
+    (extension) => pathname.endsWith(extension),
+  );
+}
+
+function getUrlPathname(assetUrl: string): string {
+  try {
+    return new URL(assetUrl).pathname;
+  } catch {
+    return assetUrl.split("?")[0] ?? assetUrl;
+  }
+}
+
 function getCitationDisplayKey(
   citation: ChatCitationView,
   label: string,

src/domains/chat/contracts.ts

@@ -2,6 +2,7 @@ import type { RetrievalQueryParams, RetrievalQueryResponse } from "@ontos-ai/kno
 
 import type { Source } from "@/infrastructure/db/schema"
 import type { ChatCitationView } from "@/domains/chat/types"
+import type { LoadSourceAssetUrls } from "./media-assets"
 
 export type RetrievalClient = {
   query(params: RetrievalQueryParams): Promise<RetrievalQueryResponse>
@@ -25,6 +26,7 @@ export type GenerateAnswer = (input: {
   retrievalQuery: string
   messages: readonly ChatHistoryMessage[]
   evidenceText: string
+  mediaAssetContext?: string
 }) => Promise<string>
 
 export type AnswerQuestionInput = {
@@ -35,11 +37,11 @@ export type AnswerQuestionInput = {
   retrieval: RetrievalClient
   generateRetrievalQuery: GenerateRetrievalQuery
   generateAnswer: GenerateAnswer
+  loadSourceAssetUrls?: LoadSourceAssetUrls
   messages: readonly ChatHistoryMessage[]
 }
 
 export type AnswerQuestionResult = {
   answer: string
   citations: ChatCitationView[]
 }
-

src/domains/chat/index.test.ts

@@ -188,6 +188,82 @@ describe("answerQuestionWithRetrieval", () => {
     ]);
   });
 
+  it("passes retrieved image asset URLs to the answer prompt and citations", async () => {
+    const result = makeRetrievalResult({
+      chunkType: "image",
+      source: {
+        documentId: "doc_spacex",
+        sourceFileName: "document-generated.pdf",
+        sectionPath: "Assets / images / image-9-Night Rocket Launch.jpg",
+      },
+    });
+    const retrieval = {
+      query: vi.fn().mockResolvedValue({
+        results: [result],
+        evidenceText: "A SpaceX rocket launches at night.",
+        referencedChunks: [],
+        namespace: "notebook-workspace",
+        query: "SpaceX rocket photos",
+        routerUsed: "workflow_single_step",
+        answerText: null,
+      }),
+    };
+    const generateAnswer = vi
+      .fn()
+      .mockResolvedValue(
+        "Use this launch photo. https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg",
+      );
+    const generateRetrievalQuery = vi.fn().mockResolvedValue("SpaceX rocket photos");
+    const loadSourceAssetUrls = vi.fn().mockResolvedValue({
+      "images/image-9-Night Rocket Launch.jpg":
+        "https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg",
+    });
+
+    const answer = await Effect.runPromise(
+      answerQuestionWithRetrieval({
+        question: "Show me the SpaceX rocket photos.",
+        namespace: "notebook-workspace",
+        sources: [
+          makeSource({
+            id: "source_spacex",
+            title: "spacex-s1.pdf",
+            knowhereDocumentId: "doc_spacex",
+          }),
+        ],
+        excludedSourceIds: [],
+        retrieval,
+        generateRetrievalQuery,
+        generateAnswer,
+        loadSourceAssetUrls,
+        messages: [],
+      }),
+    );
+
+    expect(loadSourceAssetUrls).toHaveBeenCalledWith(
+      expect.objectContaining({ id: "source_spacex" }),
+    );
+    expect(generateAnswer).toHaveBeenCalledWith({
+      question: "Show me the SpaceX rocket photos.",
+      retrievalQuery: "SpaceX rocket photos",
+      messages: [],
+      evidenceText: "A SpaceX rocket launches at night.",
+      mediaAssetContext:
+        "- spacex-s1.pdf / Assets / images / image-9-Night Rocket Launch.jpg: https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg",
+    });
+    expect(answer.answer).toBe("Use this launch photo.");
+    expect(answer.citations).toEqual([
+      {
+        ...result,
+        assetUrl:
+          "https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg",
+        source: {
+          ...result.source,
+          sourceFileName: "spacex-s1.pdf",
+        },
+      },
+    ]);
+  });
+
   it("returns a deterministic no-results answer without calling the model", async () => {
     const retrieval = {
       query: vi.fn().mockResolvedValue({
@@ -374,6 +450,26 @@ describe("buildGroundedPrompt", () => {
       "If the sources are related but incomplete, answer what you can and briefly say what is not covered.",
     );
   });
+
+  it("includes retrieved media asset references as internal metadata", () => {
+    const prompt = buildGroundedPrompt({
+      question: "Show me the launch image.",
+      evidenceText: "A launch image was retrieved.",
+      mediaAssetContext:
+        "- spacex-s1.pdf / Assets / images / launch.jpg: https://blob.example/images/launch.jpg",
+    });
+
+    expect(prompt).toContain(
+      "Retrieved media asset references (internal; do not quote raw URLs):",
+    );
+    expect(prompt).toContain(
+      "When retrieved image or table asset references are relevant to the user's request, cite the matching source label; the UI renders media from citation metadata.",
+    );
+    expect(prompt).toContain(
+      "Do not write raw media asset URLs in the answer. They are internal metadata only.",
+    );
+    expect(prompt).toContain("https://blob.example/images/launch.jpg");
+  });
 });
 
 describe("buildRetrievalQueryPrompt", () => {

src/domains/chat/index.ts

@@ -10,6 +10,11 @@ import {
   excludeDocuments,
   normalizeRetrievalQuery,
 } from "./retrieval"
+import {
+  enrichRetrievalResultsWithAssetUrls,
+  formatRetrievedMediaAssetContext,
+  removeRetrievedMediaAssetUrls,
+} from "./media-assets"
 
 const DEFAULT_TOP_K = 8
 const NO_RESULTS_ANSWER = "I couldn't find that in your sources."
@@ -65,18 +70,27 @@ export const answerQuestionWithRetrieval = (
       return { answer: NO_RESULTS_ANSWER, citations: [] as ChatCitationView[] }
     }
 
-    const results = useNotebookSourceTitles(response.results, input.sources)
-    const answer = yield* Effect.tryPromise(() =>
-      input.generateAnswer({
-        question,
-        retrievalQuery: query,
-        messages: input.messages,
-        evidenceText,
+    const results = yield* Effect.tryPromise(() =>
+      enrichRetrievalResultsWithAssetUrls({
+        results: useNotebookSourceTitles(response.results, input.sources),
+        sources: input.sources,
+        loadSourceAssetUrls: input.loadSourceAssetUrls,
       }),
     )
+    const mediaAssetContext = formatRetrievedMediaAssetContext(results)
+    const generateAnswerInput = {
+      question,
+      retrievalQuery: query,
+      messages: input.messages,
+      evidenceText,
+      ...(mediaAssetContext ? { mediaAssetContext } : {}),
+    }
+    const generatedAnswer = yield* Effect.tryPromise(() =>
+      input.generateAnswer(generateAnswerInput),
+    )
+    const answer = removeRetrievedMediaAssetUrls(generatedAnswer, results)
     return {
       answer,
       citations: toChatCitationViews(results, answer),
     }
   })
-