DA-831: Update LangChain to use CouchbaseSearchVectorStore

.gitignore

@@ -20,6 +20,7 @@
 .DS_Store
 *.pem
 .vscode
+.idea/
 
 # debug
 npm-debug.log*
@@ -28,6 +29,7 @@ yarn-error.log*
 
 # local env files
 .env*.local
+.env
 
 # vercel
 .vercel
@@ -36,4 +38,4 @@ yarn-error.log*
 *.tsbuildinfo
 next-env.d.ts
 
-/public/assets
+/public/assets

README.md

@@ -22,7 +22,7 @@ For RAG, we are using LangChain, Couchbase Vector Search & OpenAI. We fetch part
 
 - #### Set the environment secrets
 
-  Copy the `.env.template` file in and rename it to `.env` (`.env.local` in case of local development) and replace the placeholders with the actual values for your environment
+  Copy the `.env.template` file in and rename it to `.env` (`.env.local` in case of local development) and replace the placeholders with the actual values for your environment. **Note: the DB_CONN_STR must start with couchbase:// and should not contain the port number.**
 
   ```
   OPENAI_API_KEY=<open_ai_api_key>

app/api/chat/route.ts

@@ -1,16 +1,16 @@
 import {
-  CouchbaseVectorStore,
-  CouchbaseVectorStoreArgs,
-} from "@langchain/community/vectorstores/couchbase";
+  CouchbaseSearchVectorStore,
+} from "@langchain/community/vectorstores/couchbase_search";
 import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
 import { createCouchbaseCluster } from "@/lib/couchbase-connection";
 import { Message as VercelChatMessage } from "ai";
 import { HumanMessage, AIMessage, ChatMessage } from "@langchain/core/messages";
-import { createHistoryAwareRetriever } from "langchain/chains/history_aware_retriever";
 import { ChatPromptTemplate, MessagesPlaceholder } from "@langchain/core/prompts";
-import { createStuffDocumentsChain } from "langchain/chains/combine_documents";
-import { createRetrievalChain } from "langchain/chains/retrieval";
+import { createHistoryAwareRetriever } from "@langchain/classic/chains/history_aware_retriever";
+import { createStuffDocumentsChain } from "@langchain/classic/chains/combine_documents";
+import { createRetrievalChain } from "@langchain/classic/chains/retrieval";
 import { Document } from '@langchain/core/documents';
+import {NextResponse} from 'next/server';
 
 const formatVercelMessages = (message: VercelChatMessage) => {
   if (message.role === "user") {
@@ -26,6 +26,12 @@ const formatVercelMessages = (message: VercelChatMessage) => {
 };
 
 export async function POST(request: Request) {
+  // Load environment variables explicitly
+  const openaiApiKey = process.env.OPENAI_API_KEY;
+  if (!openaiApiKey) {
+    return NextResponse.json({ error: "OPENAI_API_KEY not set" }, { status: 500 });
+  }
+
   const body = await request.json();
   const messages = body.messages ?? [];
   if (!messages.length) {
@@ -38,10 +44,12 @@ export async function POST(request: Request) {
   const currentMessageContent = messages[messages.length - 1].content;
   try {
     const model = new ChatOpenAI({
-      model: "gpt-4"
+      model: "gpt-4",
+      openAIApiKey: openaiApiKey,
     });
     const embeddings = new OpenAIEmbeddings({
-      openAIApiKey: process.env.OPENAI_API_KEY,
+      apiKey: openaiApiKey,
+      model: "text-embedding-3-small",
     });
 
     const historyAwarePrompt = ChatPromptTemplate.fromMessages([
@@ -78,7 +86,12 @@ export async function POST(request: Request) {
     const scopedIndex = true;
 
     const cluster = await createCouchbaseCluster();
-    const couchbaseConfig: CouchbaseVectorStoreArgs = {
+
+    if (!cluster) {
+      throw new Error("Couchbase cluster connection failed");
+    }
+
+    const couchbaseConfig = {
       cluster,
       bucketName,
       scopeName,
@@ -88,7 +101,7 @@ export async function POST(request: Request) {
       embeddingKey,
       scopedIndex,
     };
-    const couchbaseVectorStore = await CouchbaseVectorStore.initialize(
+    const couchbaseSearchVectorStore = await CouchbaseSearchVectorStore.initialize(
       embeddings,
       couchbaseConfig
     );
@@ -98,7 +111,9 @@ export async function POST(request: Request) {
       resolveWithDocuments = resolve;
     });
 
-    const retriever = couchbaseVectorStore.asRetriever({
+    const retriever = couchbaseSearchVectorStore.asRetriever({
+      searchType: "similarity",
+      searchKwargs: { k: 4 },
       callbacks: [
         {
           handleRetrieverEnd(documents) {
@@ -108,15 +123,20 @@ export async function POST(request: Request) {
           },
         },
       ],
-    }
-    );
+    });
 
+    // Create the history-aware retriever chain
     const historyAwareRetrieverChain = await createHistoryAwareRetriever({
       llm: model,
       retriever,
       rephrasePrompt: historyAwarePrompt,
     });
 
+    // Use history-aware retriever only when there's chat history
+    const retrieverToUse = formattedPreviousMessages.length > 0 
+      ? historyAwareRetrieverChain 
+      : retriever;
+
     // Create a chain that answers questions using retrieved relevant documents as context.
     const documentChain = await createStuffDocumentsChain({
       llm: model,
@@ -125,7 +145,7 @@ export async function POST(request: Request) {
 
     // Create a chain that combines the above retriever and question answering chains.
     const conversationalRetrievalChain = await createRetrievalChain({
-      retriever: historyAwareRetrieverChain,
+      retriever: retrieverToUse,
       combineDocsChain: documentChain,
     });
 
@@ -164,5 +184,6 @@ export async function POST(request: Request) {
 
   } catch (err) {
     console.log("Error Received ", err);
+    return NextResponse.json({ error: "Failed to process chat message" });
   }
 }

app/api/ingestPdf/route.ts

@@ -1,10 +1,9 @@
 import { NextResponse } from "next/server";
-import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
+import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
 import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf';
 import {
-  CouchbaseVectorStore,
-  CouchbaseVectorStoreArgs,
-} from "@langchain/community/vectorstores/couchbase";
+  CouchbaseSearchVectorStore, CouchbaseSearchVectorStoreArgs,
+} from "@langchain/community/vectorstores/couchbase_search";
 import { OpenAIEmbeddings } from "@langchain/openai";
 import { createCouchbaseCluster } from "@/lib/couchbase-connection";
 import { writeFile } from "fs/promises";
@@ -52,7 +51,12 @@ export async function POST(request: Request) {
     const scopedIndex = true;
 
     const cluster = await createCouchbaseCluster();
-    const couchbaseConfig: CouchbaseVectorStoreArgs = {
+
+    if (!cluster) {
+      throw new Error("Couchbase cluster connection failed");
+    }
+
+    const couchbaseConfig: CouchbaseSearchVectorStoreArgs = {
       cluster,
       bucketName,
       scopeName,
@@ -62,7 +66,7 @@ export async function POST(request: Request) {
       embeddingKey,
       scopedIndex,
     };
-    await CouchbaseVectorStore.fromDocuments(docs, embeddings, couchbaseConfig);
+    await CouchbaseSearchVectorStore.fromDocuments(docs, embeddings, couchbaseConfig);
 
     console.log("creating vector store...");
   } catch (error) {

app/chatPage/page.tsx

@@ -7,7 +7,7 @@ import type {
   TransformToolbarSlot,
 } from "@react-pdf-viewer/toolbar";
 import { toolbarPlugin } from "@react-pdf-viewer/toolbar";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useRef, useState, Suspense } from "react";
 import { useSearchParams } from "next/navigation";
 import Image from "next/image";
 import { useChat } from "ai/react";
@@ -16,7 +16,7 @@ import LoadingDots from "@/components/LoadingDots";
 import "@react-pdf-viewer/core/lib/styles/index.css";
 import "@react-pdf-viewer/default-layout/lib/styles/index.css";
 
-const ChatPage = () => {
+const ChatContent = ({ searchParams }: { searchParams: URLSearchParams }) => {
   const [sourcesForMessages, setSourcesForMessages] = useState<
     Record<string, any>
   >({});
@@ -25,7 +25,7 @@ const ChatPage = () => {
   const toolbarPluginInstance = toolbarPlugin();
   const pageNavigationPluginInstance = pageNavigationPlugin();
   const { renderDefaultToolbar, Toolbar } = toolbarPluginInstance;
-  const searchParams = useSearchParams();
+
   useEffect(() => {
     setPdfUrl(("/assets/" + searchParams.get("fileName")) as string);
   }, [searchParams]);
@@ -237,4 +237,15 @@ const ChatPage = () => {
   );
 };
 
-export default ChatPage;
+export default function ChatPage() {
+  return (
+    <Suspense fallback={<div>Loading...</div>}>
+      <ChatPageContentWrapper />
+    </Suspense>
+  );
+}
+
+function ChatPageContentWrapper() {
+  const searchParams = useSearchParams();
+  return <ChatContent searchParams={searchParams} />;
+}

components/PDFUploader.tsx

@@ -1,13 +1,12 @@
 "use client";
 
 import { useRouter, useSearchParams } from "next/navigation";
-import React, { useCallback, useEffect, useState } from "react";
+import React, { useCallback, useEffect, useState, Suspense } from "react";
 import { useDropzone } from "react-dropzone";
 import { Loader } from "./Loader";
 
-const PDFUploader = () => {
+const PDFUploaderContent = ({ searchParams }: { searchParams: URLSearchParams }) => {
   const router = useRouter();
-  const searchParams = useSearchParams();
 
   const createQueryString = useCallback(
     (name: string, value: string) => {
@@ -98,4 +97,15 @@ const PDFUploader = () => {
   );
 };
 
-export default PDFUploader;
+export default function PDFUploader() {
+  return (
+    <Suspense fallback={<div>Loading...</div>}>
+      <PDFUploaderWrapper />
+    </Suspense>
+  );
+}
+
+function PDFUploaderWrapper() {
+  const searchParams = useSearchParams();
+  return <PDFUploaderContent searchParams={searchParams} />;
+}

lib/couchbase-connection.ts

@@ -1,6 +1,6 @@
 import { connect, Cluster } from "couchbase";
 
-export async function createCouchbaseCluster(): Promise<Cluster> {
+export async function createCouchbaseCluster(): Promise<Cluster | void> {
   const connectionString = process.env.DB_CONN_STR;
   const databaseUsername = process.env.DB_USERNAME;
   const databasePassword = process.env.DB_PASSWORD;
@@ -23,11 +23,16 @@ export async function createCouchbaseCluster(): Promise<Cluster> {
     );
   }
 
-  const cluster = await connect(connectionString, {
-    username: databaseUsername,
-    password: databasePassword,
-    configProfile: "wanDevelopment",
-  });
+  try {
+    return await connect(connectionString, {
+      username: databaseUsername,
+      password: databasePassword,
+      configProfile: "wanDevelopment",
+    });
+  } catch (e) {
+    throw new Error(
+      `Could not connect to the Couchbase cluster. Please check your DB_CONN_STR, DB_USERNAME, and DB_PASSWORD environment variables. \n${e}`,
+    )
+  }
 
-  return cluster;
 }