LLM initialized , and packages configured

Author: arnab2001

package.json

@@ -17,7 +17,14 @@
     "src/**/*.{js,jsx,ts,tsx}": "npm run lint-prettier"
   },
   "dependencies": {
+    "@langchain/pinecone": "^0.0.7",
+    "@pinecone-database/pinecone": "2.2.2",
+    "@langchain/community": "^0.2.15",
     "@radix-ui/react-slot": "^1.1.0",
+    "ai": "^2.1.34",
+    "ai-stream-experimental": "^2.2.2",
+    "langchain": "^0.2.8",
+    "git-repo-parser": "^2.0.6",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",
     "framer-motion": "^11.2.12",

src/lib/config.ts

@@ -0,0 +1,12 @@
+import z from 'zod';
+
+const envSchema = z.object({
+  OPENAI_API_KEY: z.string().trim().min(1),
+  PINECONE_API_KEY: z.string().trim().min(1),
+  PINECONE_ENVIRONMENT: z.string().trim().min(1),
+  PINECONE_INDEX_NAME: z.string().trim().min(1),
+  PDF_PATH: z.string().trim().min(1),
+  INDEX_INIT_TIMEOUT: z.coerce.number().min(1),
+});
+
+export const env = envSchema.parse(process.env);

src/lib/langchain.ts

@@ -0,0 +1,75 @@
+import { ConversationalRetrievalQAChain } from "langchain/chains";
+import { getVectorStore } from "./vector-store";
+import { getPineconeClient } from "./pinecone-client";
+import {
+  StreamingTextResponse,
+  experimental_StreamData,
+  LangChainStream,
+} from "ai-stream-experimental";
+import { streamingModel, nonStreamingModel } from "./llm";
+import { STANDALONE_QUESTION_TEMPLATE, QA_TEMPLATE } from "./prompt-templates";
+
+type callChainArgs = {
+  question: string;
+  chatHistory: string;
+};
+
+export async function callChain({ question, chatHistory }: callChainArgs) {
+  try {
+    // Open AI recommendation
+    const sanitizedQuestion = question.trim().replaceAll("\n", " ");
+    const pineconeClient = await getPineconeClient();
+    const vectorStore = await getVectorStore(pineconeClient);
+    const retriever = vectorStore.asRetriever({
+      searchKwargs: { k: 5 },  // Increase from default 4 to 5
+      searchType: "mmr",  // Use Maximum Marginal Relevance for diverse results
+      filter: { type: "code" }  // Add a filter if you've categorized your embeddings
+    });
+    const { stream, handlers } = LangChainStream({
+      experimental_streamData: true,
+    });
+    const data = new experimental_StreamData();
+
+    const chain = ConversationalRetrievalQAChain.fromLLM(
+      streamingModel,
+      vectorStore.asRetriever(),
+      {
+        qaTemplate: QA_TEMPLATE,
+        questionGeneratorTemplate: STANDALONE_QUESTION_TEMPLATE,
+        returnSourceDocuments: true, //default 4
+        questionGeneratorChainOptions: {
+          llm: nonStreamingModel,
+        },
+      }
+    );
+
+    // Question using chat-history
+    // Reference https://js.langchain.com/docs/modules/chains/popular/chat_vector_db#externally-managed-memory
+    chain
+      .call(
+        {
+          question: sanitizedQuestion,
+          chat_history: chatHistory,
+        },
+        [handlers]
+      )
+      .then(async (res) => {
+        const sourceDocuments = res?.sourceDocuments;
+        const firstTwoDocuments = sourceDocuments.slice(0, 2);
+        const pageContents = firstTwoDocuments.map(
+          ({ pageContent }: { pageContent: string }) => pageContent
+        );
+        console.log("already appended ", data);
+        data.append({
+          sources: pageContents,
+        });
+        data.close();
+      });
+
+    // Return the readable stream
+    return new StreamingTextResponse(stream, {}, data);
+  } catch (e) {
+    console.error(e);
+    throw new Error("Call chain method failed to execute successfully!!");
+  }
+}

src/lib/llm.ts

@@ -0,0 +1,36 @@
+import { ChatOpenAI } from "@langchain/openai";
+
+export const streamingModel = new ChatOpenAI({
+  modelName: "gpt-3.5-turbo",
+  streaming: true,
+  verbose: true,
+  temperature: 0,
+});
+
+export const nonStreamingModel = new ChatOpenAI({
+  modelName: "gpt-3.5-turbo",
+  verbose: true,
+  temperature: 0,
+});
+
+
+
+// const model = new ChatGoogleGenerativeAI({
+//   model: "gemini-pro",
+//   maxOutputTokens: 2048,
+//   apiKey,
+// });
+// export const streamingModel = new ChatGoogleGenerativeAI({
+//   model: "gemini-pro", // Replace with your desired Gemini model name
+//   streaming: true,
+//   verbose: true,
+//   temperature: 0.5,
+//   apiKey,
+// });
+
+// export const nonStreamingModel = new ChatGoogleGenerativeAI({
+//   model: "gemini-pro", // Replace with your desired Gemini model name
+//   verbose: true,
+//   temperature: 0.5,
+//   apiKey,
+// });

src/lib/pdf-loader.ts

@@ -0,0 +1,23 @@
+import { PDFLoader } from "@langchain/community/document_loaders/fs/pdf";
+import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
+import { env } from "./config";
+
+export async function getChunkedDocsFromPDF() {
+  try {
+    const loader = new PDFLoader(env.PDF_PATH);
+    const docs = await loader.load();
+
+    // From the docs https://www.pinecone.io/learn/chunking-strategies/
+    const textSplitter = new RecursiveCharacterTextSplitter({
+      chunkSize: 1000,
+      chunkOverlap: 200,
+    });
+
+    const chunkedDocs = await textSplitter.splitDocuments(docs);
+
+    return chunkedDocs;
+  } catch (e) {
+    console.error(e);
+    throw new Error("PDF docs chunking failed !");
+  }
+}

src/lib/pinecone-client.ts

@@ -0,0 +1,28 @@
+import { Pinecone } from '@pinecone-database/pinecone';
+import { env } from "./config";
+
+let pineconeClientInstance: Pinecone | null = null;
+
+// Initialize Pinecone client and connect to the index.
+async function initPineconeClient() {
+  try {
+    const pineconeClient = new Pinecone({
+      apiKey: env.PINECONE_API_KEY,
+    });
+
+    return pineconeClient;
+  } catch (error) {
+    console.error("Error initializing Pinecone client:", error);
+    throw new Error("Failed to initialize Pinecone Client");
+  }
+}
+
+export async function getPineconeClient() {
+  if (!pineconeClientInstance) {
+    pineconeClientInstance = await initPineconeClient();
+  }
+
+  return pineconeClientInstance;
+}
+
+// Example usage to ensure the client is working

src/lib/prompt-templates.ts

@@ -0,0 +1,27 @@
+// Creates a standalone question from the chat-history and the current question
+export const STANDALONE_QUESTION_TEMPLATE = `Given the following conversation history and a follow-up question, rephrase the follow-up question to be a standalone question that can be understood without the conversation context.
+
+Chat History:
+{chat_history}
+Follow-Up Input: {question}
+Standalone Question:`;
+
+// Actual question you ask the chat and send the response to client
+export const QA_TEMPLATE = `You are an AI assistant specializing in software development and project management. You have been provided with the codebase of a project. Use the following context to answer the question at the end.
+
+If the question is about code:
+1. Provide a brief explanation of the relevant code snippet.
+2. Suggest improvements or best practices if applicable.
+3. If asked about a specific function or feature, explain its purpose and how it fits into the larger project.
+
+If the question is about project management:
+1. Offer insights based on the code structure and organization.
+2. Suggest potential improvements in project architecture or workflow.
+
+If you don't know the answer, just say you don't know. DO NOT try to make up an answer.
+Always base your answers on the provided context, but you can briefly supplement with your general knowledge if relevant.
+
+{context}
+
+Question: {question}
+Helpful answer in markdown:`;

src/lib/utils.ts

@@ -1,6 +1,67 @@
 import { clsx, type ClassValue } from 'clsx';
 import { twMerge } from 'tailwind-merge';
+import { Message } from "ai";
+
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
 }
