add back image citations

web/src/app/chat/ChatPage.tsx

@@ -2165,6 +2165,7 @@ export function ChatPage({
                                     }
                                   >
                                     <AIMessage
+                                      setPopup={setPopup}
                                       hasChildAI={hasChildAI}
                                       hasParentAI={hasParentAI}
                                       continueGenerating={

web/src/app/chat/interfaces.ts

@@ -143,3 +143,10 @@ export interface StreamingError {
   error: string;
   stack_trace: string;
 }
+
+export interface ImageGenerationResult {
+  revised_prompt: string;
+  url: string;
+}
+
+export type ImageGenerationResults = ImageGenerationResult[];

web/src/app/chat/message/Messages.tsx

@@ -26,7 +26,12 @@ import { SourceIcon } from "@/components/SourceIcon";
 import { SkippedSearch } from "./SkippedSearch";
 import remarkGfm from "remark-gfm";
 import { CopyButton } from "@/components/CopyButton";
-import { ChatFileType, FileDescriptor, ToolCallMetadata } from "../interfaces";
+import {
+  ChatFileType,
+  FileDescriptor,
+  ImageGenerationResults,
+  ToolCallMetadata,
+} from "../interfaces";
 import {
   IMAGE_GENERATION_TOOL_NAME,
   SEARCH_TOOL_NAME,
@@ -46,7 +51,11 @@ import { AssistantIcon } from "@/components/assistants/AssistantIcon";
 import { Citation } from "@/components/search/results/Citation";
 import { DocumentMetadataBlock } from "@/components/search/DocumentDisplay";
 
-import { LikeFeedback, DislikeFeedback } from "@/components/icons/icons";
+import {
+  LikeFeedback,
+  DislikeFeedback,
+  ToolCallIcon,
+} from "@/components/icons/icons";
 import {
   CustomTooltip,
   TooltipGroup,
@@ -60,6 +69,9 @@ import GeneratingImageDisplay from "../tools/GeneratingImageDisplay";
 import RegenerateOption from "../RegenerateOption";
 import { LlmOverride } from "@/lib/hooks";
 import { ContinueGenerating } from "./ContinueMessage";
+import DualPromptDisplay from "../tools/ImagePromptCitaiton";
+import { PopupSpec } from "@/components/admin/connectors/Popup";
+import { Popover } from "@/components/popover/Popover";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -145,6 +157,7 @@ export const AIMessage = ({
   currentPersona,
   otherMessagesCanSwitchTo,
   onMessageSelection,
+  setPopup,
 }: {
   hasChildAI?: boolean;
   hasParentAI?: boolean;
@@ -175,7 +188,9 @@ export const AIMessage = ({
   retrievalDisabled?: boolean;
   overriddenModel?: string;
   regenerate?: (modelOverRide: LlmOverride) => Promise<void>;
+  setPopup?: (popupSpec: PopupSpec | null) => void;
 }) => {
+  const [isPopoverOpen, setIsPopoverOpen] = useState(false);
   const toolCallGenerating = toolCall && !toolCall.tool_result;
   const processContent = (content: string | JSX.Element) => {
     if (typeof content !== "string") {
@@ -198,6 +213,9 @@ export const AIMessage = ({
         return content;
       }
     }
+    if (toolCall?.tool_result) {
+      return content + ` [${toolCall.tool_name}]()`;
+    }
 
     return content + (!isComplete && !toolCallGenerating ? " [*]() " : "");
   };
@@ -373,6 +391,47 @@ export const AIMessage = ({
                                   const { node, ...rest } = props;
                                   const value = rest.children;
 
+                                  if (
+                                    value
+                                      ?.toString()
+                                      .startsWith(IMAGE_GENERATION_TOOL_NAME)
+                                  ) {
+                                    const imageGenerationResult =
+                                      toolCall?.tool_result as ImageGenerationResults;
+
+                                    return (
+                                      <Popover
+                                        open={isPopoverOpen}
+                                        onOpenChange={() => null} // only allow closing from the icon
+                                        content={
+                                          <button
+                                            onMouseDown={() => {
+                                              setIsPopoverOpen(!isPopoverOpen);
+                                            }}
+                                          >
+                                            <ToolCallIcon className="cursor-pointer flex-none text-blue-500 hover:text-blue-700 !h-4 !w-4 inline-block" />
+                                          </button>
+                                        }
+                                        popover={
+                                          <DualPromptDisplay
+                                            arg="Prompt"
+                                            setPopup={setPopup!}
+                                            prompt1={
+                                              imageGenerationResult[0]
+                                                .revised_prompt
+                                            }
+                                            prompt2={
+                                              imageGenerationResult[1]
+                                                .revised_prompt
+                                            }
+                                          />
+                                        }
+                                        side="top"
+                                        align="center"
+                                      />
+                                    );
+                                  }
+
                                   if (value?.toString().startsWith("*")) {
                                     return (
                                       <div className="flex-none bg-background-800 inline-block rounded-full h-3 w-3 ml-2" />