05Sessions and MongoDB

Sessions and MongoDB

This is the point where things actually get interesting. We will be creating sessions to store the context of the conversation and store messages and sessions in MongoDB, so that we can retrieve the conversation history and continue the conversation, just the way it is done in the official ChatGPT app.

Setting up MongoDB

Go to MongoDB and sign up for a free account. Then, create a new cluster and get the connection string from the Connect tab. Once you have the connection string, store the MONGODB_URI in the root directory of your project. Now, we will create the library which will help our application connect to mongodb.

lib/mongodb.ts


import { MongoClient } from "mongodb";

const uri = process.env.MONGODB_URI!;
const options = {};

/**
 * Global is used here to maintain a cached connection across hot reloads in development
*/
declare global {
  var _mongoClientPromise: Promise<MongoClient> | undefined;
}

let client: MongoClient;
let clientPromise: Promise<MongoClient>;

if (!process.env.MONGODB_URI) {
  throw new Error("Please add MONGODB_URI to .env");
}

if (process.env.NODE_ENV === "development") {
  // Reuse client in dev
  if (!(global)._mongoClientPromise) {
    client = new MongoClient(uri, options);
    (global)._mongoClientPromise = client.connect();
  }
  clientPromise = (global)._mongoClientPromise;
} else {
  client = new MongoClient(uri, options);
  clientPromise = client.connect();
}

export default clientPromise;

Once we have the connection library completed, we will create the chatMessages.ts file which will help us store and retrieve messages from the database.

lib/chatMessages.ts

import clientPromise from "@/lib/mongodb";
import { ObjectId } from "mongodb";

export type Citation = {
    filePath: string;
    chunkIndex: number;
};

export type ChatMessage = {
    _id: ObjectId;
    sessionId: ObjectId;
    role: "user" | "assistant";
    content: string;
    citations?: Citation[];
    createdAt: Date;
};

/**
 * Save a message
 */
export async function saveMessage({
    sessionId,
    role,
    content,
    citations,
}: {
    sessionId: string;
    role: "user" | "assistant";
    content: string;
    citations?: Citation[];
}) {
    try {
        const client = await clientPromise;
        const db = client.db("secondbrain");

        console.log(`[DB] Saving ${role} message for session ${sessionId}...`);

        const res = await db.collection("messages").insertOne({
            sessionId: new ObjectId(sessionId),
            role,
            content,
            citations,
            createdAt: new Date(),
        });

        console.log(`[DB] Successfully saved message with ID: ${res.insertedId}`);

        // bump session updatedAt
        const updateRes = await db.collection("sessions").updateOne(
            { _id: new ObjectId(sessionId) },
            { $set: { updatedAt: new Date() } }
        );

        if (updateRes.matchedCount === 0) {
            console.warn(`[DB] Warning: No session found with ID ${sessionId} to update.`);
        } else {
            console.log(`[DB] Updated session timestamp for ${sessionId}`);
        }
    } catch (error) {
        console.error(`[DB] Error saving message:`, error);
        throw error; // Re-throw to ensure the caller knows it failed
    }
}

/**
 * Load all messages for a session
 */
export async function getMessagesBySession(sessionId: string) {
    const client = await clientPromise;
    const db = client.db("secondbrain");

    return db
        .collection<ChatMessage>("messages")
        .find({ sessionId: new ObjectId(sessionId) })
        .sort({ createdAt: 1 })
        .toArray();
}

Create a Session Id Page

This page will be the page containing all the chat conversations for each of the session

app/chat/[id]/page.tsx

"use client";
import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";
import { useState, useRef, useEffect } from "react";
import Markdown from "react-markdown";
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
import TypingIndicator from "@/app/components/TypingIndicator";
import { Button } from "@/app/components/ui/Button";
import { Input } from "@/app/components/ui/Input";
import { Send, User, Bot, Brain } from "lucide-react";
import { motion, AnimatePresence } from "framer-motion";
import { use } from "react";
import { useSearchParams } from "next/navigation";