+
+export function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+export function scrollToBottom(containerRef: React.RefObject<HTMLElement>) {
+  if (containerRef.current) {
+    const lastMessage = containerRef.current.lastElementChild;
+    if (lastMessage) {
+      const scrollOptions: ScrollIntoViewOptions = {
+        behavior: "smooth",
+        block: "end",
+      };
+      lastMessage.scrollIntoView(scrollOptions);
+    }
+  }
+}
+
+// Reference:
+// github.com/hwchase17/langchainjs/blob/357d6fccfc78f1332b54d2302d92e12f0861c12c/examples/src/guides/expression_language/cookbook_conversational_retrieval.ts#L61
+export const formatChatHistory = (chatHistory: [string, string][]) => {
+  const formattedDialogueTurns = chatHistory.map(
+    (dialogueTurn) => `Human: ${dialogueTurn[0]}\nAssistant: ${dialogueTurn[1]}`
+  );
+
+  return formattedDialogueTurns.join("\n");
+};
+
+export function formattedText(inputText: string) {
+  return inputText
+    .replace(/\n+/g, " ") // Replace multiple consecutive new lines with a single space
+    .replace(/(\w) - (\w)/g, "$1$2") // Join hyphenated words together
+    .replace(/\s+/g, " "); // Replace multiple consecutive spaces with a single space
+}
+
+// Default UI Message
+export const initialMessages: Message[] = [
+  {
+    role: "assistant",
+    id: "0",
+    content:
+      "Hi! I am your Project Assistant. I am happy to help with your questions about your project's Codebase.",
+  },
+];
+
+interface Data {
+  sources: string[];
+}
+
+// Maps the sources with the right ai-message
+export const getSources = (data: Data[], role: string, index: number) => {
+  if (role === "assistant" && index >= 2 && (index - 2) % 2 === 0) {
+    const sourcesIndex = (index - 2) / 2;
+    if (data[sourcesIndex] && data[sourcesIndex].sources) {
+      return data[sourcesIndex].sources;
+    }
+  }
+  return [];
+};

