diff --git a/apps/cf-ai-backend/src/routes.ts b/apps/cf-ai-backend/src/routes.ts index 4a2d2827..841f107a 100644 --- a/apps/cf-ai-backend/src/routes.ts +++ b/apps/cf-ai-backend/src/routes.ts @@ -2,6 +2,7 @@ import { CloudflareVectorizeStore } from '@langchain/cloudflare'; import * as apiAdd from './routes/add'; import * as apiQuery from "./routes/query" import * as apiAsk from "./routes/ask" +import * as apiChat from "./routes/chat" import { OpenAIEmbeddings } from './OpenAIEmbedder'; import { GenerativeModel } from '@google/generative-ai'; import { Request } from '@cloudflare/workers-types'; @@ -17,6 +18,8 @@ routeMap.set('/query', apiQuery); routeMap.set('/ask', apiAsk); +routeMap.set('/chat', apiChat); + // Add more route mappings as needed // routeMap.set('/api/otherRoute', { ... }); diff --git a/apps/cf-ai-backend/src/routes/chat.ts b/apps/cf-ai-backend/src/routes/chat.ts new file mode 100644 index 00000000..1082998f --- /dev/null +++ b/apps/cf-ai-backend/src/routes/chat.ts @@ -0,0 +1,91 @@ +import { Content, GenerativeModel } from "@google/generative-ai"; +import { OpenAIEmbeddings } from "../OpenAIEmbedder"; +import { CloudflareVectorizeStore } from "@langchain/cloudflare"; +import { Request } from "@cloudflare/workers-types"; + +export async function POST(request: Request, _: CloudflareVectorizeStore, embeddings: OpenAIEmbeddings, model: GenerativeModel, env?: Env) { + const queryparams = new URL(request.url).searchParams; + const query = queryparams.get("q"); + const topK = parseInt(queryparams.get("topK") ?? "5"); + const user = queryparams.get("user") + const space = queryparams.get("space") + + const sourcesOnly = (queryparams.get("sourcesOnly") ?? "false") + + if (!user) { + return new Response(JSON.stringify({ message: "Invalid User" }), { status: 400 }); + } + + if (!query) { + return new Response(JSON.stringify({ message: "Invalid Query" }), { status: 400 }); + } + + const filter: VectorizeVectorMetadataFilter = { + user + } + + if (space) { + filter.space + } + + const queryAsVector = await embeddings.embedQuery(query); + + const resp = await env!.VECTORIZE_INDEX.query(queryAsVector, { + topK, + filter + }); + + if (resp.count === 0) { + return new Response(JSON.stringify({ message: "No Results Found" }), { status: 404 }); + } + + const highScoreIds = resp.matches.filter(({ score }) => score > 0.3).map(({ id }) => id) + + if (sourcesOnly === "true") { + return new Response(JSON.stringify({ ids: highScoreIds }), { status: 200 }); + } + + const vec = await env!.VECTORIZE_INDEX.getByIds(highScoreIds) + + const preparedContext = vec.slice(0, 3).map(({ metadata }) => `Website title: ${metadata!.title}\nDescription: ${metadata!.description}\nURL: ${metadata!.url}\nContent: ${metadata!.text}`).join("\n\n"); + + const body = await request.json() as { + chatHistory?: Content[] + }; + + const defaultHistory = [ + { + role: "user", + parts: [{ text: `You are an agent that summarizes a page based on the query. don't say 'based on the context'. I expect you to be like a 'Second Brain'. you will be provided with the context (old saved posts) and questions. Answer accordingly. Answer in markdown format` }], + }, + { + role: "model", + parts: [{ text: "Ok, I am a personal assistant, and will act as a second brain to help with user's queries." }], + }, + ] as Content[]; + + const chat = model.startChat({ + history: [...defaultHistory, ...(body.chatHistory ?? [])], + }); + + const prompt = `Context:\n${preparedContext}\n\nQuestion: ${query}\nAnswer:`; + + const output = await chat.sendMessageStream(prompt); + + const response = new Response( + new ReadableStream({ + async start(controller) { + const converter = new TextEncoder(); + for await (const chunk of output.stream) { + const chunkText = await chunk.text(); + const encodedChunk = converter.encode("data: " + JSON.stringify({ "response": chunkText }) + "\n\n"); + controller.enqueue(encodedChunk); + } + const doneChunk = converter.encode("data: [DONE]"); + controller.enqueue(doneChunk); + controller.close(); + } + }) + ); + return response; +} diff --git a/apps/cf-ai-backend/src/routes/query.ts b/apps/cf-ai-backend/src/routes/query.ts index e02d0150..be237d7d 100644 --- a/apps/cf-ai-backend/src/routes/query.ts +++ b/apps/cf-ai-backend/src/routes/query.ts @@ -36,7 +36,7 @@ export async function GET(request: Request, _: CloudflareVectorizeStore, embeddi }); if (resp.count === 0) { - return new Response(JSON.stringify({ message: "No Results Found" }), { status: 400 }); + return new Response(JSON.stringify({ message: "No Results Found" }), { status: 404 }); } const highScoreIds = resp.matches.filter(({ score }) => score > 0.3).map(({ id }) => id) 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 ( +
+
+ + {user === 'ai' ? ( + AI + ) : user?.image ? ( + <> + + + {user?.name?.split(' ').map((n) => n[0])}{' '} + + + ) : ( + + )} + +
{message}
+
+
+
+ ); +} + +export { ChatMessage }; diff --git a/apps/web/src/components/Main.tsx b/apps/web/src/components/Main.tsx index b77106a9..3c338526 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 { @@ -26,7 +29,34 @@ export default function Main({ sidebarOpen }: { sidebarOpen: boolean }) { const [searchResults, setSearchResults] = useState([]); const [isAiLoading, setIsAiLoading] = useState(false); + const { data: session } = useSession(); + + // Variable to keep track of the chat history in this session + const [chatHistory, setChatHistory] = useState([]); + + // 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 + >({}); + + // 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(null); @@ -155,6 +185,15 @@ export default function Main({ sidebarOpen }: { sidebarOpen: boolean }) { hide ? "" : "main-hidden", )} > +
+ {chatHistory.map((chat, index) => ( + + ))} +

Ask your Second brain

diff --git a/apps/web/types/memory.tsx b/apps/web/types/memory.tsx index f184615a..e71e92c9 100644 --- a/apps/web/types/memory.tsx +++ b/apps/web/types/memory.tsx @@ -53,3 +53,8 @@ export type CollectedSpaces = { title: string; content: StoredContent[]; }; + +export type ChatHistory = { + role: 'user' | 'model'; + parts: [{ text: string }]; +};