export default function ChatSessionPage({ params }: { params: Promise<{ id: string }> }) {
    const { id } = use(params);
    const [input, setInput] = useState("");
    const { messages, setMessages, sendMessage, status } = useChat({
        transport: new DefaultChatTransport({
            api: `/api/chat?sessionId=${id}`
        }),
        body: {
            sessionId: id
        }
    } as any);

    const [historyLoaded, setHistoryLoaded] = useState(false);

    useEffect(() => {
        async function fetchHistory() {
            try {
                const res = await fetch(`/api/messages/${id}`);
                const data = await res.json();
                if (Array.isArray(data) && data.length > 0) {
                    const mappedMessages = data.map((m: any) => ({
                        id: m._id,
                        role: m.role as any,
                        parts: [{ type: "text" as const, text: m.content }],
                    }));
                    setMessages(mappedMessages);
                }
                setHistoryLoaded(true);
            } catch (error) {
                console.error("Failed to fetch history:", error);
                setHistoryLoaded(true);
            }
        }
        fetchHistory();
    }, [id, setMessages]);

    const bottomRef = useRef<HTMLDivElement>(null);

    useEffect(() => {
        bottomRef.current?.scrollIntoView({ behavior: "smooth" });
    }, [messages]);

    const onSubmit = async (e: React.FormEvent) => {
        e.preventDefault();
        if (!input.trim() || status === "submitted" || status === "streaming") return;
        await sendMessage({ text: input, metadata: { sessionId: id } });
        setInput("");
    };

    const isLoading = status === "submitted" || status === "streaming";

    const searchParams = useSearchParams();
    const firstQuery = searchParams.get("q");
    const sentRef = useRef(false);

    useEffect(() => {
        if (historyLoaded && firstQuery && !sentRef.current && messages.length === 0) {
            sentRef.current = true;
            sendMessage({
                text: firstQuery,
                metadata: { sessionId: id }
            })
        }
    }, [historyLoaded, firstQuery, sendMessage, messages.length, id]);

    return (
        <div className="flex flex-col h-screen bg-[#0a0a0a]">
            {/* Header */}
            <header className="h-16 border-b border-white/5 flex items-center justify-between px-8 bg-black/50 backdrop-blur-xl shrink-0">
                <div className="flex items-center gap-3">
                    <div className="w-8 h-8 rounded-lg bg-zinc-800 flex items-center justify-center border border-white/10">
                        <Brain className="w-4 h-4 text-blue-400" />
                    </div>
                    <div>
                        <h1 className="text-sm font-semibold text-white">Knowledge Session</h1>
                        <p className="text-[10px] text-zinc-500 uppercase tracking-widest">ID: {id.slice(-6)}</p>
                    </div>
                </div>
            </header>

            {/* Messages */}
            <main className="flex-1 overflow-y-auto px-4 md:px-0 py-8">
                <div className="max-w-3xl mx-auto space-y-8 pb-12">
                    {messages.length === 0 && (
                        <motion.div
                            initial={{ opacity: 0, y: 20 }}
                            animate={{ opacity: 1, y: 0 }}
                            className="flex flex-col items-center justify-center py-20 text-center"
                        >
                            <div className="w-16 h-16 rounded-3xl bg-blue-600/10 border border-blue-500/20 flex items-center justify-center mb-6">
                                <Brain className="w-8 h-8 text-blue-500" />
                            </div>
                            <h2 className="text-2xl font-bold text-white mb-2">Deep Knowledge Retrieval</h2>
                            <p className="text-zinc-500 max-w-sm">
                                This session is ready. Ask anything about your indexed documents and notes.
                            </p>
                        </motion.div>
                    )}

                    <AnimatePresence initial={false}>
                        {messages.map((m) => (
                            <motion.div
                                key={m.id}
                                initial={{ opacity: 0, y: 10 }}
                                animate={{ opacity: 1, y: 0 }}
                                className={`flex gap-4 ${m.role === "user" ? "flex-row-reverse" : "flex-row"}`}
                            >
                                <div className={`w-8 h-8 rounded-full flex items-center justify-center shrink-0 border ${m.role === "user"
                                    ? "bg-blue-600/10 border-blue-500/20 text-blue-400"
                                    : "bg-zinc-800 border-white/10 text-zinc-400"
                                    }`}>
                                    {m.role === "user" ? <User className="w-4 h-4" /> : <Bot className="w-4 h-4" />}
                                </div>

                                <div className={`flex flex-col max-w-[85%] ${m.role === "user" ? "items-end" : "items-start"}`}>
                                    <div className={`px-5 py-3 rounded-2xl ${m.role === "user"
                                        ? "bg-blue-600 text-white"
                                        : "bg-zinc-900 border border-white/5 text-zinc-200"
                                        }`}>
                                        {m.parts.map((p, i) => (
                                            p.type === "text" ? (
                                                <div key={i} className="prose prose-sm prose-invert max-w-none">
                                                    <Markdown
                                                        components={{
                                                            code({ node, inline, className, children, ...props }: any) {
                                                                const match = /language-(w+)/.exec(className || '');
                                                                return !inline && match ? (
                                                                    <SyntaxHighlighter
                                                                        style={vscDarkPlus}
                                                                        language={match[1]}
                                                                        PreTag="div"
                                                                        className="rounded-xl my-4 text-xs"
                                                                        {...props}
                                                                    >
                                                                        {String(children).replace(/
$/, '')}
                                                                    </SyntaxHighlighter>
                                                                ) : (
                                                                    <code className="bg-white/10 px-1.5 py-0.5 rounded text-sm" {...props}>
                                                                        {children}
                                                                    </code>
                                                                );
                                                            }
                                                        }}
                                                    >
                                                        {p.text}
                                                    </Markdown>
                                                </div>
                                            ) : null
                                        ))}
                                    </div>
                                    <div className="mt-1 px-2 text-[10px] text-zinc-600">
                                        {m.role === "user" ? "You" : "Second Brain"}
                                    </div>
                                </div>
                            </motion.div>
                        ))}
                    </AnimatePresence>

                    {isLoading && (
                        <div className="flex gap-4">
                            <div className="w-8 h-8 rounded-full bg-zinc-800 border border-white/10 flex items-center justify-center text-zinc-400">
                                <Bot className="w-4 h-4" />
                            </div>
                            <div className="bg-zinc-900/50 border border-white/5 rounded-2xl px-4 py-2">
                                <TypingIndicator />
                            </div>
                        </div>
                    )}
                    <div ref={bottomRef} className="h-4" />
                </div>
            </main>

            {/* Input Area */}
            <div className="p-6 bg-linear-to-t from-black via-black to-transparent">
                <form onSubmit={onSubmit} className="max-w-3xl mx-auto relative group">
                    <Input
                        value={input}
                        onChange={(e) => setInput(e.target.value)}
                        placeholder="Search your brain..."
                        disabled={isLoading}
                        className="pr-14 h-14 rounded-2xl bg-zinc-900/50 border-white/10 focus:border-blue-500/50 focus:ring-blue-500/20 shadow-2xl transition-all"
                    />
                    <div className="absolute right-2 top-2">
                        <Button
                            type="submit"
                            size="icon"
                            disabled={isLoading || !input.trim()}
                            className="h-10 w-10 rounded-xl"
                        >
                            <Send className="w-4 h-4" />
                        </Button>
                    </div>
                </form>
                <p className="text-center mt-3 text-[10px] text-zinc-600 tracking-wider">
                    AI generated responses may be inaccurate. Check citations for verification.
                </p>
            </div>
        </div>
    );
}