src/lib/vector-store.ts

@@ -0,0 +1,53 @@
+import { env } from './config';
+import { OpenAIEmbeddings } from '@langchain/openai';
+import { PineconeStore } from "@langchain/pinecone";
+import { Pinecone } from '@pinecone-database/pinecone';
+
+export async function embedAndStoreDocs(
+  client: Pinecone,
+  docs: { id: string, text: string }[] // Adjusted to match the document structure
+) {
+  try {
+    console.log("Initializing embeddings...");
+    const embeddings = new OpenAIEmbeddings();
+    console.log("Embeddings initialized.");
+
+    console.log("Connecting to Pinecone index...");
+    const pineconeIndex = client.Index(env.PINECONE_INDEX_NAME);
+    console.log("Connected to Pinecone index.");
+
+    console.log("Embedding and storing documents...");
+    const documents = docs.map(doc => ({
+      id: doc.id,
+      text: doc.text,
+      pageContent: doc.text, // Adding pageContent as text
+      metadata: {} // Adding empty metadata
+    }));
+    await PineconeStore.fromDocuments(documents, embeddings, {
+      pineconeIndex,
+      textKey: 'text',
+    });
+    console.log("Documents embedded and stored.");
+  } catch (error) {
+    console.error("Error while embedding and storing documents: ", error);
+    throw new Error('Failed to load your docs!');
+  }
+}
+
+// Returns vector-store handle to be used as retrievers on langchains
+export async function getVectorStore(client: Pinecone) {
+  try {
+    const embeddings = new OpenAIEmbeddings();
+    const pineconeIndex = client.Index(env.PINECONE_INDEX_NAME);
+
+    const vectorStore = await PineconeStore.fromExistingIndex(embeddings, {
+      pineconeIndex,
+      textKey: 'text',
+    });
+
+    return vectorStore;
+  } catch (error) {
+    console.log('error ', error);
+    throw new Error('Something went wrong while getting vector store!');
+  }
+}