feat: ensure retriever returns an image and send it to the LLM base64 encoded

Author: marcusschiesser

examples/multimodal/data/1.jpg

@@ -31,8 +31,8 @@ async function main() {
 
   const queryEngine = index.asQueryEngine({
     responseSynthesizer: new MultiModalResponseSynthesizer({ serviceContext }),
-    // TODO: set imageSimilarityTopK: 1,
-    retriever: index.asRetriever({ similarityTopK: 2 }),
+    // TODO: set text similarity to a higher value than image similarity
+    retriever: index.asRetriever({ similarityTopK: 1 }),
   });
   const result = await queryEngine.query(
     "what are Vincent van Gogh's famous paintings",

examples/multimodal/data/2.jpg

@@ -5,7 +5,6 @@ import {
   TextNode,
   VectorStoreIndex,
 } from "llamaindex";
-import * as path from "path";
 
 export async function createIndex() {
   // set up vector store index with two vector stores, one for text, the other for images
@@ -37,7 +36,7 @@ async function main() {
       continue;
     }
     if (node instanceof ImageNode) {
-      console.log(`Image: ${path.join(__dirname, node.id_)}`);
+      console.log(`Image: ${node.getUrl()}`);
     } else if (node instanceof TextNode) {
       console.log("Text:", (node as TextNode).text.substring(0, 128));
     }

examples/multimodal/data/3.jpg

@@ -10,6 +10,7 @@
     "@xenova/transformers": "^2.10.0",
     "assemblyai": "^4.0.0",
     "crypto-js": "^4.2.0",
+    "file-type": "^18.7.0",
     "js-tiktoken": "^1.0.8",
     "lodash": "^4.17.21",
     "mammoth": "^1.6.0",

examples/multimodal/data/San Francisco.txt

@@ -1,4 +1,5 @@
 import CryptoJS from "crypto-js";
+import path from "path";
 import { v4 as uuidv4 } from "uuid";
 
 export enum NodeRelationship {
@@ -304,6 +305,12 @@ export class ImageNode<T extends Metadata = Metadata> extends TextNode<T> {
   getType(): ObjectType {
     return ObjectType.IMAGE;
   }
+
+  getUrl(): URL {
+    // id_ stores the relative path, convert it to the URL of the file
+    const absPath = path.resolve(this.id_);
+    return new URL(`file://${absPath}`);
+  }
 }
 
 export class ImageDocument<T extends Metadata = Metadata> extends ImageNode<T> {

examples/multimodal/rag.ts

@@ -1,7 +1,7 @@
 import _ from "lodash";
 import { ImageType } from "../Node";
 import { DEFAULT_SIMILARITY_TOP_K } from "../constants";
-import { VectorStoreQueryMode } from "../storage";
+import { DEFAULT_FS, VectorStoreQueryMode } from "../storage";
 import { SimilarityType } from "./types";
 
 /**
@@ -185,6 +185,16 @@ export function getTopKMMREmbeddings(
   return [resultSimilarities, resultIds];
 }
 
+async function blobToDataUrl(input: Blob) {
+  const { fileTypeFromBuffer } = await import("file-type");
+  const buffer = Buffer.from(await input.arrayBuffer());
+  const type = await fileTypeFromBuffer(buffer);
+  if (!type) {
+    throw new Error("Unsupported image type");
+  }
+  return "data:" + type.mime + ";base64," + buffer.toString("base64");
+}
+
 export async function readImage(input: ImageType) {
   const { RawImage } = await import("@xenova/transformers");
   if (input instanceof Blob) {
@@ -195,3 +205,53 @@ export async function readImage(input: ImageType) {
     throw new Error(`Unsupported input type: ${typeof input}`);
   }
 }
+
+export async function imageToString(input: ImageType): Promise<string> {
+  if (input instanceof Blob) {
+    // if the image is a Blob, convert it to a base64 data URL
+    return await blobToDataUrl(input);
+  } else if (_.isString(input)) {
+    return input;
+  } else if (input instanceof URL) {
+    return input.toString();
+  } else {
+    throw new Error(`Unsupported input type: ${typeof input}`);
+  }
+}
+
+export function stringToImage(input: string): ImageType {
+  if (input.startsWith("data:")) {
+    // if the input is a base64 data URL, convert it back to a Blob
+    const base64Data = input.split(",")[1];
+    const byteArray = Buffer.from(base64Data, "base64");
+    return new Blob([byteArray]);
+  } else if (input.startsWith("http://") || input.startsWith("https://")) {
+    return new URL(input);
+  } else if (_.isString(input)) {
+    return input;
+  } else {
+    throw new Error(`Unsupported input type: ${typeof input}`);
+  }
+}
+
+export async function imageToDataUrl(input: ImageType): Promise<string> {
+  // first ensure, that the input is a Blob
+  if (
+    (input instanceof URL && input.protocol === "file:") ||
+    _.isString(input)
+  ) {
+    // string or file URL
+    const fs = DEFAULT_FS;
+    const dataBuffer = await fs.readFile(
+      input instanceof URL ? input.pathname : input,
+    );
+    input = new Blob([dataBuffer]);
+  } else if (!(input instanceof Blob)) {
+    if (input instanceof URL) {
+      throw new Error(`Unsupported URL with protocol: ${input.protocol}`);
+    } else {
+      throw new Error(`Unsupported input type: ${typeof input}`);
+    }
+  }
+  return await blobToDataUrl(input);
+}

examples/multimodal/retrieve.ts

@@ -1,10 +1,10 @@
-import { Event } from "../../callbacks/CallbackManager";
-import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
-import { BaseEmbedding } from "../../embeddings";
 import { globalsHelper } from "../../GlobalsHelper";
-import { Metadata, NodeWithScore } from "../../Node";
+import { ImageNode, Metadata, NodeWithScore } from "../../Node";
 import { BaseRetriever } from "../../Retriever";
 import { ServiceContext } from "../../ServiceContext";
+import { Event } from "../../callbacks/CallbackManager";
+import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
+import { BaseEmbedding } from "../../embeddings";
 import {
   VectorStoreQuery,
   VectorStoreQueryMode,
@@ -108,6 +108,12 @@ export class VectorIndexRetriever implements BaseRetriever {
       }
 
       const node = this.index.indexStruct.nodesDict[result.ids[i]];
+      // XXX: Hack, if it's an image node, we reconstruct the image from the URL
+      // Alternative: Store image in doc store and retrieve it here
+      if (node instanceof ImageNode) {
+        node.image = node.getUrl();
+      }
+
       nodesWithScores.push({
         node: node,
         score: result.similarities[i],