setup for multi chat

3 files changed, 157 insertions, 7 deletions

diff --git a/apps/web/src/app/api/chat/route.ts b/apps/web/src/app/api/chat/route.ts
new file mode 100644
index 00000000..2cb03186
--- /dev/null
+++ b/apps/web/src/app/api/chat/route.ts
@@ -0,0 +1,62 @@
+import { db } from "@/server/db";
+import { eq } from "drizzle-orm";
+import { sessions, users } from "@/server/db/schema";
+import { type NextRequest, NextResponse } from "next/server";
+import { env } from "@/env";
+import { ChatHistory } from "../../../../types/memory";
+
+export const runtime = "edge";
+
+export async function POST(req: NextRequest) {
+    const token = req.cookies.get("next-auth.session-token")?.value ?? req.cookies.get("__Secure-authjs.session-token")?.value ?? req.cookies.get("authjs.session-token")?.value ?? req.headers.get("Authorization")?.replace("Bearer ", "");
+
+    const sessionData = await db.select().from(sessions).where(eq(sessions.sessionToken, token!))
+
+    if (!sessionData || sessionData.length === 0) {
+        return new Response(JSON.stringify({ message: "Invalid Key, session not found." }), { status: 404 });
+    }
+
+    const user = await db.select().from(users).where(eq(users.id, sessionData[0].userId)).limit(1)
+
+    if (!user || user.length === 0) {
+        return NextResponse.json({ message: "Invalid Key, session not found." }, { status: 404 });
+    }
+
+    const session = { session: sessionData[0], user: user[0] }
+
+    const query = new URL(req.url).searchParams.get("q");
+    const sourcesOnly = new URL(req.url).searchParams.get("sourcesOnly") ?? "false";
+
+    const chatHistory = await req.json() as {
+        chatHistory: ChatHistory[]
+    };
+
+
+    if (!query) {
+        return new Response(JSON.stringify({ message: "Invalid query" }), { status: 400 });
+    }
+
+    const resp = await fetch(`https://cf-ai-backend.dhravya.workers.dev/chat?q=${query}&user=${session.user.email ?? session.user.name}&sourcesOnly=${sourcesOnly}`, {
+        headers: {
+            "X-Custom-Auth-Key": env.BACKEND_SECURITY_KEY,
+        },
+        method: "POST",
+        body: JSON.stringify({
+            chatHistory
+        })
+    })
+
+    console.log(resp.status)
+
+    if (resp.status !== 200 || !resp.ok) {
+        const errorData = await resp.json();
+        console.log(errorData)
+        return new Response(JSON.stringify({ message: "Error in CF function", error: errorData }), { status: resp.status });
+    }
+
+    // Stream the response back to the client
+    const { readable, writable } = new TransformStream();
+    resp && resp.body!.pipeTo(writable);
+
+    return new Response(readable, { status: 200 });
+}
\ No newline at end of file
diff --git a/apps/web/src/components/ChatMessage.tsx b/apps/web/src/components/ChatMessage.tsx
new file mode 100644
index 00000000..a8199758
--- /dev/null
+++ b/apps/web/src/components/ChatMessage.tsx
@@ -0,0 +1,50 @@
+import React from 'react';
+import { Avatar, AvatarFallback, AvatarImage } from './ui/avatar';
+import { User } from 'next-auth';
+import { User2 } from 'lucide-react';
+import Image from 'next/image';
+
+function ChatMessage({
+  message,
+  user,
+}: {
+  message: string;
+  user: User | 'ai';
+}) {
+  return (
+    <div className="flex flex-col gap-4">
+      <div
+        className={`font-bold ${!(user === 'ai') && 'text-xl '} flex flex-col md:flex-row items-center gap-4`}
+      >
+        <Avatar>
+          {user === 'ai' ? (
+            <Image
+              src="/logo.png"
+              width={48}
+              height={48}
+              alt="AI"
+              className="rounded-md w-12 h-12"
+            />
+          ) : user?.image ? (
+            <>
+              <AvatarImage
+                className="h-6 w-6 rounded-lg"
+                src={user?.image}
+                alt="user pfp"
+              />
+              <AvatarFallback>
+                {user?.name?.split(' ').map((n) => n[0])}{' '}
+              </AvatarFallback>
+            </>
+          ) : (
+            <User2 strokeWidth={1.3} className="h-6 w-6" />
+          )}
+        </Avatar>
+        <div className="ml-4">{message}</div>
+      </div>
+      <div className="w-full h-0.5 bg-gray-700 my-2 md:my-0 md:mx-4 mt-8"></div>
+    </div>
+  );
+}
+
+export { ChatMessage };
diff --git a/apps/web/src/components/Main.tsx b/apps/web/src/components/Main.tsx
index 86679dcf..a9111494 100644
--- a/apps/web/src/components/Main.tsx
+++ b/apps/web/src/components/Main.tsx
@@ -1,5 +1,5 @@
 'use client';
