Feat: RAG search (#6)

* chore: package updates

* chore: prisma chat thread and suggested articles

* feat: new ui components

* feat: context chat store

* feat: vector store lib

* feat: chat lib methods

* feat: page and api updates

* chore: tailwind css

* chore: radix ui packages

* feat: delete chat

* chore: vercel ai

* feat: chat completion

* chore: package updates

* feat: data migration script update

* feat: chat lib updates

* feat: metadata lib update

* feat: page & component updates

* feat: trigger dev script updates

* feat: chat history keep previous data

* chore: pulled prisma schema

* chore: pinecone and llamaindex packages

* feat: answear generation improvements

* fix: missing schema

* Revert "chore: pinecone and llamaindex packages"

This reverts commit 3c56e49.

* chore: pinecone package

* feat: local chat history handling

* fix: temporal weight

* feat: extending prompt context

* feat: links and performance

* feat: demo chat component

* fix: array reverse fix

* fix: default prompt change

.env.template

@@ -21,4 +21,9 @@ NEXT_PUBLIC_POSTHOG_KEY=
 NEXT_PUBLIC_POSTHOG_HOST=
 
 PINECONE_API_KEY=
+PINECONE_INDEX_NAME=
+PINECONE_ENVIRONMENT=
+PINECONE_CHUNK_SIZE=
+PINECONE_CHUNK_OVERLAP=
+
 SEQUIN_WEBHOOK_SECRET=

app/api/articles/[id]/summary/route.ts

@@ -13,13 +13,16 @@ const getSummary: NextRouteFunction<Params> = async (_, { params }) => {
     return Response.json({ data: summary.generated_summary }, { status: 200 });
   }
 
-  const pdfLink = await articleService.getArticlePdfLink(params.id);
+  const metadata = await articleService.getArticleWithPdfLink(params.id);
 
-  if (!pdfLink) {
+  if (!metadata?.pdfLink) {
     return Response.json('Article not found', { status: 404 });
   }
 
-  const pdfNodes = await loadPDF(pdfLink, { metadata_id: params.id });
+  const pdfNodes = await loadPDF(metadata.pdfLink, {
+    metadata_id: params.id,
+    published: metadata.published,
+  });
 
   logger.log('PDF file loaded successfully!');
 

app/api/categories/get-by-groups/route.ts

@@ -1,4 +1,4 @@
-import { NextRouteFunction, routeValidator } from '@/lib/route-validator.server';
+import { NextRouteFunction } from '@/lib/route-validator.server';
 import * as categoryService from '@/lib/categories/categories.server';
 
 const getCategoriesByGroup: NextRouteFunction<{}> = async (request) => {

app/api/chat/[slug]/completion/route.ts

@@ -0,0 +1,53 @@
+import { LlamaIndexAdapter, StreamData } from 'ai';
+import { NextRouteFunction } from '@/lib/route-validator.server';
+import * as logger from '@/lib/logger';
+import { chatCompletionSchema } from './schema';
+import * as threadService from '@/lib/chat/thread.server';
+import * as chatService from '@/lib/chat/chat.server';
+import * as queryService from '@/lib/chat/query.server';
+import { chat_message_role } from '@prisma/client';
+
+type Params = { params: { slug: string } };
+
+const chatCompletion: NextRouteFunction<Params> = async (req, { params }) => {
+  const reqJson = await req.json();
+
+  const parsedParams = chatCompletionSchema.parse(reqJson);
+
+  const temporalAnalysis = await queryService.analyzeTemporalQuery(parsedParams.prompt);
+
+  const thread = await threadService.getMessagesByThreadSlug(params.slug);
+
+  const suggestions = await chatService.getDocumentSuggestions(
+    thread.chat_message,
+    parsedParams.prompt,
+    temporalAnalysis
+  );
+
+  const stream = await chatService.chatCompletion(thread.chat_message, parsedParams.prompt, true, suggestions);
+
+  const data = new StreamData();
+
+  suggestions.forEach((document) => {
+    data.appendMessageAnnotation(JSON.stringify({ type: 'article_metadata', article_metadata: document }));
+  });
+
+  return LlamaIndexAdapter.toDataStreamResponse(stream, {
+    data,
+    callbacks: {
+      onCompletion: async (response) => {
+        await threadService.createMessageWithSuggestions(
+          params.slug,
+          response,
+          chat_message_role.ASSISTANT,
+          suggestions
+        );
+      },
+      onFinal: async () => {
+        await data.close();
+      },
+    },
+  });
+};
+
+export const POST = chatCompletion;

app/api/chat/[slug]/completion/schema.ts

@@ -0,0 +1,8 @@
+import { escapeHtml } from '@/lib/utils';
+import { z } from 'zod';
+
+export const chatCompletionSchema = z.object({
+  prompt: z.string().transform(escapeHtml),
+});
+
+export type ChatCompletion = z.infer<typeof chatCompletionSchema>;

app/api/chat/create-thread/route.ts

@@ -0,0 +1,17 @@
+import { NextRouteFunction } from '@/lib/route-validator.server';
+import * as threadService from '@/lib/chat/thread.server';
+import { createNewThreadSchema } from './schema';
+
+type Params = { params: { prompt: string } };
+
+const createNewThread: NextRouteFunction<{}, Params> = async (req) => {
+  const reqJson = await req.json();
+
+  const parsedParams = createNewThreadSchema.parse(reqJson);
+
+  const thread = await threadService.create(parsedParams.prompt);
+
+  return Response.json({ data: thread }, { status: 200 });
+};
+
+export const POST = createNewThread;

app/api/chat/create-thread/schema.ts

@@ -0,0 +1,8 @@
+import { escapeHtml } from '@/lib/utils';
+import { z } from 'zod';
+
+export const createNewThreadSchema = z.object({
+  prompt: z.string().transform(escapeHtml),
+});
+
+export type CreateNewThread = z.infer<typeof createNewThreadSchema>;

app/api/chat/delete/[slug]/route.ts

@@ -0,0 +1,16 @@
+import { NextRouteFunction } from '@/lib/route-validator.server';
+import * as threadService from '@/lib/chat/thread.server';
+
+type Params = { params: { slug: string } };
+
+const deleteThread: NextRouteFunction<Params> = async (_, { params }) => {
+  try {
+    await threadService.deleteThread(params.slug);
+  } catch (error) {
+    return Response.json({ status: 500, error });
+  }
+
+  return Response.json({ status: 200 });
+};
+
+export const DELETE = deleteThread;

app/api/chat/route.ts

@@ -0,0 +1,17 @@
+import { NextRouteFunction } from '@/lib/route-validator.server';
+import * as chatService from '@/lib/chat/chat.server';
+
+const getThreads: NextRouteFunction<{}> = async (req) => {
+  const searchParams = req.nextUrl.searchParams;
+  const limit = searchParams.get('limit');
+
+  try {
+    const threads = await chatService.getChatHistory({ limit: limit ? parseInt(limit) : undefined });
+
+    return Response.json({ status: 200, data: threads });
+  } catch (error) {
+    return Response.json({ status: 500, error });
+  }
+};
+
+export const GET = getThreads;

app/api/embeddings/[id]/route.ts

@@ -1,19 +1,22 @@
 import { NextRouteFunction } from '@/lib/route-validator.server';
-import { getArticlePdfLink } from '@/lib/article-metadata/metadata.server';
+import { getArticleWithPdfLink } from '@/lib/article-metadata/metadata.server';
 import * as logger from '@/lib/logger';
 import { loadPDF } from '@/lib/file-handlers';
 import { addNewEmbeddings } from '@/lib/embeddings/embeddings.server';
 
 type Params = { params: { id: string } };
 
 const generateEmbeddings: NextRouteFunction<Params> = async (_, { params }) => {
-  const pdfLink = await getArticlePdfLink(params.id);
+  const metadata = await getArticleWithPdfLink(params.id);
 
-  if (!pdfLink) {
+  if (!metadata?.pdfLink) {
     return Response.json('Article not found', { status: 404 });
   }
 
-  const pdfNodes = await loadPDF(pdfLink, { metadata_id: params.id });
+  const pdfNodes = await loadPDF(metadata.pdfLink, {
+    metadata_id: params.id,
+    published: metadata.published,
+  });
 
   logger.log('PDF file loaded successfully!');
 

app/articles/[slug]/page.tsx

@@ -15,7 +15,7 @@ type Props = {
   params: { slug: string };
 };
 
-export default async function ArticleDetails({ params }: Props) {
+export default async function ArticleDetails({ params }: Readonly<Props>) {
   const article = await getArticleMetadataBySlug(params.slug);
 
   if (!article) {

app/articles/components/articles.tsx

@@ -18,9 +18,7 @@ export const Articles: React.FC<Props> = async ({ searchParams }) => {
 
   return (
     <div className="flex flex-col px-8 py-12 sm:py-16 gap-4">
-      <Suspense>
-        <Toolbar categoryTree={categoryTree} authors={authorList} searchParams={parsedSearchParams} />
-      </Suspense>
+      <Toolbar categoryTree={categoryTree} authors={authorList} searchParams={parsedSearchParams} />
 
       <NewsletterSection closable />
 

app/articles/components/toolbar.tsx

@@ -1,7 +1,7 @@
 'use client';
 
-import React, { useEffect, useState } from 'react';
-import { ChevronsDownUpIcon, ChevronsUpDownIcon, Settings2Icon } from 'lucide-react';
+import React, { startTransition, useEffect, useState } from 'react';
+import { ChevronsDownUpIcon, ChevronsUpDownIcon, MessageCircleIcon, Settings2Icon } from 'lucide-react';
 import { usePathname, useRouter } from 'next/navigation';
 import { useDebounce } from 'use-debounce';
 import posthog from 'posthog-js';
@@ -36,7 +36,7 @@ const isFiltered = (searchParams: ArticleMetadataSearch): boolean => {
 
 export const Toolbar: React.FC<Props> = ({ categoryTree, authors, searchParams }) => {
   const [isFilterOpen, setIsFilterOpen] = useState(isFiltered(searchParams));
-  const { isCompactMode, toggleCompactMode } = useBoundStore();
+  const { isCompactMode, toggleCompactMode, isContextChatOpen, toggleContextChat } = useBoundStore();
 
   const pathname = usePathname();
   const { replace } = useRouter();
@@ -71,7 +71,9 @@ export const Toolbar: React.FC<Props> = ({ categoryTree, authors, searchParams }
       to: searchState.to,
     });
 
-    replace(`${pathname}?${queryParams}`);
+    startTransition(() => {
+      replace(`${pathname}?${queryParams}`);
+    });
   }, [queryParams]);
 
   const resetFilters = () =>
@@ -138,6 +140,13 @@ export const Toolbar: React.FC<Props> = ({ categoryTree, authors, searchParams }
           onClick={toggleCompactMode}>
           {isCompactMode ? <ChevronsUpDownIcon className="h-4 w-4" /> : <ChevronsDownUpIcon className="h-4 w-4" />}
         </Button>
+        <Button
+          variant={isContextChatOpen ? 'default' : 'outline'}
+          size="sm"
+          className="h-8 font-normal"
+          onClick={toggleContextChat}>
+          <MessageCircleIcon className="h-4 w-4" />
+        </Button>
       </div>
       {isFilterOpen ? (
         <div className="flex flex-col gap-2 md:flex-row md:gap-0 md:items-center md:justify-between">
@@ -179,9 +188,6 @@ export const Toolbar: React.FC<Props> = ({ categoryTree, authors, searchParams }
               </Button>
             )}
           </div>
-          {/* <div className="flex items-center space-x-2 gap-2">
-            <AIToggle />
-          </div> */}
         </div>
       ) : null}
     </div>

app/articles/layout.tsx

@@ -1,3 +1,3 @@
-export default async function ArticlesLayout(props: { children: React.ReactNode }) {
+export default async function ArticlesLayout(props: Readonly<{ children: React.ReactNode }>) {
   return <>{props.children}</>;
 }

app/chat/[slug]/chat-completion.tsx

@@ -0,0 +1,62 @@
+'use client';
+
+import { useEffect } from 'react';
+
+import { Label } from '@/components/ui/label';
+import { MessageBubble } from '@/components/ui/message';
+import { chat_message } from '@prisma/client';
+import { useCompletion } from '@/hooks/use-completion';
+import { CardSmall } from '@/components/article-metadata/card-small';
+import { ChatThreadWithMessages } from '@/stores/chat-history';
+import { ChatDemo } from './chat-demo';
+
+type Props = {
+  slug: string;
+  thread: ChatThreadWithMessages;
+};
+
+export const ChatCompletion: React.FC<Props> = ({ slug, thread }) => {
+  const [trigger, { isLoading }] = useCompletion(slug);
+
+  useEffect(() => {
+    if (thread.chat_message.length === 1 && thread.chat_message[0].role === 'USER') {
+      trigger({ prompt: thread.chat_message[0].content });
+    }
+  }, []);
+
+  const getMessageContent = (message: chat_message, isLast: boolean, isLoading: boolean) => {
+    if (message.role === 'ASSISTANT') {
+      return isLoading && isLast && !message.content ? '...' : message.content;
+    }
+    return message.content;
+  };
+
+  return (
+    <>
+      <div className="flex flex-col p-4 justify-between min-h-screen">
+        <div className="flex flex-col gap-4">
+          {thread.chat_message.map((message, index) => (
+            <MessageBubble variant={message.role === 'ASSISTANT' ? 'answer' : 'question'} key={message.id}>
+              {getMessageContent(message, index === thread.chat_message.length - 1, isLoading)}
+            </MessageBubble>
+          ))}
+        </div>
+
+        <ChatDemo />
+      </div>
+      <div className="flex flex-col gap-4 p-4 max-w-[360px]">
+        <Label>Suggested Articles</Label>
+        {thread.suggested_articles?.map(({ article_metadata: article }) => (
+          <CardSmall
+            key={article.id}
+            id={article.id}
+            slug={article.slug}
+            title={article.title}
+            abstract={article.abstract}
+            published={new Date(article.published).toDateString()}
+          />
+        ))}
+      </div>
+    </>
+  );
+};