integrate langchain for RAG (#22)

* integrate langchain for RAG

* fix pipeline

.env.example

@@ -0,0 +1,10 @@
+OPENAI_API_KEY=""
+
+KV_URL=""
+KV_REST_API_URL=""
+KV_REST_API_TOKEN=""
+KV_REST_API_READ_ONLY_TOKEN=""
+
+SUPABASE_URL=""
+SUPABASE_ANON_KEY=""
+

.github/workflows/ci.yml

@@ -23,5 +23,3 @@ jobs:
         run: bun run check-types
       - name: Lint
         run: bun run lint
-      - name: Build
-        run: bun run build

.github/workflows/sync.yml

@@ -0,0 +1,20 @@
+name: SYNC
+
+on:
+  push:
+    branches:
+      - master
+
+jobs:
+  sync:
+    name: Generate Embeddings
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup Bun
+        uses: oven-sh/setup-bun@v1
+      - name: Install Dependencies
+        run: bun install
+      - name: Run Script
+        run: bun run ./generate-embeddings.ts

app/api/chat/route.ts

@@ -0,0 +1,159 @@
+import { kv } from '@vercel/kv';
+import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
+import { Ratelimit } from '@upstash/ratelimit';
+import { LangChainAdapter, Message, StreamingTextResponse } from 'ai';
+import { AIMessage, HumanMessage } from '@langchain/core/messages';
+import { Document } from '@langchain/core/documents';
+import { PromptTemplate } from '@langchain/core/prompts';
+import { createClient } from '@supabase/supabase-js';
+import { SupabaseVectorStore } from '@langchain/community/vectorstores/supabase';
+import { RunnableSequence } from '@langchain/core/runnables';
+import { StringOutputParser } from '@langchain/core/output_parsers';
+import { RedisCache } from '@langchain/community/caches/ioredis';
+import { getIp } from '@/lib/server/utils';
+import { env } from '@/config/env';
+
+export const dynamic = 'force-dynamic';
+export const maxDuration = 60;
+
+function combineDocumentsFn(docs: Document[]) {
+  const serializedDocs = docs.map((doc) => doc.pageContent);
+  return serializedDocs.join('\n\n');
+}
+
+function formatChatHistory(chatHistory: Message[]) {
+  const formattedDialogueTurns = chatHistory.map((m) =>
+    m.role === 'user' ? new HumanMessage(m.content) : new AIMessage(m.content)
+  );
+  return formattedDialogueTurns.join('\n\n');
+}
+
+const CONDENSE_QUESTION_TEMPLATE = `Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
+
+<chat_history>
+  {chat_history}
+</chat_history>
+
+Follow Up Input: {question}
+Standalone question:
+`;
+const condenseQuestionPrompt = PromptTemplate.fromTemplate(
+  CONDENSE_QUESTION_TEMPLATE
+);
+
+const ANSWER_TEMPLATE = `You are an AI Assistant called Roronova Zoro. You are the chat bot prensent on personal portfolio website and you answer questions only related to the portfolio.
+The name of the Owner of this Website is Nikhil S.
+
+Whenever it makes sense, provide links to pages that contain more information about the topic from the given context. If the question is out of context inform user accordingly.
+
+Format your messages in markdown format.
+
+Answer the question based only on the following context and chat history:
+<context>
+  {context}
+</context>
+
+<chat_history>
+  {chat_history}
+</chat_history>
+
+Question: {question}
+`;
+const answerPrompt = PromptTemplate.fromTemplate(ANSWER_TEMPLATE);
+
+const ratelimit = new Ratelimit({
+  redis: kv,
+  limiter: Ratelimit.fixedWindow(2, '1m'),
+});
+
+export async function POST(req: Request) {
+  try {
+    const ip = getIp();
+
+    const { success, limit, reset, remaining } = await ratelimit.limit(
+      `ratelimit_${ip ?? 'anonymous'}`
+    );
+
+    if (!success) {
+      return new Response('You have reached your request limit for the day.', {
+        status: 429,
+        headers: {
+          'X-RateLimit-Limit': limit.toString(),
+          'X-RateLimit-Remaining': remaining.toString(),
+          'X-RateLimit-Reset': reset.toString(),
+        },
+      });
+    }
+
+    const {
+      messages,
+    }: {
+      messages: Message[];
+    } = await req.json();
+
+    const previousMessages = messages.slice(0, -1);
+    const currentMessageContent = messages[messages.length - 1].content;
+
+    const cache = new RedisCache({
+      client: kv,
+    });
+
+    const model = new ChatOpenAI({
+      model: 'gpt-3.5-turbo',
+      temperature: 0,
+      streaming: true,
+      verbose: true,
+      cache,
+    });
+
+    const client = createClient(env.SUPABASE_URL, env.SUPABASE_ANON_KEY);
+
+    const vectorstore = new SupabaseVectorStore(new OpenAIEmbeddings(), {
+      client,
+      tableName: 'documents',
+      queryName: 'match_documents',
+    });
+
+    const standaloneQuestionChain = RunnableSequence.from([
+      condenseQuestionPrompt,
+      model,
+      new StringOutputParser(),
+    ]);
+
+    const retriever = vectorstore.asRetriever();
+
+    const retrievalChain = retriever.pipe(combineDocumentsFn);
+
+    const answerChain = RunnableSequence.from([
+      {
+        context: RunnableSequence.from([
+          (input) => input.question,
+          retrievalChain,
+        ]),
+        chat_history: (input) => input.chat_history,
+        question: (input) => input.question,
+      },
+      answerPrompt,
+      model,
+    ]);
+
+    const conversationalRetrievalQAChain = RunnableSequence.from([
+      {
+        question: standaloneQuestionChain,
+        chat_history: (input) => input.chat_history,
+      },
+      answerChain,
+    ]);
+
+    const stream = await conversationalRetrievalQAChain.stream({
+      question: currentMessageContent,
+      chat_history: formatChatHistory(previousMessages),
+    });
+
+    return new StreamingTextResponse(LangChainAdapter.toAIStream(stream));
+  } catch (error: any) {
+    return new Response(error?.message ?? 'Something went wrong', {
+      status: error?.status ?? 500,
+    });
+  }
+}

app/bot/chat.tsx

@@ -1,93 +1,86 @@
 'use client';
 
-import { useState } from 'react';
-import { AssistantContent, UserContent, type CoreMessage } from 'ai';
-import { readStreamableValue } from 'ai/rsc';
+import { useChat } from 'ai/react';
 import { toast } from 'sonner';
+import Markdown from 'react-markdown';
 import { ScrollArea } from '@/components/ui/scroll-area';
 import { Input } from '@/components/ui/input';
 import { Button } from '@/components/ui/button';
-import { continueConversation } from './actions';
+import { useEffect, useRef } from 'react';
+import { LoadingSpinner } from '@/assets/icons';
 
-function BotMessage({ content }: { content: AssistantContent }) {
+function BotMessage({ content }: { content: string }) {
   return (
     <div className='max-w-full'>
       <p className='max-w-max whitespace-pre-wrap rounded-lg bg-gray-100 p-2 dark:bg-gray-800'>
-        {content as string}
+        <Markdown>{content}</Markdown>
       </p>
     </div>
   );
 }
 
-function UserMessage({ content }: { content: UserContent }) {
+function UserMessage({ content }: { content: string }) {
   return (
     <div className='max-w-full'>
       <p className='ml-auto max-w-max whitespace-pre-wrap rounded-lg bg-gray-800 p-2 text-gray-100 dark:bg-gray-100 dark:text-gray-800'>
-        {content as string}
+        {content}
       </p>
     </div>
   );
 }
 
 export function Chat() {
-  const [messages, setMessages] = useState<CoreMessage[]>([]);
-  const [input, setInput] = useState('');
+  const { messages, input, handleInputChange, handleSubmit, isLoading } =
+    useChat({
+      onError: (e) => {
+        toast(e.message);
+      },
+    });
+
+  const scrollRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    const scrollArea = scrollRef.current?.querySelector(
+      '[data-radix-scroll-area-viewport]'
+    );
+    if (scrollArea) {
+      scrollArea.scrollTop = scrollArea.scrollHeight;
+    }
+  }, [messages]);
+
   return (
-    <div className='flex h-[55vh] w-full flex-col gap-4'>
+    <div className='flex h-[70vh] w-full flex-col gap-4'>
       <div className='flex-1 overflow-auto'>
-        <ScrollArea className='h-full w-full rounded-md border'>
+        <ScrollArea className='h-full w-full rounded-md border' ref={scrollRef}>
           <div className='p-4 text-sm'>
             <div className='grid gap-4'>
-              {/* <BotMessage content={'Coming soon...'} /> */}
-              {messages.map((m, i) =>
-                m.role === 'user' ? (
-                  <UserMessage key={i} content={m.content} />
+              {messages.map(({ id, content, role }) =>
+                role === 'user' ? (
+                  <UserMessage key={id} content={content} />
                 ) : (
-                  <BotMessage key={i} content={m.content as AssistantContent} />
+                  <BotMessage key={id} content={content} />
                 )
               )}
             </div>
           </div>
         </ScrollArea>
       </div>
-      <form
-        className='flex items-center'
-        action={async () => {
-          const newMessages: CoreMessage[] = [
-            ...messages,
-            { content: input, role: 'user' },
-          ];
-
-          setMessages(newMessages);
-          setInput('');
-
-          const result = await continueConversation(newMessages);
-
-          if ('error' in result) {
-            setMessages((messages) => messages.slice(0, -1));
-            toast.error(result.error as string);
-            return;
-          }
-
-          for await (const content of readStreamableValue(result)) {
-            setMessages([
-              ...newMessages,
-              {
-                role: 'assistant',
-                content: content as string,
-              },
-            ]);
-          }
-        }}
-      >
+      <form className='flex items-center gap-3' onSubmit={handleSubmit}>
         <Input
-          className='mr-4 flex-1'
+          className=''
           placeholder='Type your message...'
           type='text'
           value={input}
-          onChange={(e) => setInput(e.target.value)}
+          onChange={handleInputChange}
+          disabled={isLoading}
         />
-        <Button size='sm'>Ask Zoro</Button>
+        <Button
+          disabled={isLoading}
+          size='sm'
+          className='w-[10%] shrink-0 text-center'
+        >
+          {isLoading ? <LoadingSpinner /> : 'Ask Zoro'}
+        </Button>
       </form>
     </div>
   );