-import { useEffect, useRef, useState } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
 import { FilterCombobox } from './Sidebar/FilterCombobox';
 import { Textarea2 } from './ui/textarea';
 import { ArrowRight } from 'lucide-react';
@@ -8,6 +8,9 @@ import useViewport from '@/hooks/useViewport';
 import { motion } from 'framer-motion';
 import { cn } from '@/lib/utils';
 import SearchResults from './SearchResults';
+import { ChatHistory } from '../../types/memory';
+import { ChatMessage } from './ChatMessage';
+import { useSession } from 'next-auth/react';
 
 function supportsDVH() {
   try {
@@ -24,7 +27,34 @@ export default function Main({ sidebarOpen }: { sidebarOpen: boolean }) {
   const [searchResults, setSearchResults] = useState<string[]>([]);
   const [isAiLoading, setIsAiLoading] = useState(false);
 
+  const { data: session } = useSession();
+
+  // Variable to keep track of the chat history in this session
+  const [chatHistory, setChatHistory] = useState<ChatHistory[]>([]);
+
+  // TEMPORARY solution: Basically this is to just keep track of the sources used for each chat message
+  // Not a great solution
+  const [chatTextSourceDict, setChatTextSourceDict] = useState<
+    Record<string, string>
+  >({});
+
+  // helper function to append a new msg
+  const appendToChatHistory = useCallback(
+    (role: 'user' | 'model', content: string) => {
+      setChatHistory((prev) => [
+        ...prev,
+        {
+          role,
+          parts: [{ text: content }],
+        },
+      ]);
+    },
+    [],
+  );
+
+  // This is the streamed AI response we get from the server.
   const [aiResponse, setAIResponse] = useState('');
+  
   const [toBeParsed, setToBeParsed] = useState('');
 
   const textArea = useRef<HTMLTextAreaElement>(null);
@@ -153,17 +183,28 @@ export default function Main({ sidebarOpen }: { sidebarOpen: boolean }) {
         hide ? '' : 'main-hidden',
       )}
     >
+      <div className="flex flex-col w-full">
+        {chatHistory.map((chat, index) => (
+          <ChatMessage
+            key={index}
+            message={chat.parts[0].text}
+            user={chat.role === 'model' ? 'ai' : session?.user!}
+          />
+        ))}
+      </div>
       <h1 className="text-rgray-11 mt-auto w-full text-center text-3xl md:mt-0">
         Ask your Second brain
       </h1>
-      <form onSubmit={async (e) => await getSearchResults(e)}>
+      <form
+        className="mt-auto h-max min-h-[3em] w-full resize-y flex-row items-start justify-center overflow-none py-5 md:mt-0 md:h-[20vh] md:resize-none md:flex-col md:items-center md:justify-center md:p-2 md:pb-2 md:pt-2"
+        onSubmit={async (e) => await getSearchResults(e)}
+      >
         <Textarea2
           ref={textArea}
-          className="mt-auto h-max max-h-[30em] min-h-[3em] resize-y flex-row items-start justify-center overflow-auto py-5 md:mt-0 md:h-[20vh] md:resize-none md:flex-col md:items-center md:justify-center md:p-2 md:pb-2 md:pt-2"
           textAreaProps={{
             placeholder: 'Ask your SuperMemory...',
             className:
-              'h-auto overflow-auto md:h-full md:resize-none text-lg py-0 px-2 md:py-0 md:p-5 resize-y text-rgray-11 w-full min-h-[1em]',
+              'h-auto overflow-auto md:h-full md:resize-none text-lg py-0 px-2 pt-2 md:py-0 md:p-5 resize-y text-rgray-11 w-full min-h-[1em]',
             value,
             autoFocus: true,
             onChange: (e) => setValue(e.target.value),
@@ -181,9 +222,6 @@ export default function Main({ sidebarOpen }: { sidebarOpen: boolean }) {
           </div>
         </Textarea2>
       </form>
-      {searchResults && (
-        <SearchResults aiResponse={aiResponse} sources={searchResults} />
-      )}
       {width <= 768 && <MemoryDrawer hide={hide} />}
     </motion.main>
   );