Release 0.1.39 (#243)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

.github/workflows/lint_on_push_or_pull.yml

@@ -30,3 +30,13 @@ jobs:
 
       - name: Run Prettier
         run: pnpm run format
+
+      - name: Run Python format check
+        uses: chartboost/ruff-action@v1
+        with:
+          args: "format --check"
+
+      - name: Run Python lint
+        uses: chartboost/ruff-action@v1
+        with:
+          args: "check"

CHANGELOG.md

@@ -1,5 +1,29 @@
 # create-llama
 
+## 0.1.39
+
+### Patch Changes
+
+- 3fb93c7: Use LlamaCloud pipeline for data ingestion in TS (private file uploads and generate script)
+
+## 0.1.38
+
+### Patch Changes
+
+- bd5e39a: Fix error that files in sub folders of 'data' are not displayed
+
+## 0.1.37
+
+### Patch Changes
+
+- 9fd832c: Add in-text citation references
+
+## 0.1.36
+
+### Patch Changes
+
+- 2b7a5d8: Fix: private file upload not working in Python without LlamaCloud
+
 ## 0.1.35
 
 ### Patch Changes

helpers/env-variables.ts

@@ -4,6 +4,7 @@ import { TOOL_SYSTEM_PROMPT_ENV_VAR, Tool } from "./tools";
 import {
   InstallTemplateArgs,
   ModelConfig,
+  TemplateDataSource,
   TemplateFramework,
   TemplateObservability,
   TemplateType,
@@ -159,7 +160,7 @@ const getVectorDBEnvs = (
         {
           name: "LLAMA_CLOUD_ORGANIZATION_ID",
           description:
-            "The organization ID for the LlamaCloud project (uses default organization if not specified - Python only)",
+            "The organization ID for the LlamaCloud project (uses default organization if not specified)",
         },
         ...(framework === "nextjs"
           ? // activate index selector per default (not needed for non-NextJS backends as it's handled by createFrontendEnvFile)
@@ -423,7 +424,11 @@ const getToolEnvs = (tools?: Tool[]): EnvVar[] => {
   return toolEnvs;
 };
 
-const getSystemPromptEnv = (tools?: Tool[]): EnvVar => {
+const getSystemPromptEnv = (
+  tools?: Tool[],
+  dataSources?: TemplateDataSource[],
+  framework?: TemplateFramework,
+): EnvVar[] => {
   const defaultSystemPrompt =
     "You are a helpful assistant who helps users with their questions.";
 
@@ -442,11 +447,49 @@ const getSystemPromptEnv = (tools?: Tool[]): EnvVar => {
     ? `\"${toolSystemPrompt}\"`
     : defaultSystemPrompt;
 
-  return {
-    name: "SYSTEM_PROMPT",
-    description: "The system prompt for the AI model.",
-    value: systemPrompt,
-  };
+  const systemPromptEnv = [
+    {
+      name: "SYSTEM_PROMPT",
+      description: "The system prompt for the AI model.",
+      value: systemPrompt,
+    },
+  ];
+
+  // Citation only works with FastAPI along with the chat engine and data source provided for now.
+  if (
+    framework === "fastapi" &&
+    tools?.length == 0 &&
+    (dataSources?.length ?? 0 > 0)
+  ) {
+    const citationPrompt = `'You have provided information from a knowledge base that has been passed to you in nodes of information.
+Each node has useful metadata such as node ID, file name, page, etc.
+Please add the citation to the data node for each sentence or paragraph that you reference in the provided information.
+The citation format is: . [citation:<node_id>]()
+Where the <node_id> is the unique identifier of the data node.
+
+Example:
+We have two nodes:
+  node_id: xyz
+  file_name: llama.pdf
+  
+  node_id: abc
+  file_name: animal.pdf
+
+User question: Tell me a fun fact about Llama.
+Your answer:
+A baby llama is called "Cria" [citation:xyz]().
+It often live in desert [citation:abc]().
+It\\'s cute animal.
+'`;
+    systemPromptEnv.push({
+      name: "SYSTEM_CITATION_PROMPT",
+      description:
+        "An additional system prompt to add citation when responding to user questions.",
+      value: citationPrompt,
+    });
+  }
+
+  return systemPromptEnv;
 };
 
 const getTemplateEnvs = (template?: TemplateType): EnvVar[] => {
@@ -525,7 +568,7 @@ export const createBackendEnvFile = async (
     ...getToolEnvs(opts.tools),
     ...getTemplateEnvs(opts.template),
     ...getObservabilityEnvs(opts.observability),
-    getSystemPromptEnv(opts.tools),
+    ...getSystemPromptEnv(opts.tools, opts.dataSources, opts.framework),
   ];
   // Render and write env file
   const content = renderEnvVar(envVars);

package.json

@@ -1,6 +1,6 @@
 {
   "name": "create-llama",
-  "version": "0.1.35",
+  "version": "0.1.39",
   "description": "Create LlamaIndex-powered apps with one command",
   "keywords": [
     "rag",

templates/components/engines/python/agent/tools/__init__.py

@@ -1,8 +1,6 @@
 import os
 import yaml
-import json
 import importlib
-from cachetools import cached, LRUCache
 from llama_index.core.tools.tool_spec.base import BaseToolSpec
 from llama_index.core.tools.function_tool import FunctionTool
 
@@ -13,7 +11,6 @@ class ToolType:
 
 
 class ToolFactory:
-
     TOOL_SOURCE_PACKAGE_MAP = {
         ToolType.LLAMAHUB: "llama_index.tools",
         ToolType.LOCAL: "app.engine.tools",

templates/components/engines/python/agent/tools/interpreter.py

@@ -3,7 +3,7 @@ import logging
 import base64
 import uuid
 from pydantic import BaseModel
-from typing import List, Tuple, Dict, Optional
+from typing import List, Dict, Optional
 from llama_index.core.tools import FunctionTool
 from e2b_code_interpreter import CodeInterpreter
 from e2b_code_interpreter.models import Logs
@@ -26,7 +26,6 @@ class E2BToolOutput(BaseModel):
 
 
 class E2BCodeInterpreter:
-
     output_dir = "output/tool"
 
     def __init__(self, api_key: str = None):

templates/components/engines/python/chat/__init__.py

@@ -1,11 +1,20 @@
 import os
+
 from app.engine.index import get_index
+from app.engine.node_postprocessors import NodeCitationProcessor
 from fastapi import HTTPException
+from llama_index.core.chat_engine import CondensePlusContextChatEngine
 
 
 def get_chat_engine(filters=None, params=None):
     system_prompt = os.getenv("SYSTEM_PROMPT")
-    top_k = os.getenv("TOP_K", 3)
+    citation_prompt = os.getenv("SYSTEM_CITATION_PROMPT", None)
+    top_k = int(os.getenv("TOP_K", 3))
+
+    node_postprocessors = []
+    if citation_prompt:
+        node_postprocessors = [NodeCitationProcessor()]
+        system_prompt = f"{system_prompt}\n{citation_prompt}"
 
     index = get_index(params)
     if index is None:
@@ -16,9 +25,13 @@ def get_chat_engine(filters=None, params=None):
             ),
         )
 
-    return index.as_chat_engine(
-        similarity_top_k=int(top_k),
-        system_prompt=system_prompt,
-        chat_mode="condense_plus_context",
+    retriever = index.as_retriever(
+        similarity_top_k=top_k,
         filters=filters,
     )
+
+    return CondensePlusContextChatEngine.from_defaults(
+        system_prompt=system_prompt,
+        retriever=retriever,
+        node_postprocessors=node_postprocessors,
+    )

templates/components/engines/python/chat/node_postprocessors.py

@@ -0,0 +1,21 @@
+from typing import List, Optional
+
+from llama_index.core import QueryBundle
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.schema import NodeWithScore
+
+
+class NodeCitationProcessor(BaseNodePostprocessor):
+    """
+    Append node_id into metadata for citation purpose.
+    Config SYSTEM_CITATION_PROMPT in your runtime environment variable to enable this feature.
+    """
+
+    def _postprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        for node_score in nodes:
+            node_score.node.metadata["node_id"] = node_score.node.node_id
+        return nodes

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

@@ -11,7 +11,20 @@ const MIME_TYPE_TO_EXT: Record<string, string> = {
 
 const UPLOADED_FOLDER = "output/uploaded";
 
-export async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
+export async function storeAndParseFile(fileBuffer: Buffer, mimeType: string) {
+  const documents = await loadDocuments(fileBuffer, mimeType);
+  const { filename } = await saveDocument(fileBuffer, mimeType);
+  for (const document of documents) {
+    document.metadata = {
+      ...document.metadata,
+      file_name: filename,
+      private: "true", // to separate private uploads from public documents
+    };
+  }
+  return documents;
+}
+
+async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
   const extractors = getExtractors();
   const reader = extractors[MIME_TYPE_TO_EXT[mimeType]];
 
@@ -22,7 +35,7 @@ export async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
   return await reader.loadDataAsContent(fileBuffer);
 }
 
-export async function saveDocument(fileBuffer: Buffer, mimeType: string) {
+async function saveDocument(fileBuffer: Buffer, mimeType: string) {
   const fileExt = MIME_TYPE_TO_EXT[mimeType];
   if (!fileExt) throw new Error(`Unsupported document type: ${mimeType}`);
 

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

@@ -5,34 +5,24 @@ import {
   SimpleNodeParser,
   VectorStoreIndex,
 } from "llamaindex";
-import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
 
 export async function runPipeline(
-  currentIndex: VectorStoreIndex | LlamaCloudIndex,
+  currentIndex: VectorStoreIndex,
   documents: Document[],
 ) {
-  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 currentIndex.insertNodes(nodes);
-    currentIndex.storageContext.docStore.persist();
-    console.log("Added nodes to the vector store.");
-  }
-
+  // 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 currentIndex.insertNodes(nodes);
+  currentIndex.storageContext.docStore.persist();
+  console.log("Added nodes to the vector store.");
   return documents.map((document) => document.id_);
 }

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

@@ -1,26 +1,32 @@
-import { VectorStoreIndex } from "llamaindex";
+import { LLamaCloudFileService, VectorStoreIndex } from "llamaindex";
 import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
-import { loadDocuments, saveDocument } from "./helper";
+import { storeAndParseFile } from "./helper";
 import { runPipeline } from "./pipeline";
 
 export async function uploadDocument(
   index: VectorStoreIndex | LlamaCloudIndex,
+  filename: string,
   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);
 
-  // Update documents with metadata
-  for (const document of documents) {
-    document.metadata = {
-      ...document.metadata,
-      file_name: filename,
-      private: "true", // to separate private uploads from public documents
-    };
+  if (index instanceof LlamaCloudIndex) {
+    // trigger LlamaCloudIndex API to upload the file and run the pipeline
+    const projectId = await index.getProjectId();
+    const pipelineId = await index.getPipelineId();
+    return [
+      await LLamaCloudFileService.addFileToPipeline(
+        projectId,
+        pipelineId,
+        new File([fileBuffer], filename, { type: mimeType }),
+        { private: "true" },
+      ),
+    ];
   }
 
-  return await runPipeline(index, documents);
+  // run the pipeline for other vector store indexes
+  const documents = await storeAndParseFile(fileBuffer, mimeType);
+  return runPipeline(index, documents);
 }

templates/components/llamaindex/typescript/streaming/events.ts

@@ -1,13 +1,18 @@
 import { StreamData } from "ai";
 import {
   CallbackManager,
+  LLamaCloudFileService,
   Metadata,
   MetadataMode,
   NodeWithScore,
   ToolCall,
   ToolOutput,
 } from "llamaindex";
-import { LLamaCloudFileService } from "./service";
+import path from "node:path";
+import { DATA_DIR } from "../../engine/loader";
+import { downloadFile } from "./file";
+
+const LLAMA_CLOUD_DOWNLOAD_FOLDER = "output/llamacloud";
 
 export function appendSourceData(
   data: StreamData,
@@ -84,7 +89,7 @@ export function createCallbackManager(stream: StreamData) {
       stream,
       `Retrieved ${nodes.length} sources to use as context for the query`,
     );
-    LLamaCloudFileService.downloadFiles(nodes); // don't await to avoid blocking chat streaming
+    downloadFilesFromNodes(nodes); // don't await to avoid blocking chat streaming
   });
 
   callbackManager.on("llm-tool-call", (event) => {
@@ -116,15 +121,71 @@ function getNodeUrl(metadata: Metadata) {
   if (fileName && process.env.FILESERVER_URL_PREFIX) {
     // file_name exists and file server is configured
     const pipelineId = metadata["pipeline_id"];
-    if (pipelineId && metadata["private"] == null) {
-      // file is from LlamaCloud and was not ingested locally
-      const name = LLamaCloudFileService.toDownloadedName(pipelineId, fileName);
-      return `${process.env.FILESERVER_URL_PREFIX}/output/llamacloud/${name}`;
+    if (pipelineId) {
+      const name = toDownloadedName(pipelineId, fileName);
+      return `${process.env.FILESERVER_URL_PREFIX}/${LLAMA_CLOUD_DOWNLOAD_FOLDER}/${name}`;
     }
     const isPrivate = metadata["private"] === "true";
-    const folder = isPrivate ? "output/uploaded" : "data";
-    return `${process.env.FILESERVER_URL_PREFIX}/${folder}/${fileName}`;
+    if (isPrivate) {
+      return `${process.env.FILESERVER_URL_PREFIX}/output/uploaded/${fileName}`;
+    }
+    const filePath = metadata["file_path"];
+    const dataDir = path.resolve(DATA_DIR);
+
+    if (filePath && dataDir) {
+      const relativePath = path.relative(dataDir, filePath);
+      return `${process.env.FILESERVER_URL_PREFIX}/data/${relativePath}`;
+    }
   }
   // fallback to URL in metadata (e.g. for websites)
   return metadata["URL"];
 }
+
+async function downloadFilesFromNodes(nodes: NodeWithScore<Metadata>[]) {
+  try {
+    const files = nodesToLlamaCloudFiles(nodes);
+    for (const { pipelineId, fileName, downloadedName } of files) {
+      const downloadUrl = await LLamaCloudFileService.getFileUrl(
+        pipelineId,
+        fileName,
+      );
+      if (downloadUrl) {
+        await downloadFile(
+          downloadUrl,
+          downloadedName,
+          LLAMA_CLOUD_DOWNLOAD_FOLDER,
+        );
+      }
+    }
+  } catch (error) {
+    console.error("Error downloading files from nodes:", error);
+  }
+}
+
+function nodesToLlamaCloudFiles(nodes: NodeWithScore<Metadata>[]) {
+  const files: Array<{
+    pipelineId: string;
+    fileName: string;
+    downloadedName: string;
+  }> = [];
+  for (const node of nodes) {
+    const pipelineId = node.node.metadata["pipeline_id"];
+    const fileName = node.node.metadata["file_name"];
+    if (!pipelineId || !fileName) continue;
+    const isDuplicate = files.some(
+      (f) => f.pipelineId === pipelineId && f.fileName === fileName,
+    );
+    if (!isDuplicate) {
+      files.push({
+        pipelineId,
+        fileName,
+        downloadedName: toDownloadedName(pipelineId, fileName),
+      });
+    }
+  }
+  return files;
+}
+
+function toDownloadedName(pipelineId: string, fileName: string) {
+  return `${pipelineId}$${fileName}`;
+}

templates/components/llamaindex/typescript/streaming/file.ts

@@ -0,0 +1,35 @@
+import fs from "node:fs";
+import https from "node:https";
+import path from "node:path";
+
+export async function downloadFile(
+  urlToDownload: string,
+  filename: string,
+  folder = "output/uploaded",
+) {
+  try {
+    const downloadedPath = path.join(folder, filename);
+
+    // Check if file already exists
+    if (fs.existsSync(downloadedPath)) return;
+
+    const file = fs.createWriteStream(downloadedPath);
+    https
+      .get(urlToDownload, (response) => {
+        response.pipe(file);
+        file.on("finish", () => {
+          file.close(() => {
+            console.log("File downloaded successfully");
+          });
+        });
+      })
+      .on("error", (err) => {
+        fs.unlink(downloadedPath, () => {
+          console.error("Error downloading file:", err);
+          throw err;
+        });
+      });
+  } catch (error) {
+    throw new Error(`Error downloading file: ${error}`);
+  }
+}

templates/components/llamaindex/typescript/streaming/service.ts

@@ -1,187 +0,0 @@
-import { Metadata, NodeWithScore } from "llamaindex";
-import fs from "node:fs";
-import https from "node:https";
-import path from "node:path";
-
-const LLAMA_CLOUD_OUTPUT_DIR = "output/llamacloud";
-const LLAMA_CLOUD_BASE_URL = "https://cloud.llamaindex.ai/api/v1";
-const FILE_DELIMITER = "$"; // delimiter between pipelineId and filename
-
-type LlamaCloudFile = {
-  name: string;
-  file_id: string;
-  project_id: string;
-};
-
-type LLamaCloudProject = {
-  id: string;
-  organization_id: string;
-  name: string;
-  is_default: boolean;
-};
-
-type LLamaCloudPipeline = {
-  id: string;
-  name: string;
-  project_id: string;
-};
-
-export class LLamaCloudFileService {
-  private static readonly headers = {
-    Accept: "application/json",
-    Authorization: `Bearer ${process.env.LLAMA_CLOUD_API_KEY}`,
-  };
-
-  public static async getAllProjectsWithPipelines() {
-    try {
-      const projects = await LLamaCloudFileService.getAllProjects();
-      const pipelines = await LLamaCloudFileService.getAllPipelines();
-      return projects.map((project) => ({
-        ...project,
-        pipelines: pipelines.filter((p) => p.project_id === project.id),
-      }));
-    } catch (error) {
-      console.error("Error listing projects and pipelines:", error);
-      return [];
-    }
-  }
-
-  public static async downloadFiles(nodes: NodeWithScore<Metadata>[]) {
-    const files = LLamaCloudFileService.nodesToDownloadFiles(nodes);
-    if (!files.length) return;
-    console.log("Downloading files from LlamaCloud...");
-    for (const file of files) {
-      await LLamaCloudFileService.downloadFile(file.pipelineId, file.fileName);
-    }
-  }
-
-  public static toDownloadedName(pipelineId: string, fileName: string) {
-    return `${pipelineId}${FILE_DELIMITER}${fileName}`;
-  }
-
-  /**
-   * This function will return an array of unique files to download from LlamaCloud
-   * We only download files that are uploaded directly in LlamaCloud datasources (don't have `private` in metadata)
-   * Files are uploaded directly in LlamaCloud datasources don't have `private` in metadata (public docs)
-   * Files are uploaded from local via `generate` command will have `private=false` (public docs)
-   * Files are uploaded from local via `/chat/upload` endpoint will have `private=true` (private docs)
-   *
-   * @param nodes
-   * @returns list of unique files to download
-   */
-  private static nodesToDownloadFiles(nodes: NodeWithScore<Metadata>[]) {
-    const downloadFiles: Array<{
-      pipelineId: string;
-      fileName: string;
-    }> = [];
-    for (const node of nodes) {
-      const isLocalFile = node.node.metadata["private"] != null;
-      const pipelineId = node.node.metadata["pipeline_id"];
-      const fileName = node.node.metadata["file_name"];
-      if (isLocalFile || !pipelineId || !fileName) continue;
-      const isDuplicate = downloadFiles.some(
-        (f) => f.pipelineId === pipelineId && f.fileName === fileName,
-      );
-      if (!isDuplicate) {
-        downloadFiles.push({ pipelineId, fileName });
-      }
-    }
-    return downloadFiles;
-  }
-
-  private static async downloadFile(pipelineId: string, fileName: string) {
-    try {
-      const downloadedName = LLamaCloudFileService.toDownloadedName(
-        pipelineId,
-        fileName,
-      );
-      const downloadedPath = path.join(LLAMA_CLOUD_OUTPUT_DIR, downloadedName);
-
-      // Check if file already exists
-      if (fs.existsSync(downloadedPath)) return;
-
-      const urlToDownload = await LLamaCloudFileService.getFileUrlByName(
-        pipelineId,
-        fileName,
-      );
-      if (!urlToDownload) throw new Error("File not found in LlamaCloud");
-
-      const file = fs.createWriteStream(downloadedPath);
-      https
-        .get(urlToDownload, (response) => {
-          response.pipe(file);
-          file.on("finish", () => {
-            file.close(() => {
-              console.log("File downloaded successfully");
-            });
-          });
-        })
-        .on("error", (err) => {
-          fs.unlink(downloadedPath, () => {
-            console.error("Error downloading file:", err);
-            throw err;
-          });
-        });
-    } catch (error) {
-      throw new Error(`Error downloading file from LlamaCloud: ${error}`);
-    }
-  }
-
-  private static async getFileUrlByName(
-    pipelineId: string,
-    name: string,
-  ): Promise<string | null> {
-    const files = await LLamaCloudFileService.getAllFiles(pipelineId);
-    const file = files.find((file) => file.name === name);
-    if (!file) return null;
-    return await LLamaCloudFileService.getFileUrlById(
-      file.project_id,
-      file.file_id,
-    );
-  }
-
-  private static async getFileUrlById(
-    projectId: string,
-    fileId: string,
-  ): Promise<string> {
-    const url = `${LLAMA_CLOUD_BASE_URL}/files/${fileId}/content?project_id=${projectId}`;
-    const response = await fetch(url, {
-      method: "GET",
-      headers: LLamaCloudFileService.headers,
-    });
-    const data = (await response.json()) as { url: string };
-    return data.url;
-  }
-
-  private static async getAllFiles(
-    pipelineId: string,
-  ): Promise<LlamaCloudFile[]> {
-    const url = `${LLAMA_CLOUD_BASE_URL}/pipelines/${pipelineId}/files`;
-    const response = await fetch(url, {
-      method: "GET",
-      headers: LLamaCloudFileService.headers,
-    });
-    const data = await response.json();
-    return data;
-  }
-
-  private static async getAllProjects(): Promise<LLamaCloudProject[]> {
-    const url = `${LLAMA_CLOUD_BASE_URL}/projects`;
-    const response = await fetch(url, {
-      method: "GET",
-      headers: LLamaCloudFileService.headers,
-    });
-    const data = (await response.json()) as LLamaCloudProject[];
-    return data;
-  }
-
-  private static async getAllPipelines(): Promise<LLamaCloudPipeline[]> {
-    const url = `${LLAMA_CLOUD_BASE_URL}/pipelines`;
-    const response = await fetch(url, {
-      method: "GET",
-      headers: LLamaCloudFileService.headers,
-    });
-    const data = (await response.json()) as LLamaCloudPipeline[];
-    return data;
-  }
-}

templates/components/loaders/python/db.py

@@ -1,8 +1,6 @@
-import os
 import logging
 from typing import List
-from pydantic import BaseModel, validator
-from llama_index.core.indices.vector_store import VectorStoreIndex
+from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 

templates/components/loaders/python/file.py

@@ -2,21 +2,16 @@ import os
 import logging
 from typing import Dict
 from llama_parse import LlamaParse
-from pydantic import BaseModel, validator
+from pydantic import BaseModel
+
+from app.config import DATA_DIR
 
 logger = logging.getLogger(__name__)
 
 
 class FileLoaderConfig(BaseModel):
-    data_dir: str = "data"
     use_llama_parse: bool = False
 
-    @validator("data_dir")
-    def data_dir_must_exist(cls, v):
-        if not os.path.isdir(v):
-            raise ValueError(f"Directory '{v}' does not exist")
-        return v
-
 
 def llama_parse_parser():
     if os.getenv("LLAMA_CLOUD_API_KEY") is None:
@@ -54,7 +49,7 @@ def get_file_documents(config: FileLoaderConfig):
 
             file_extractor = llama_parse_extractor()
         reader = SimpleDirectoryReader(
-            config.data_dir,
+            DATA_DIR,
             recursive=True,
             filename_as_id=True,
             raise_on_error=True,

templates/components/loaders/python/web.py

@@ -1,5 +1,3 @@
-import os
-import json
 from pydantic import BaseModel, Field
 
 

templates/components/settings/python/llmhub.py

@@ -6,11 +6,13 @@ import os
 DEFAULT_MODEL = "gpt-3.5-turbo"
 DEFAULT_EMBEDDING_MODEL = "text-embedding-3-large"
 
+
 class TSIEmbedding(OpenAIEmbedding):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._query_engine = self._text_engine = self.model_name
 
+
 def llm_config_from_env() -> Dict:
     from llama_index.core.constants import DEFAULT_TEMPERATURE
 
@@ -32,7 +34,7 @@ def llm_config_from_env() -> Dict:
 
 def embedding_config_from_env() -> Dict:
     from llama_index.core.constants import DEFAULT_EMBEDDING_DIM
-    
+
     model = os.getenv("EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL)
     dimension = os.getenv("EMBEDDING_DIM", DEFAULT_EMBEDDING_DIM)
     api_key = os.getenv("T_SYSTEMS_LLMHUB_API_KEY")
@@ -46,6 +48,7 @@ def embedding_config_from_env() -> Dict:
     }
     return config
 
+
 def init_llmhub():
     from llama_index.llms.openai_like import OpenAILike
 
@@ -58,4 +61,4 @@ def init_llmhub():
         is_chat_model=True,
         is_function_calling_model=False,
         context_window=4096,
-    )
+    )

templates/components/vectordbs/python/llamacloud/generate.py

@@ -1,3 +1,4 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 
 from app.engine.index import get_index

templates/components/vectordbs/python/none/generate.py

@@ -1,3 +1,4 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 
 load_dotenv()

templates/components/vectordbs/python/weaviate/vectordb.py

@@ -17,9 +17,11 @@ def _create_weaviate_client():
     client = weaviate.connect_to_weaviate_cloud(cluster_url, auth_credentials)
     return client
 
+
 # Global variable to store the Weaviate client
 client = None
 
+
 def get_vector_store():
     global client
     if client is None:

templates/components/vectordbs/typescript/llamacloud/generate.ts

@@ -1,24 +1,44 @@
 import * as dotenv from "dotenv";
-import { LlamaCloudIndex } from "llamaindex";
+import * as fs from "fs/promises";
+import { LLamaCloudFileService } from "llamaindex";
+import * as path from "path";
 import { getDataSource } from "./index";
-import { getDocuments } from "./loader";
+import { DATA_DIR } from "./loader";
 import { initSettings } from "./settings";
 import { checkRequiredEnvVars } from "./shared";
 
 dotenv.config();
 
-async function loadAndIndex() {
-  const documents = await getDocuments();
+async function* walk(dir: string): AsyncGenerator<string> {
+  const directory = await fs.opendir(dir);
 
-  await getDataSource();
-  await LlamaCloudIndex.fromDocuments({
-    documents,
-    name: process.env.LLAMA_CLOUD_INDEX_NAME!,
-    projectName: process.env.LLAMA_CLOUD_PROJECT_NAME!,
-    apiKey: process.env.LLAMA_CLOUD_API_KEY,
-    baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
-  });
-  console.log(`Successfully created embeddings!`);
+  for await (const dirent of directory) {
+    const entryPath = path.join(dir, dirent.name);
+
+    if (dirent.isDirectory()) {
+      yield* walk(entryPath); // Recursively walk through directories
+    } else if (dirent.isFile()) {
+      yield entryPath; // Yield file paths
+    }
+  }
+}
+
+async function loadAndIndex() {
+  const index = await getDataSource();
+  const projectId = await index.getProjectId();
+  const pipelineId = await index.getPipelineId();
+
+  // walk through the data directory and upload each file to LlamaCloud
+  for await (const filePath of walk(DATA_DIR)) {
+    const buffer = await fs.readFile(filePath);
+    const filename = path.basename(filePath);
+    const file = new File([buffer], filename);
+    await LLamaCloudFileService.addFileToPipeline(projectId, pipelineId, file, {
+      private: "false",
+    });
+  }
+
+  console.log(`Successfully uploaded documents to LlamaCloud!`);
 }
 
 (async () => {

templates/components/vectordbs/typescript/llamacloud/index.ts

@@ -18,6 +18,7 @@ export async function getDataSource(params?: LlamaCloudDataSourceParams) {
     );
   }
   const index = new LlamaCloudIndex({
+    organizationId: process.env.LLAMA_CLOUD_ORGANIZATION_ID,
     name: pipelineName,
     projectName,
     apiKey,

templates/components/vectordbs/typescript/llamacloud/queryFilter.ts

@@ -12,7 +12,7 @@ export function generateFilters(documentIds: string[]): MetadataFilters {
   if (!documentIds.length) return { filters: [publicDocumentsFilter] };
 
   const privateDocumentsFilter: MetadataFilter = {
-    key: "doc_id",
+    key: "file_id", // Note: LLamaCloud uses "file_id" to reference private document ids as "doc_id" is a restricted field in LlamaCloud
     value: documentIds,
     operator: "in",
   };

templates/types/extractor/fastapi/main.py

@@ -1,3 +1,4 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 
 load_dotenv()

templates/types/multiagent/fastapi/app/utils.py

@@ -5,4 +5,4 @@ def load_from_env(var: str, throw_error: bool = True) -> str:
     res = os.getenv(var)
     if res is None and throw_error:
         raise ValueError(f"Missing environment variable: {var}")
-    return res
+    return res

templates/types/multiagent/fastapi/main.py

@@ -1,3 +1,4 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 from app.settings import init_settings
 

templates/types/streaming/express/package.json

@@ -20,7 +20,7 @@
     "dotenv": "^16.3.1",
     "duck-duck-scrape": "^2.2.5",
     "express": "^4.18.2",
-    "llamaindex": "0.5.17",
+    "llamaindex": "0.5.19",
     "pdf2json": "3.0.5",
     "ajv": "^8.12.0",
     "@e2b/code-interpreter": "^0.0.5",

templates/types/streaming/express/src/controllers/chat-config.controller.ts

@@ -1,5 +1,5 @@
 import { Request, Response } from "express";
-import { LLamaCloudFileService } from "./llamaindex/streaming/service";
+import { LLamaCloudFileService } from "llamaindex";
 
 export const chatConfig = async (_req: Request, res: Response) => {
   let starterQuestions = undefined;
@@ -15,6 +15,11 @@ export const chatConfig = async (_req: Request, res: Response) => {
 };
 
 export const chatLlamaCloudConfig = async (_req: Request, res: Response) => {
+  if (!process.env.LLAMA_CLOUD_API_KEY) {
+    return res.status(500).json({
+      error: "env variable LLAMA_CLOUD_API_KEY is required to use LlamaCloud",
+    });
+  }
   const config = {
     projects: await LLamaCloudFileService.getAllProjectsWithPipelines(),
     pipeline: {

templates/types/streaming/express/src/controllers/chat-upload.controller.ts

@@ -3,12 +3,16 @@ import { getDataSource } from "./engine";
 import { uploadDocument } from "./llamaindex/documents/upload";
 
 export const chatUpload = async (req: Request, res: Response) => {
-  const { base64, params }: { base64: string; params?: any } = req.body;
-  if (!base64) {
+  const {
+    filename,
+    base64,
+    params,
+  }: { filename: string; base64: string; params?: any } = req.body;
+  if (!base64 || !filename) {
     return res.status(400).json({
-      error: "base64 is required in the request body",
+      error: "base64 and filename is required in the request body",
     });
   }
   const index = await getDataSource(params);
-  return res.status(200).json(await uploadDocument(index, base64));
+  return res.status(200).json(await uploadDocument(index, filename, base64));
 };

templates/types/streaming/fastapi/app/api/routers/models.py

@@ -1,12 +1,14 @@
 import logging
 import os
-from typing import Any, Dict, List, Literal, Optional, Set
+from typing import Any, Dict, List, Literal, Optional
 
 from llama_index.core.llms import ChatMessage, MessageRole
 from llama_index.core.schema import NodeWithScore
 from pydantic import BaseModel, Field, validator
 from pydantic.alias_generators import to_camel
 
+from app.config import DATA_DIR
+
 logger = logging.getLogger("uvicorn")
 
 
@@ -175,6 +177,7 @@ class SourceNodes(BaseModel):
                 "Warning: FILESERVER_URL_PREFIX not set in environment variables. Can't use file server"
             )
         file_name = metadata.get("file_name")
+
         if file_name and url_prefix:
             # file_name exists and file server is configured
             pipeline_id = metadata.get("pipeline_id")
@@ -184,11 +187,17 @@ class SourceNodes(BaseModel):
                 return f"{url_prefix}/output/llamacloud/{file_name}"
             is_private = metadata.get("private", "false") == "true"
             if is_private:
+                # file is a private upload
                 return f"{url_prefix}/output/uploaded/{file_name}"
-            return f"{url_prefix}/data/{file_name}"
-        else:
-            # fallback to URL in metadata (e.g. for websites)
-            return metadata.get("URL")
+            # file is from calling the 'generate' script
+            # Get the relative path of file_path to data_dir
+            file_path = metadata.get("file_path")
+            data_dir = os.path.abspath(DATA_DIR)
+            if file_path and data_dir:
+                relative_path = os.path.relpath(file_path, data_dir)
+                return f"{url_prefix}/data/{relative_path}"
+        # fallback to URL in metadata (e.g. for websites)
+        return metadata.get("URL")
 
     @classmethod
     def from_source_nodes(cls, source_nodes: List[NodeWithScore]):

templates/types/streaming/fastapi/app/api/routers/upload.py

@@ -14,14 +14,16 @@ logger = logging.getLogger("uvicorn")
 class FileUploadRequest(BaseModel):
     base64: str
     filename: str
-    params: Any
+    params: Any = None
 
 
 @r.post("")
 def upload_file(request: FileUploadRequest) -> List[str]:
     try:
         logger.info("Processing file")
-        return PrivateFileService.process_file(request.filename, request.base64, request.params)
+        return PrivateFileService.process_file(
+            request.filename, request.base64, request.params
+        )
     except Exception as e:
         logger.error(f"Error processing file: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail="Error processing file")

templates/types/streaming/fastapi/app/api/services/file.py

@@ -3,9 +3,7 @@ import mimetypes
 import os
 from io import BytesIO
 from pathlib import Path
-import time
-from typing import Any, Dict, List, Tuple
-from uuid import uuid4
+from typing import Any, List, Tuple
 
 
 from app.engine.index import get_index
@@ -14,7 +12,6 @@ from llama_index.core.ingestion import IngestionPipeline
 from llama_index.core.readers.file.base import (
     _try_loading_included_file_formats as get_file_loaders_map,
 )
-from llama_index.core.readers.file.base import default_file_metadata_func
 from llama_index.core.schema import Document
 from llama_index.indices.managed.llama_cloud.base import LlamaCloudIndex
 from llama_index.readers.file import FlatReader
@@ -50,12 +47,9 @@ class PrivateFileService:
         return base64.b64decode(data), extension
 
     @staticmethod
-    def store_and_parse_file(file_data, extension) -> List[Document]:
+    def store_and_parse_file(file_name, file_data, extension) -> List[Document]:
         # Store file to the private directory
         os.makedirs(PrivateFileService.PRIVATE_STORE_PATH, exist_ok=True)
-
-        # random file name
-        file_name = f"{uuid4().hex}{extension}"
         file_path = Path(os.path.join(PrivateFileService.PRIVATE_STORE_PATH, file_name))
 
         # write file
@@ -106,7 +100,9 @@ class PrivateFileService:
             ]
         else:
             # First process documents into nodes
-            documents = PrivateFileService.store_and_parse_file(file_data, extension)
+            documents = PrivateFileService.store_and_parse_file(
+                file_name, file_data, extension
+            )
             pipeline = IngestionPipeline()
             nodes = pipeline.run(documents=documents)
 

templates/types/streaming/fastapi/app/api/services/suggestion.py

@@ -25,7 +25,6 @@ class NextQuestions(BaseModel):
 
 
 class NextQuestionSuggestion:
-
     @staticmethod
     async def suggest_next_questions(
         messages: List[Message],

templates/types/streaming/fastapi/app/config.py

@@ -0,0 +1 @@
+DATA_DIR = "data"

templates/types/streaming/fastapi/app/engine/generate.py

@@ -1,3 +1,4 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -21,7 +22,6 @@ STORAGE_DIR = os.getenv("STORAGE_DIR", "storage")
 
 
 def get_doc_store():
-
     # If the storage directory is there, load the document store from it.
     # If not, set up an in-memory document store since we can't load from a directory that doesn't exist.
     if os.path.exists(STORAGE_DIR):

templates/types/streaming/fastapi/main.py

@@ -1,5 +1,8 @@
+# flake8: noqa: E402
 from dotenv import load_dotenv
 
+from app.config import DATA_DIR
+
 load_dotenv()
 
 import logging
@@ -42,15 +45,16 @@ if environment == "dev":
 
 def mount_static_files(directory, path):
     if os.path.exists(directory):
-        for dir, _, _ in os.walk(directory):
-            relative_path = os.path.relpath(dir, directory)
-            mount_path = path if relative_path == "." else f"{path}/{relative_path}"
-            logger.info(f"Mounting static files '{dir}' at {mount_path}")
-            app.mount(mount_path, StaticFiles(directory=dir), name=f"{dir}-static")
+        logger.info(f"Mounting static files '{directory}' at '{path}'")
+        app.mount(
+            path,
+            StaticFiles(directory=directory, check_dir=False),
+            name=f"{directory}-static",
+        )
 
 
 # Mount the data files to serve the file viewer
-mount_static_files("data", "/api/files/data")
+mount_static_files(DATA_DIR, "/api/files/data")
 # Mount the output files from tools
 mount_static_files("output", "/api/files/output")
 

templates/types/streaming/nextjs/app/api/chat/config/llamacloud/route.ts

@@ -1,10 +1,18 @@
+import { LLamaCloudFileService } from "llamaindex";
 import { NextResponse } from "next/server";
-import { LLamaCloudFileService } from "../../llamaindex/streaming/service";
 
 /**
  * This API is to get config from the backend envs and expose them to the frontend
  */
 export async function GET() {
+  if (!process.env.LLAMA_CLOUD_API_KEY) {
+    return NextResponse.json(
+      {
+        error: "env variable LLAMA_CLOUD_API_KEY is required to use LlamaCloud",
+      },
+      { status: 500 },
+    );
+  }
   const config = {
     projects: await LLamaCloudFileService.getAllProjectsWithPipelines(),
     pipeline: {

templates/types/streaming/nextjs/app/api/chat/upload/route.ts

@@ -10,11 +10,15 @@ export const dynamic = "force-dynamic";
 
 export async function POST(request: NextRequest) {
   try {
-    const { base64, params }: { base64: string; params?: any } =
+    const {
+      filename,
+      base64,
+      params,
+    }: { filename: string; base64: string; params?: any } =
       await request.json();
-    if (!base64) {
+    if (!base64 || !filename) {
       return NextResponse.json(
-        { error: "base64 is required in the request body" },
+        { error: "base64 and filename is required in the request body" },
         { status: 400 },
       );
     }
@@ -24,7 +28,7 @@ export async function POST(request: NextRequest) {
         `StorageContext is empty - call 'npm run generate' to generate the storage first`,
       );
     }
-    return NextResponse.json(await uploadDocument(index, base64));
+    return NextResponse.json(await uploadDocument(index, filename, base64));
   } catch (error) {
     console.error("[Upload API]", error);
     return NextResponse.json(

templates/types/streaming/nextjs/app/api/files/[...slug]/route.ts

@@ -1,6 +1,7 @@
 import { readFile } from "fs/promises";
 import { NextRequest, NextResponse } from "next/server";
 import path from "path";
+import { DATA_DIR } from "../../chat/engine/loader";
 
 /**
  * This API is to get file data from allowed folders
@@ -28,7 +29,11 @@ export async function GET(
   }
 
   try {
-    const filePath = path.join(process.cwd(), folder, path.join(...pathTofile));
+    const filePath = path.join(
+      process.cwd(),
+      folder === "data" ? DATA_DIR : folder,
+      path.join(...pathTofile),
+    );
     const blob = await readFile(filePath);
 
     return new NextResponse(blob, {

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-sources.tsx

@@ -13,8 +13,6 @@ import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
 import { DocumentFileType, SourceData, SourceNode } from "../index";
 import PdfDialog from "../widgets/PdfDialog";
 
-const SCORE_THRESHOLD = 0.25;
-
 type Document = {
   url: string;
   sources: SourceNode[];
@@ -24,15 +22,11 @@ export function ChatSources({ data }: { data: SourceData }) {
   const documents: Document[] = useMemo(() => {
     // group nodes by document (a document must have a URL)
     const nodesByUrl: Record<string, SourceNode[]> = {};
-    data.nodes
-      .filter((node) => (node.score ?? 1) > SCORE_THRESHOLD)
-      .filter((node) => isValidUrl(node.url))
-      .sort((a, b) => (b.score ?? 1) - (a.score ?? 1))
-      .forEach((node) => {
-        const key = node.url!.replace(/\/$/, ""); // remove trailing slash
-        nodesByUrl[key] ??= [];
-        nodesByUrl[key].push(node);
-      });
+    data.nodes.forEach((node) => {
+      const key = node.url;
+      nodesByUrl[key] ??= [];
+      nodesByUrl[key].push(node);
+    });
 
     // convert to array of documents
     return Object.entries(nodesByUrl).map(([url, sources]) => ({
@@ -55,11 +49,51 @@ export function ChatSources({ data }: { data: SourceData }) {
   );
 }
 
-function SourceNumberButton({ index }: { index: number }) {
+export function SourceInfo({
+  node,
+  index,
+}: {
+  node?: SourceNode;
+  index: number;
+}) {
+  if (!node) return <SourceNumberButton index={index} />;
   return (
-    <div className="text-xs w-5 h-5 rounded-full bg-gray-100 flex items-center justify-center hover:text-white hover:bg-primary ">
+    <HoverCard>
+      <HoverCardTrigger
+        className="cursor-default"
+        onClick={(e) => {
+          e.preventDefault();
+          e.stopPropagation();
+        }}
+      >
+        <SourceNumberButton
+          index={index}
+          className="hover:text-white hover:bg-primary"
+        />
+      </HoverCardTrigger>
+      <HoverCardContent className="w-[400px]">
+        <NodeInfo nodeInfo={node} />
+      </HoverCardContent>
+    </HoverCard>
+  );
+}
+
+export function SourceNumberButton({
+  index,
+  className,
+}: {
+  index: number;
+  className?: string;
+}) {
+  return (
+    <span
+      className={cn(
+        "text-xs w-5 h-5 rounded-full bg-gray-100 inline-flex items-center justify-center",
+        className,
+      )}
+    >
       {index + 1}
-    </div>
+    </span>
   );
 }
 
@@ -89,20 +123,7 @@ function DocumentInfo({ document }: { document: Document }) {
           {sources.map((node: SourceNode, index: number) => {
             return (
               <div key={node.id}>
-                <HoverCard>
-                  <HoverCardTrigger
-                    className="cursor-default"
-                    onClick={(e) => {
-                      e.preventDefault();
-                      e.stopPropagation();
-                    }}
-                  >
-                    <SourceNumberButton index={index} />
-                  </HoverCardTrigger>
-                  <HoverCardContent className="w-[400px]">
-                    <NodeInfo nodeInfo={node} />
-                  </HoverCardContent>
-                </HoverCard>
+                <SourceInfo node={node} index={index} />
               </div>
             );
           })}
@@ -125,13 +146,7 @@ function DocumentInfo({ document }: { document: Document }) {
 
   if (url.endsWith(".pdf")) {
     // open internal pdf dialog for pdf files when click document card
-    return (
-      <PdfDialog
-        documentId={document.url}
-        url={document.url}
-        trigger={DocumentDetail}
-      />
-    );
+    return <PdfDialog documentId={url} url={url} trigger={DocumentDetail} />;
   }
   // open external link when click document card for other file types
   return <div onClick={() => window.open(url, "_blank")}>{DocumentDetail}</div>;
@@ -179,13 +194,3 @@ function NodeInfo({ nodeInfo }: { nodeInfo: SourceNode }) {
     </div>
   );
 }
-
-function isValidUrl(url?: string): boolean {
-  if (!url) return false;
-  try {
-    new URL(url);
-    return true;
-  } catch (_) {
-    return false;
-  }
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx

@@ -11,10 +11,10 @@ import {
   ImageData,
   MessageAnnotation,
   MessageAnnotationType,
-  SourceData,
   SuggestedQuestionsData,
   ToolData,
   getAnnotationData,
+  getSourceAnnotationData,
 } from "../index";
 import ChatAvatar from "./chat-avatar";
 import { ChatEvents } from "./chat-events";
@@ -54,10 +54,9 @@ function ChatMessageContent({
     annotations,
     MessageAnnotationType.EVENTS,
   );
-  const sourceData = getAnnotationData<SourceData>(
-    annotations,
-    MessageAnnotationType.SOURCES,
-  );
+
+  const sourceData = getSourceAnnotationData(annotations);
+
   const toolData = getAnnotationData<ToolData>(
     annotations,
     MessageAnnotationType.TOOLS,
@@ -91,7 +90,7 @@ function ChatMessageContent({
     },
     {
       order: 0,
-      component: <Markdown content={message.content} />,
+      component: <Markdown content={message.content} sources={sourceData[0]} />,
     },
     {
       order: 3,

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/markdown.tsx

@@ -5,6 +5,8 @@ import rehypeKatex from "rehype-katex";
 import remarkGfm from "remark-gfm";
 import remarkMath from "remark-math";
 
+import { SourceData } from "..";
+import { SourceNumberButton } from "./chat-sources";
 import { CodeBlock } from "./codeblock";
 
 const MemoizedReactMarkdown: FC<Options> = memo(
@@ -34,12 +36,48 @@ const preprocessMedia = (content: string) => {
   return content.replace(/(sandbox|attachment|snt):/g, "");
 };
 
-const preprocessContent = (content: string) => {
-  return preprocessMedia(preprocessLaTeX(content));
+/**
+ * Update the citation flag [citation:id]() to the new format [citation:index](url)
+ */
+const preprocessCitations = (content: string, sources?: SourceData) => {
+  if (sources) {
+    const citationRegex = /\[citation:(.+?)\]\(\)/g;
+    let match;
+    // Find all the citation references in the content
+    while ((match = citationRegex.exec(content)) !== null) {
+      const citationId = match[1];
+      // Find the source node with the id equal to the citation-id, also get the index of the source node
+      const sourceNode = sources.nodes.find((node) => node.id === citationId);
+      // If the source node is found, replace the citation reference with the new format
+      if (sourceNode !== undefined) {
+        content = content.replace(
+          match[0],
+          `[citation:${sources.nodes.indexOf(sourceNode)}]()`,
+        );
+      } else {
+        // If the source node is not found, remove the citation reference
+        content = content.replace(match[0], "");
+      }
+    }
+  }
+  return content;
 };
 
-export default function Markdown({ content }: { content: string }) {
-  const processedContent = preprocessContent(content);
+const preprocessContent = (content: string, sources?: SourceData) => {
+  return preprocessCitations(
+    preprocessMedia(preprocessLaTeX(content)),
+    sources,
+  );
+};
+
+export default function Markdown({
+  content,
+  sources,
+}: {
+  content: string;
+  sources?: SourceData;
+}) {
+  const processedContent = preprocessContent(content, sources);
 
   return (
     <MemoizedReactMarkdown
@@ -80,6 +118,23 @@ export default function Markdown({ content }: { content: string }) {
             />
           );
         },
+        a({ href, children }) {
+          // If a text link starts with 'citation:', then render it as a citation reference
+          if (
+            Array.isArray(children) &&
+            typeof children[0] === "string" &&
+            children[0].startsWith("citation:")
+          ) {
+            const index = Number(children[0].replace("citation:", ""));
+            if (!isNaN(index)) {
+              return <SourceNumberButton index={index} />;
+            } else {
+              // citation is not looked up yet, don't render anything
+              return <></>;
+            }
+          }
+          return <a href={href}>{children}</a>;
+        },
       }}
     >
       {processedContent}

templates/types/streaming/nextjs/app/components/ui/chat/index.ts

@@ -1,4 +1,5 @@
 import { JSONValue } from "ai";
+import { isValidUrl } from "../lib/utils";
 import ChatInput from "./chat-input";
 import ChatMessages from "./chat-messages";
 
@@ -42,7 +43,7 @@ export type SourceNode = {
   metadata: Record<string, unknown>;
   score?: number;
   text: string;
-  url?: string;
+  url: string;
 };
 
 export type SourceData = {
@@ -83,9 +84,41 @@ export type MessageAnnotation = {
   data: AnnotationData;
 };
 
+const NODE_SCORE_THRESHOLD = 0.25;
+
 export function getAnnotationData<T extends AnnotationData>(
   annotations: MessageAnnotation[],
   type: MessageAnnotationType,
 ): T[] {
   return annotations.filter((a) => a.type === type).map((a) => a.data as T);
 }
+
+export function getSourceAnnotationData(
+  annotations: MessageAnnotation[],
+): SourceData[] {
+  const data = getAnnotationData<SourceData>(
+    annotations,
+    MessageAnnotationType.SOURCES,
+  );
+  if (data.length > 0) {
+    const sourceData = data[0] as SourceData;
+    if (sourceData.nodes) {
+      sourceData.nodes = preprocessSourceNodes(sourceData.nodes);
+    }
+  }
+  return data;
+}
+
+function preprocessSourceNodes(nodes: SourceNode[]): SourceNode[] {
+  // Filter source nodes has lower score
+  nodes = nodes
+    .filter((node) => (node.score ?? 1) > NODE_SCORE_THRESHOLD)
+    .filter((node) => isValidUrl(node.url))
+    .sort((a, b) => (b.score ?? 1) - (a.score ?? 1))
+    .map((node) => {
+      // remove trailing slash for node url if exists
+      node.url = node.url.replace(/\/$/, "");
+      return node;
+    });
+  return nodes;
+}

templates/types/streaming/nextjs/app/components/ui/lib/utils.ts

@@ -4,3 +4,13 @@ import { twMerge } from "tailwind-merge";
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
 }
+
+export function isValidUrl(url?: string): boolean {
+  if (!url) return false;
+  try {
+    new URL(url);
+    return true;
+  } catch (_) {
+    return false;
+  }
+}

templates/types/streaming/nextjs/package.json

@@ -25,7 +25,7 @@
     "duck-duck-scrape": "^2.2.5",
     "formdata-node": "^6.0.3",
     "got": "^14.4.1",
-    "llamaindex": "0.5.17",
+    "llamaindex": "0.5.19",
     "lucide-react": "^0.294.0",
     "next": "^14.2.4",
     "react": "^18.2.0",