feat: use latest create llama and llamaindex to support file upload (#101)

---------
Co-authored-by: Marcus Schiesser <marcus.schiesser@googlemail.com>
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.gitignore

@@ -51,3 +51,6 @@ dev
 app/api/chat/config/
 app/api/files/
 cl/
+
+# uploaded files
+output/

app/api/chat/engine/chat.ts

@@ -1,30 +1,29 @@
-import { ContextChatEngine, Settings, SimpleChatEngine } from "llamaindex";
+import { ContextChatEngine, Settings } from "llamaindex";
 import { getDataSource } from "./index";
-import { STORAGE_CACHE_DIR } from "./shared";
+import { generateFilters } from "@/cl/app/api/chat/engine/chat";
 
 interface ChatEngineOptions {
-  datasource?: string;
+  datasource: string;
+  documentIds?: string[];
 }
 
-export async function createChatEngine({ datasource }: ChatEngineOptions) {
-  if (datasource) {
-    const index = await getDataSource(datasource);
-    if (!index) {
-      throw new Error(
-        `No datasources found in storage cache folder: ${STORAGE_CACHE_DIR}/${datasource}. Run generate it first.`,
-      );
-    }
-    const retriever = index.asRetriever({
-      similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : 3,
-    });
-    return new ContextChatEngine({
-      chatModel: Settings.llm,
-      retriever,
-      systemPrompt: process.env.SYSTEM_PROMPT,
-    });
+export async function createChatEngine({
+  datasource,
+  documentIds,
+}: ChatEngineOptions) {
+  const index = await getDataSource(datasource);
+  if (!index) {
+    throw new Error(
+      `StorageContext is empty - call 'pnpm run generate ${datasource}' to generate the storage first`,
+    );
   }
-
-  return new SimpleChatEngine({
-    llm: Settings.llm,
+  const retriever = index.asRetriever({
+    similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : 3,
+    filters: generateFilters(documentIds || []),
+  });
+  return new ContextChatEngine({
+    chatModel: Settings.llm,
+    retriever,
+    systemPrompt: process.env.SYSTEM_PROMPT,
   });
 }

app/api/chat/engine/generate.ts

@@ -1,11 +1,8 @@
-import { VectorStoreIndex } from "llamaindex";
-import { storageContextFromDefaults } from "llamaindex/storage/StorageContext";
-
 import * as dotenv from "dotenv";
-
 import { getDocuments } from "./loader";
 import { initSettings } from "./settings";
-import { STORAGE_CACHE_DIR } from "./shared";
+import { storageContextFromDefaults, VectorStoreIndex } from "llamaindex";
+import { STORAGE_CACHE_DIR } from "@/cl/app/api/chat/engine/shared";
 
 // Load environment variables from local .env.development.local file
 dotenv.config({ path: ".env.development.local" });
@@ -30,7 +27,11 @@ async function generateDatasource() {
     const storageContext = await storageContextFromDefaults({
       persistDir: `${STORAGE_CACHE_DIR}/${datasource}`,
     });
-    const documents = await getDocuments();
+    const documents = await getDocuments(datasource);
+    //  Set private=false to mark the document as public (required for filtering)
+    documents.forEach((doc) => {
+      doc.metadata["private"] = "false";
+    });
     await VectorStoreIndex.fromDocuments(documents, {
       storageContext,
     });

app/api/chat/engine/index.ts

@@ -1,6 +1,6 @@
 import { SimpleDocumentStore, VectorStoreIndex } from "llamaindex";
 import { storageContextFromDefaults } from "llamaindex/storage/StorageContext";
-import { STORAGE_CACHE_DIR } from "./shared";
+import { STORAGE_CACHE_DIR } from "@/cl/app/api/chat/engine/shared";
 
 export async function getDataSource(datasource: string) {
   console.log(`Using datasource: ${datasource}`);

app/api/chat/engine/loader.ts

@@ -1,9 +1,9 @@
-import { SimpleDirectoryReader } from "llamaindex";
+import { SimpleDirectoryReader } from "llamaindex/readers/SimpleDirectoryReader";
 
 export const DATA_DIR = "./datasources";
 
-export async function getDocuments() {
+export async function getDocuments(datasource: string) {
   return await new SimpleDirectoryReader().loadData({
-    directoryPath: DATA_DIR,
+    directoryPath: `${DATA_DIR}/${datasource}`,
   });
 }

app/api/chat/engine/shared.ts

@@ -1 +0,0 @@
-export const STORAGE_CACHE_DIR = "./cache";

app/api/chat/route.ts

@@ -1,4 +1,4 @@
-import { Message, StreamData, StreamingTextResponse } from "ai";
+import { JSONValue, Message, StreamData, StreamingTextResponse } from "ai";
 import {
   ChatMessage,
   OpenAI,
@@ -9,14 +9,15 @@ import {
 import { NextRequest, NextResponse } from "next/server";
 import { createChatEngine } from "./engine/chat";
 import { initSettings } from "./engine/settings";
+import { LlamaIndexStream } from "@/cl/app/api/chat/llamaindex/streaming/stream";
 import {
-  LlamaIndexStream,
   convertMessageContent,
-} from "@/cl/app/api/chat/llamaindex-stream";
+  retrieveDocumentIds,
+} from "@/cl/app/api/chat/llamaindex/streaming/annotations";
 import {
   createCallbackManager,
   createStreamTimeout,
-} from "@/cl/app/api/chat/stream-helper";
+} from "@/cl/app/api/chat/llamaindex/streaming/events";
 import { LLMConfig } from "@/app/store/bot";
 
 initSettings();
@@ -41,22 +42,21 @@ export async function POST(request: NextRequest) {
     const { messages, context, modelConfig, datasource } =
       body as ChatRequestBody;
     const userMessage = messages.pop();
-    if (!messages || !userMessage || userMessage.role !== "user") {
+    if (
+      !messages ||
+      !userMessage ||
+      userMessage.role !== "user" ||
+      !datasource
+    ) {
       return NextResponse.json(
         {
           error:
-            "messages are required in the request body and the last message must be from the user",
+            "datasource and messages are required in the request body and the last message must be from the user",
         },
         { status: 400 },
       );
     }
 
-    // Create chat engine instance with llm config from request
-    const llm = new OpenAI(modelConfig);
-    const chatEngine = await Settings.withLLM(llm, async () => {
-      return await createChatEngine({ datasource });
-    });
-
     let annotations = userMessage.annotations;
     if (!annotations) {
       // the user didn't send any new annotations with the last message
@@ -70,6 +70,21 @@ export async function POST(request: NextRequest) {
         )?.annotations;
     }
 
+    // retrieve document Ids from the annotations of all messages (if any) and create chat engine with index
+    const allAnnotations: JSONValue[] = [...messages, userMessage].flatMap(
+      (message) => {
+        return message.annotations ?? [];
+      },
+    );
+
+    const ids = retrieveDocumentIds(allAnnotations);
+
+    // Create chat engine instance with llm config from request
+    const llm = new OpenAI(modelConfig);
+    const chatEngine = await Settings.withLLM(llm, async () => {
+      return await createChatEngine({ datasource, documentIds: ids });
+    });
+
     // Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
     const userMessageContent = convertMessageContent(
       userMessage.content,
@@ -95,7 +110,7 @@ export async function POST(request: NextRequest) {
     });
 
     // Transform LlamaIndex stream to Vercel/AI format
-    const stream = LlamaIndexStream(response, vercelStreamData);
+    const stream = LlamaIndexStream(response, vercelStreamData, chatMessages);
 
     // Return a StreamingTextResponse, which can be consumed by the Vercel/AI client
     return new StreamingTextResponse(stream, {}, vercelStreamData);

app/api/chat/upload/route.ts

@@ -0,0 +1,36 @@
+import { NextRequest, NextResponse } from "next/server";
+import { initSettings } from "../engine/settings";
+import { uploadDocument } from "@/cl/app/api/chat/llamaindex/documents/upload";
+import { getDataSource } from "../engine";
+
+initSettings();
+
+export const runtime = "nodejs";
+export const dynamic = "force-dynamic";
+
+// Custom upload API to use datasource from request body
+export async function POST(request: NextRequest) {
+  try {
+    const { base64, datasource }: { base64: string; datasource: string } =
+      await request.json();
+    if (!base64 || !datasource) {
+      return NextResponse.json(
+        { error: "base64 and datasource is required in the request body" },
+        { status: 400 },
+      );
+    }
+    const index = await getDataSource(datasource);
+    if (!index) {
+      throw new Error(
+        `StorageContext is empty - call 'pnpm run generate ${datasource}' to generate the storage first`,
+      );
+    }
+    return NextResponse.json(await uploadDocument(index, base64));
+  } catch (error) {
+    console.error("[Upload API]", error);
+    return NextResponse.json(
+      { error: (error as Error).message },
+      { status: 500 },
+    );
+  }
+}

app/components/chat/chat-session.tsx

@@ -1,7 +1,8 @@
 "use client";
 
+import { useBotStore } from "@/app/store/bot";
 import { useChatSession } from "./useChatSession";
-import { ChatInput, ChatMessages } from "@/cl/app/components/ui/chat";
+import { ChatMessages, ChatInput } from "@/cl/app/components/ui/chat";
 
 // Custom ChatSection for ChatLlamaindex
 export default function ChatSection() {
@@ -16,6 +17,8 @@ export default function ChatSection() {
     append,
     setInput,
   } = useChatSession();
+  const botStore = useBotStore();
+  const bot = botStore.currentBot();
   return (
     <div className="space-y-4 w-full h-full flex flex-col">
       <ChatMessages
@@ -33,6 +36,8 @@ export default function ChatSection() {
         messages={messages}
         append={append}
         setInput={setInput}
+        requestParams={{ datasource: bot.datasource }}
+        onFileError={(errMsg) => alert(errMsg)}
       />
     </div>
   );

app/components/chat/useChatSession.ts

@@ -13,7 +13,7 @@ export function useChatSession() {
   const { updateBotSession } = botStore;
 
   const [isFinished, setIsFinished] = useState(false);
-  const { chatAPI } = useClientConfig();
+  const { backend } = useClientConfig();
   const {
     messages,
     setMessages,
@@ -26,7 +26,7 @@ export function useChatSession() {
     append,
     setInput,
   } = useChat({
-    api: chatAPI,
+    api: `${backend}/api/chat`,
     headers: {
       "Content-Type": "application/json", // using JSON because of vercel/ai 2.2.26
     },

app/store/bot.data.ts

@@ -21,6 +21,7 @@ export const DEMO_BOTS: DemoBot[] = [
       sendMemory: false,
     },
     readOnly: true,
+    datasource: "documents",
   },
   {
     id: "3",
@@ -112,6 +113,7 @@ export const createEmptyBot = (): Bot => ({
   createdAt: Date.now(),
   botHello: Locale.Store.BotHello,
   session: createEmptySession(),
+  datasource: "documents",
 });
 
 export function createEmptySession(): ChatSession {

app/store/bot.ts

@@ -18,14 +18,15 @@ export const MESSAGE_ROLES: Message["role"][] = [
   "tool",
 ];
 
-export const ALL_MODELS = ["gpt-3.5-turbo", "gpt-4-turbo", "gpt-4o"] as const;
-
 export const AVAILABLE_DATASOURCES = [
+  "documents",
   "redhat",
   "watchos",
   "basic_law_germany",
 ] as const;
 
+export const ALL_MODELS = ["gpt-3.5-turbo", "gpt-4-turbo", "gpt-4o"] as const;
+
 export type ModelType = (typeof ALL_MODELS)[number];
 
 export interface LLMConfig {
@@ -52,7 +53,7 @@ export type Bot = {
   modelConfig: LLMConfig;
   readOnly: boolean;
   botHello: string | null;
-  datasource?: string;
+  datasource: string;
   share?: Share;
   createdAt?: number;
   session: ChatSession;

cache/documents/doc_store.json

@@ -0,0 +1 @@
+{"docstore/data":{"d4dee49a-1536-4348-8d05-71b6985fd634":{"__data__":"{\"id_\":\"d4dee49a-1536-4348-8d05-71b6985fd634\",\"metadata\":{},\"excludedEmbedMetadataKeys\":[],\"excludedLlmMetadataKeys\":[],\"relationships\":{},\"text\":\"blank document\",\"textTemplate\":\"\",\"metadataSeparator\":\"\\n\",\"type\":\"DOCUMENT\",\"hash\":\"N9F2ouMBTwtfSm/BOMv4qTJbZgwuWM/4d5f+KiHguP4=\"}","__type__":"DOCUMENT"},"98e3a153-0787-4b0e-8fdf-83138e1a8863":{"__data__":"{\"id_\":\"98e3a153-0787-4b0e-8fdf-83138e1a8863\",\"metadata\":{},\"excludedEmbedMetadataKeys\":[],\"excludedLlmMetadataKeys\":[],\"relationships\":{\"SOURCE\":{\"nodeId\":\"d4dee49a-1536-4348-8d05-71b6985fd634\",\"metadata\":{},\"hash\":\"N9F2ouMBTwtfSm/BOMv4qTJbZgwuWM/4d5f+KiHguP4=\"}},\"text\":\"blank document\",\"textTemplate\":\"\",\"endCharIdx\":14,\"metadataSeparator\":\"\\n\",\"type\":\"TEXT\",\"hash\":\"UrY0tomevnXgR8ns66YPzgXWWL3iNiejZ761QrrBDrA=\"}","__type__":"TEXT"}},"docstore/metadata":{"d4dee49a-1536-4348-8d05-71b6985fd634":{"docHash":"N9F2ouMBTwtfSm/BOMv4qTJbZgwuWM/4d5f+KiHguP4="},"98e3a153-0787-4b0e-8fdf-83138e1a8863":{"docHash":"UrY0tomevnXgR8ns66YPzgXWWL3iNiejZ761QrrBDrA=","refDocId":"d4dee49a-1536-4348-8d05-71b6985fd634"}},"docstore/ref_doc_info":{"d4dee49a-1536-4348-8d05-71b6985fd634":{"nodeIds":["98e3a153-0787-4b0e-8fdf-83138e1a8863"],"extraInfo":{}}}}

cache/documents/index_store.json

@@ -0,0 +1 @@
+{"docstore/data":{"4663e92f-59a2-4e7b-9ea0-bc442a7f3ef6":{"indexId":"4663e92f-59a2-4e7b-9ea0-bc442a7f3ef6","nodesDict":{"98e3a153-0787-4b0e-8fdf-83138e1a8863":{"id_":"98e3a153-0787-4b0e-8fdf-83138e1a8863","metadata":{},"excludedEmbedMetadataKeys":[],"excludedLlmMetadataKeys":[],"relationships":{"SOURCE":{"nodeId":"d4dee49a-1536-4348-8d05-71b6985fd634","metadata":{},"hash":"N9F2ouMBTwtfSm/BOMv4qTJbZgwuWM/4d5f+KiHguP4="}},"text":"blank document","textTemplate":"","endCharIdx":14,"metadataSeparator":"\n","type":"TEXT","hash":"UrY0tomevnXgR8ns66YPzgXWWL3iNiejZ761QrrBDrA="}},"type":"simple_dict"}}}

create-llama.sh

@@ -8,7 +8,7 @@ rm -rf app/api/files
 rm -rf cl
 
 # Run the node command with specified options
-npx -y create-llama@0.1.10 \
+npx -y create-llama@0.1.27 \
     --framework nextjs \
     --template streaming \
     --engine context \
@@ -30,4 +30,4 @@ cp -r cl/app/api/files app/api/files
 cp -r cl/app/api/chat/config app/api/chat/config
 
 # copy example .env file
-cp cl/.env .env.development.local
+cp cl/.env .env.development.local

next.config.mjs

@@ -5,6 +5,7 @@ const nextConfig = {
     serverComponentsExternalPackages: ["pdf-parse"],
     outputFileTracingIncludes: {
       "/*": ["./cache/**/*"],
+      "/api/**/*": ["node_modules/tiktoken/tiktoken_bg.wasm"]
     },
     outputFileTracingExcludes: {
       "/api/files/*": [".next/**/*", "node_modules/**/*", "public/**/*", "app/**/*"],

package.json

@@ -41,7 +41,7 @@
     "dotenv": "^16.4.5",
     "emoji-picker-react": "^4.9.2",
     "encoding": "^0.1.13",
-    "llamaindex": "^0.3.16",
+    "llamaindex": "0.5.12",
     "lucide-react": "^0.277.0",
     "mermaid": "^10.9.0",
     "nanoid": "^5.0.7",
@@ -79,7 +79,8 @@
     "@e2b/code-interpreter": "^0.0.5",
     "@apidevtools/swagger-parser": "^10.1.0",
     "got": "10.7.0",
-    "ajv": "^8.12.0"
+    "ajv": "^8.12.0",
+    "tiktoken": "^1.0.15"
   },
   "devDependencies": {
     "@types/node": "^20.12.7",