feat: display files from llamacloud (#153)

---------
Co-authored-by: leehuwuj <[email protected]>
Co-authored-by: Marcus Schiesser <[email protected]>

Author: 3 people

.changeset/happy-hairs-kick.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Display files in sources using LlamaCloud indexes.

helpers/index.ts

@@ -177,10 +177,9 @@ export const installTemplate = async (
     }
 
     // Create outputs directory
-    if (props.tools && props.tools.length > 0) {
-      await makeDir(path.join(props.root, "output/tools"));
-      await makeDir(path.join(props.root, "output/uploaded"));
-    }
+    await makeDir(path.join(props.root, "output/tools"));
+    await makeDir(path.join(props.root, "output/uploaded"));
+    await makeDir(path.join(props.root, "output/llamacloud"));
   } else {
     // this is a frontend for a full-stack app, create .env file with model information
     await createFrontendEnvFile(props.root, {

templates/components/llamaindex/typescript/documents/documents.ts

@@ -0,0 +1,44 @@
+import fs from "fs";
+import crypto from "node:crypto";
+import { getExtractors } from "../../engine/loader";
+
+const MIME_TYPE_TO_EXT: Record<string, string> = {
+  "application/pdf": "pdf",
+  "text/plain": "txt",
+  "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+    "docx",
+};
+
+const UPLOADED_FOLDER = "output/uploaded";
+
+export async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
+  const extractors = getExtractors();
+  const reader = extractors[MIME_TYPE_TO_EXT[mimeType]];
+
+  if (!reader) {
+    throw new Error(`Unsupported document type: ${mimeType}`);
+  }
+  console.log(`Processing uploaded document of type: ${mimeType}`);
+  return await reader.loadDataAsContent(fileBuffer);
+}
+
+export async function saveDocument(fileBuffer: Buffer, mimeType: string) {
+  const fileExt = MIME_TYPE_TO_EXT[mimeType];
+  if (!fileExt) throw new Error(`Unsupported document type: ${mimeType}`);
+
+  const filename = `${crypto.randomUUID()}.${fileExt}`;
+  const filepath = `${UPLOADED_FOLDER}/${filename}`;
+  const fileurl = `${process.env.FILESERVER_URL_PREFIX}/${filepath}`;
+
+  if (!fs.existsSync(UPLOADED_FOLDER)) {
+    fs.mkdirSync(UPLOADED_FOLDER, { recursive: true });
+  }
+  await fs.promises.writeFile(filepath, fileBuffer);
+
+  console.log(`Saved document file to ${filepath}.\nURL: ${fileurl}`);
+  return {
+    filename,
+    filepath,
+    fileurl,
+  };
+}

templates/components/llamaindex/typescript/documents/helper.ts

@@ -0,0 +1,65 @@
+import {
+  BaseNode,
+  Document,
+  IngestionPipeline,
+  Metadata,
+  Settings,
+  SimpleNodeParser,
+  storageContextFromDefaults,
+  VectorStoreIndex,
+} from "llamaindex";
+import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
+import { getDataSource } from "../../engine";
+
+export async function runPipeline(documents: Document[], filename: string) {
+  const currentIndex = await getDataSource();
+
+  // Update documents with metadata
+  for (const document of documents) {
+    document.metadata = {
+      ...document.metadata,
+      file_name: filename,
+      private: "true", // to separate from other public documents
+    };
+  }
+
+  if (currentIndex instanceof LlamaCloudIndex) {
+    // LlamaCloudIndex processes the documents automatically
+    // so we don't need ingestion pipeline, just insert the documents directly
+    for (const document of documents) {
+      await currentIndex.insert(document);
+    }
+  } else {
+    // Use ingestion pipeline to process the documents into nodes and add them to the vector store
+    const pipeline = new IngestionPipeline({
+      transformations: [
+        new SimpleNodeParser({
+          chunkSize: Settings.chunkSize,
+          chunkOverlap: Settings.chunkOverlap,
+        }),
+        Settings.embedModel,
+      ],
+    });
+    const nodes = await pipeline.run({ documents });
+    await addNodesToVectorStore(nodes, currentIndex);
+  }
+
+  return documents.map((document) => document.id_);
+}
+
+async function addNodesToVectorStore(
+  nodes: BaseNode<Metadata>[],
+  currentIndex: VectorStoreIndex | null,
+) {
+  if (currentIndex) {
+    await currentIndex.insertNodes(nodes);
+  } else {
+    // Not using vectordb and haven't generated local index yet
+    const storageContext = await storageContextFromDefaults({
+      persistDir: "./cache",
+    });
+    currentIndex = await VectorStoreIndex.init({ nodes, storageContext });
+  }
+  currentIndex.storageContext.docStore.persist();
+  console.log("Added nodes to the vector store.");
+}

templates/components/llamaindex/typescript/documents/pipeline.ts

@@ -0,0 +1,11 @@
+import { loadDocuments, saveDocument } from "./helper";
+import { runPipeline } from "./pipeline";
+
+export async function uploadDocument(raw: string): Promise<string[]> {
+  const [header, content] = raw.split(",");
+  const mimeType = header.replace("data:", "").replace(";base64", "");
+  const fileBuffer = Buffer.from(content, "base64");
+  const documents = await loadDocuments(fileBuffer, mimeType);
+  const { filename } = await saveDocument(fileBuffer, mimeType);
+  return await runPipeline(documents, filename);
+}

templates/components/llamaindex/typescript/documents/upload.ts

@@ -6,23 +6,31 @@ import {
   ToolCall,
   ToolOutput,
 } from "llamaindex";
+import { LLamaCloudFileService } from "./service";
 
-export function appendSourceData(
+export async function appendSourceData(
   data: StreamData,
   sourceNodes?: NodeWithScore<Metadata>[],
 ) {
   if (!sourceNodes?.length) return;
-  data.appendMessageAnnotation({
-    type: "sources",
-    data: {
-      nodes: sourceNodes.map((node) => ({
+  try {
+    const nodes = await Promise.all(
+      sourceNodes.map(async (node) => ({
         ...node.node.toMutableJSON(),
         id: node.node.id_,
         score: node.score ?? null,
-        url: getNodeUrl(node.node.metadata),
+        url: await getNodeUrl(node.node.metadata),
       })),
-    },
-  });
+    );
+    data.appendMessageAnnotation({
+      type: "sources",
+      data: {
+        nodes,
+      },
+    });
+  } catch (error) {
+    console.error("Error appending source data:", error);
+  }
 }
 
 export function appendEventData(data: StreamData, title?: string) {
@@ -68,9 +76,9 @@ export function createStreamTimeout(stream: StreamData) {
 export function createCallbackManager(stream: StreamData) {
   const callbackManager = new CallbackManager();
 
-  callbackManager.on("retrieve-end", (data) => {
+  callbackManager.on("retrieve-end", async (data) => {
     const { nodes, query } = data.detail.payload;
-    appendSourceData(stream, nodes);
+    await appendSourceData(stream, nodes);
     appendEventData(stream, `Retrieving context for query: '${query}'`);
     appendEventData(
       stream,
@@ -97,19 +105,27 @@ export function createCallbackManager(stream: StreamData) {
   return callbackManager;
 }
 
-function getNodeUrl(metadata: Metadata) {
-  const url = metadata["URL"];
-  if (url) return url;
-  const fileName = metadata["file_name"];
+async function getNodeUrl(metadata: Metadata) {
   if (!process.env.FILESERVER_URL_PREFIX) {
     console.warn(
       "FILESERVER_URL_PREFIX is not set. File URLs will not be generated.",
     );
-    return undefined;
   }
-  if (fileName) {
-    const folder = metadata["private"] ? "output/uploaded" : "data";
+  const fileName = metadata["file_name"];
+  if (fileName && process.env.FILESERVER_URL_PREFIX) {
+    // file_name exists and file server is configured
+    const isLocalFile = metadata["is_local_file"] === "true";
+    const pipelineId = metadata["pipeline_id"];
+    if (pipelineId && !isLocalFile) {
+      // file is from LlamaCloud and was not ingested locally
+      // TODO trigger but don't await file download and just use convention to generate the URL (see Python code)
+      // return `${process.env.FILESERVER_URL_PREFIX}/output/llamacloud/${pipelineId}\$${fileName}`;
+      return await LLamaCloudFileService.getFileUrl(fileName, pipelineId);
+    }
+    const isPrivate = metadata["private"] === "true";
+    const folder = isPrivate ? "output/uploaded" : "data";
     return `${process.env.FILESERVER_URL_PREFIX}/${folder}/${fileName}`;
   }
-  return undefined;
+  // fallback to URL in metadata (e.g. for websites)
+  return metadata["URL"];
 }