Create the post api route for sessionId for messages

This API is to get the messages for a specific session.

app/api/messages/[sessionId]/route.ts

import { NextRequest } from "next/server";
import { getMessagesBySession } from "@/lib/chatMessages";

export async function GET(_: NextRequest, props: { params: Promise<{ sessionId: string }> }) {
    const params = await props.params;
    const messages = await getMessagesBySession(params.sessionId);
    return Response.json(messages);
}

Store the session to MongoDB

This API will help us store the current session to MongoDB

app/api/sessions/route.ts

import clientPromise from "@/lib/mongodb";
import { NextRequest } from "next/server";

export async function POST(req: NextRequest) {
    const { firstMessage } = await req.json();

    if (!firstMessage?.trim()) {
        return new Response("Missing message", {
            status: 400
        });
    }

    const client = await clientPromise;
    const db = client.db("secondbrain");

    const now = new Date();

    const res = await db.collection("sessions").insertOne({
        title: firstMessage.slice(0, 40),
        createdAt: now,
        updatedAt: now,
    });

    return Response.json({
        sessionId: res.insertedId.toString(),
    });
}

Key implementation details for the Chat UI:

Vercel AI SDK (useChat): This hook handles the heavy lifting of chat state management. It tracks messages, handles the input field, and automatically processes the streaming response from the server.
DefaultChatTransport: We use this to explicitly point the useChat hook to our /api/chat endpoint. This provides a clean interface for network communication between the frontend and the AI route.
Rich Text Rendering: The UI uses react-markdown and react-syntax-highlighter to render the AI's response. This ensures that code snippets, bold text, and lists are displayed with professional formatting.
Auto-Scrolling: A combination of useRef and useEffect ensures that as new message chunks arrive, the window automatically scrolls to keep the latest content in view.

Update the Chat API to store the session

app/api/chat/route.ts

import { NextRequest } from "next/server";
import { google } from "@ai-sdk/google";
import { streamText, convertToModelMessages } from "ai";

import { getOrCreateCollection } from "@/lib/chromaClient";
import { saveMessage } from "@/lib/chatMessages";

export const runtime = "nodejs";
export const maxDuration = 40;

const COLLECTION_NAME = "secondbrain";

export async function POST(req: NextRequest) {
  const body = await req.json();

  const messages = body.messages;
  const lastMessage = messages?.[messages.length - 1];

  // ✅ sessionId extraction (Root Body > Message Metadata > URL Param)
  const sessionId: string | undefined =
    body.sessionId ||
    lastMessage?.metadata?.sessionId ||
    req.nextUrl.searchParams.get("sessionId") ||
    undefined;

  const userMessage =
    lastMessage?.parts?.find((p: any) => p.type === "text")?.text;

  if (!sessionId || !messages || !userMessage?.trim()) {
    console.error("INVALID PAYLOAD:", JSON.stringify(body, null, 2));
    return new Response(JSON.stringify({
      error: "Missing sessionId or message",
      received: { sessionId: !!sessionId, messages: !!messages, userMessage: !!userMessage }
    }), { status: 400 });
  }

  const query = userMessage.trim();

  console.log(`[API] Attempting to save user message for session ${sessionId}`);
  await saveMessage({
    sessionId,
    role: "user",
    content: query,
  });

  /* ---------------- RAG ---------------- */
  const collection = await getOrCreateCollection(COLLECTION_NAME);

  const ragResults = await collection.query({
    queryTexts: [query],
    nResults: 5,
    include: ["documents", "metadatas"],
  });

  const docs = ragResults.documents?.[0] ?? [];
  const metas = ragResults.metadatas?.[0] ?? [];

  const citations = metas.map((m: any) => ({
    filePath: m.filePath ?? m.path ?? "unknown",
    chunkIndex: m.chunkIndex ?? 0,
  }));

  const context = docs
    .map(
      (doc, i) =>
        `Source ${i + 1} (${citations[i].filePath}, chunk ${citations[i].chunkIndex}):
${doc}``
    )
    .join("

");

  const systemPrompt = `
You are Adi's personal Second Brain assistant.
Use ONLY the provided Context.
If missing, say:
"I don't have that in my Second Brain yet."
`.trim();

  const modelMessages = convertToModelMessages([
    {
      role: "system",
      parts: [{ type: "text", text: systemPrompt }],
    },
    {
      role: "user",
      parts: [
        {
          type: "text",
          text: `${query}

---
Context:
${context}`,
        },
      ],
    },
  ]);

  const result = streamText({
    model: google("gemini-2.5-flash"),
    messages: modelMessages,

    onFinish: async (event) => {
      await saveMessage({
        sessionId,
        role: "assistant",
        content: event.text,
        citations,
      });
    },
  });

  return result.toUIMessageStreamResponse();
}

Key implementation details for the Chat UI:

Vercel AI SDK (useChat): This hook handles the heavy lifting of chat state management. It tracks messages, handles the input field, and automatically processes the streaming response from the server.
DefaultChatTransport: We use this to explicitly point the useChat hook to our /api/chat endpoint. This provides a clean interface for network communication between the frontend and the AI route.
Rich Text Rendering: The UI uses react-markdown and react-syntax-highlighter to render the AI's response. This ensures that code snippets, bold text, and lists are displayed with professional formatting.
Auto-Scrolling: A combination of useRef and useEffect ensures that as new message chunks arrive, the window automatically scrolls to keep the latest content in view.

Add Sidebar to Display the Sessions

Now that we have the data getting saved to database as well, let us create a sidebar component to display the sessions on the left side of the screen.

app/components/Sidebar.tsx

import { getChatSessions } from "@/lib/chatSessions";
import Link from "next/link";
import { Button } from "./ui/Button";
import { Plus, MessageSquare, Brain} from "lucide-react";

export default async function ChatSidebar() {
    const sessions = await getChatSessions();

    return (
        <aside className="w-72 h-screen border-r border-white/5 bg-black flex flex-col">
            {/* Header */}
            <div className="p-6">
                <Link href="/" className="flex items-center gap-2 mb-8 group">
                    <div className="w-8 h-8 rounded-lg bg-blue-600 flex items-center justify-center group-hover:rotate-12 transition-transform">
                        <Brain className="w-5 h-5 text-white" />
                    </div>
                    <span className="font-bold text-lg tracking-tight">Second Brain</span>
                </Link>

                <Link href="/chat">
                    <Button variant="primary" className="w-full justify-start gap-2 h-11 rounded-xl shadow-blue-500/10">
                        <Plus className="w-4 h-4" />
                        New Chat
                    </Button>
                </Link>
            </div>

            {/* Sessions List */}
            <div className="flex-1 overflow-y-auto px-3 space-y-1">
                <div className="px-3 mb-2">
                    <span className="text-[10px] font-bold uppercase tracking-wider text-zinc-500">Recent Chats</span>
                </div>

                {sessions.length === 0 ? (
                    <div className="px-3 py-4 text-xs text-zinc-500 italic">
                        No recent sessions
                    </div>
                ) : (
                    sessions.map((s: any) => (
                        <Link
                            key={s._id.toString()}
                            href={`/chat/${s._id.toString()}`}
                            className="group flex items-center gap-3 px-3 py-2.5 rounded-xl text-zinc-400 hover:text-white hover:bg-white/5 transition-all"
                        >
                            <MessageSquare className="w-4 h-4 shrink-0 opacity-50 group-hover:opacity-100" />
                            <div className="flex-1 min-w-0">
                                <div className="text-sm font-medium truncate">{s.title || "Untitled Chat"}</div>
                                <div className="text-[10px] opacity-40 mt-0.5">
                                    {new Date(s.updatedAt).toLocaleDateString([], { month: 'short', day: 'numeric' })}
                                </div>
                            </div>
                        </Link>
                    ))
                )}
            </div>
        </aside>
    );
}

Creating chatSessions Library

lib/chatSessions.ts

import { ObjectId } from "mongodb"
import clientPromise from "./mongodb";

export type ChatSession = {
    _id: ObjectId,
    title: string;
    createdAt: Date;
    updatedAt: Date;
};

export async function getChatSessions() {
    const client = await clientPromise;
    const db = client.db("secondbrain");

    return db.collection<ChatSession>("sessions")
        .find({})
        .sort({ updatedAt: -1 })
        .toArray();
}

export async function createChatSession() {
    const client = await clientPromise;
    const db = client.db("secondbrain");

    const now = new Date();

    const res = await db.collection("sessions").insertOne({
        title: "New Chat",
        createdAt: now,
        updatedAt: now,
    });

    return {
        _id: res.insertedId,
        title: "New Chat",
    };
}

export async function getChatSessionById(id: string) {
    const client = await clientPromise;
    const db = client.db("secondbrain");

    return db.collection("sessions").findOne({
        _id: new ObjectId(id),
    });
}

Next Steps

In the next section, we’ll:

Add Input and Output Guardrails

Add input and output guardrails to the chat API

If you want to know more about this, do checkout our video guide:

Streaming Ui

Input And Output Guardrails