Streaming UI

import { getOrCreateCollection } from "@/lib/chromaClient";
import { google } from "@ai-sdk/google";
import { convertToModelMessages, generateText, streamText, type UIMessage } from 'ai';
import { NextRequest } from "next/server";

export const runtime = "nodejs";
export const maxDuration = 40; //Allow streaming for upto 40s

function extractUserMessageText(msg: any): string {
    if (Array.isArray(msg.parts)) {
        return msg.parts
            .map((p: any) => (p.type === "text" ? p.text : ""))
            .join("\n")
            .trim();
    }
    return "";
}


const COLLECTION_NAME = "secondbrain";

export async function POST(req: NextRequest) {
    const { messages } = (await req.json() as { messages: UIMessage[] });

    if (!messages || !messages?.length) {
        return new Response("Missing messages", { status: 400 });
    };

    // Get last user message as query
    const lastUserIndex = [...messages].reverse().findIndex((m) => m?.role === "user");

    if (lastUserIndex === -1) {
        return new Response("No user message found", { status: 400 });
    }

    const realIndex = messages?.length - 1 - lastUserIndex;
    const lastUserMessage = messages[realIndex];
    const query = extractUserMessageText(lastUserMessage);

    if (!query?.trim()) {
        return new Response("Empty user query", { status: 400 });
    }

    // RAG: Retrieve Top-k chunks from Chroma
    const collection = await getOrCreateCollection(COLLECTION_NAME);

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

    const docs = (ragResults?.documents?.[0] ?? []) as string[];
    const metas = (ragResults?.metadatas?.[0] ?? []) as Record<string, any>[];

    // Build a context block
    const context = docs
        ?.map((doc, i) => {
            const meta = metas[i] || {};
            const source = meta?.filePath ?? meta?.path ?? "unknown";
            const chunkIndex = meta?.chunkIndex ?? i;
            return `Source ${i + 1} (file: ${source}, chunk: ${chunkIndex}):\n${doc}`;
        }).join("\n\n");

    const systemPrompt = `
You are Adi's personal Second Brain assistant.

Use ONLY the information provided in the "Context" section below when answering.
If the answer is not clearly contained in the context, say:
"I don't have that in my Second Brain yet."

When you answer:
- Be concise but clear.
- Prefer bullet points where it helps.
- If relevant, mention which source(s) you used.
`.trim();

    const augmentedLastUser: UIMessage = {
        id: lastUserMessage?.id,
        role: "user",
        parts: [
            {
                type: "text",
                text: `${query} 
                --- 
                Context from Adi's Second Brain:
                ${context || "[no matching context found]"}
                
                Now answer the user's question strictly based on the context above.
                `
            }
        ],
    }

    // Build model messages: system + previous history + augmented last user
    const uiMessagesWithContext: UIMessage[] = [
        {
            id: "system-1",
            role: "system",
            parts: [{
                type: "text",
                text: systemPrompt
            }],
        },
        ...messages?.slice(0, realIndex),
        augmentedLastUser,
    ];

    // Stream response
    const result = streamText({
        model: google("gemini-1.5-flash"),
        messages: convertToModelMessages(uiMessagesWithContext),
    });

    // Return streaming response compatible with useChat
    return result?.toUIMessageStreamResponse();
}

"use client";
import { UIMessage, 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 "./components/TypingIndicator";

export default function ChatPage() {
  const [input, setInput] = useState("");

  const { messages, sendMessage, status } = useChat({
    transport: new DefaultChatTransport({
      api: '/api/chat'
    }),
  });

  const bottomRef = useRef<HTMLDivElement>(null);

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

  const onSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim()) return;
    await sendMessage({ text: input });
    setInput("");
  };

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

  return (
    <div className="min-h-screen flex flex-col bg-gradient-to-br from-gray-900 via-black to-gray-900 text-white">
      <header className="px-6 py-4 border-b border-white/10 backdrop-blur-sm bg-black/20 sticky top-0 z-10">
        <div className="max-w-4xl mx-auto flex items-center justify-between">
          <h1 className="text-xl font-semibold flex items-center gap-2">
            <span className="text-2xl">🧠</span>
            Second Brain
          </h1>
          <div className="text-sm text-white/50">
            {messages.length > 0 && `${messages.length} message${messages.length > 1 ? 's' : ''}`}
          </div>
        </div>
      </header>

      <main className="flex-1 overflow-y-auto p-6">
        <div className="max-w-4xl mx-auto space-y-6">
          {messages.length === 0 && (
            <div className="flex flex-col items-center justify-center min-h-[60vh] text-center animate-fadeIn">
              <div className="text-6xl mb-6">💡</div>
              <h2 className="text-3xl font-bold mb-3 bg-gradient-to-r from-blue-400 to-purple-400 bg-clip-text text-transparent">
                Welcome to Your Second Brain
              </h2>
              <p className="text-white/60 max-w-md mb-8">
                Ask me anything about your knowledge base. I'll search through your notes and provide answers with source citations.
              </p>
              <div className="grid grid-cols-1 md:grid-cols-2 gap-3 w-full max-w-2xl">
                {[
                  "What are my recent notes about?",
                  "Summarize my thoughts on...",
                  "Find information about...",
                  "What do I know about..."
                ].map((suggestion, i) => (
                  <button
                    key={i}
                    onClick={() => setInput(suggestion)}
                    className="px-4 py-3 rounded-xl bg-white/5 hover:bg-white/10 border border-white/10 hover:border-white/20 transition-all text-sm text-left"
                  >
                    {suggestion}
                  </button>
                ))}
              </div>
            </div>
          )}

          {messages.map((m, idx) => (
            <div
              key={m.id}
              className={`flex ${m.role === "user" ? "justify-end" : "justify-start"} animate-fadeIn`}
              style={{ animationDelay: `${idx * 0.05}s` }}
            >
              <div className={`max-w-[85%] ${m.role === "user" ? "ml-auto" : "mr-auto"}`}>
                <div className={`px-5 py-3 rounded-2xl ${m.role === "user"
                  ? "bg-gradient-to-r from-blue-600 to-blue-500 text-white shadow-lg shadow-blue-500/20"
                  : "bg-white/5 backdrop-blur-sm text-white border border-white/10"
                  }`}>
                  {m.parts.map((p, i) => (
                    p.type === "text" ? (
                      <div key={i} className="prose prose-invert prose-sm 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-lg my-2"
                                  {...props}
                                >
                                  {String(children).replace(/
$/, '')}
                                </SyntaxHighlighter>
                              ) : (
                                <code className="bg-white/10 px-1.5 py-0.5 rounded text-sm" {...props}>
                                  {children}
                                </code>
                              );
                            },
                            p: ({ children }) => <p className="mb-2 last:mb-0">{children}</p>,
                            ul: ({ children }) => <ul className="list-disc list-inside mb-2 space-y-1">{children}</ul>,
                            ol: ({ children }) => <ol className="list-decimal list-inside mb-2 space-y-1">{children}</ol>,
                            li: ({ children }) => <li className="text-white/90">{children}</li>,
                            strong: ({ children }) => <strong className="font-semibold text-white">{children}</strong>,
                            a: ({ children, href }) => (
                              <a href={href} className="text-blue-400 hover:text-blue-300 underline" target="_blank" rel="noopener noreferrer">
                                {children}
                              </a>
                            ),
                          }}
                        >
                          {p.text}
                        </Markdown>
                      </div>
                    ) : null
                  ))}
                </div>
              </div>
            </div>
          ))}

          {isLoading && (
            <div className="flex justify-start animate-fadeIn">
              <div className="bg-white/5 backdrop-blur-sm rounded-2xl border border-white/10">
                <TypingIndicator />
              </div>
            </div>
          )}

          <div ref={bottomRef} />
        </div>
      </main>

      <form onSubmit={onSubmit} className="p-4 border-t border-white/10 bg-black/40 backdrop-blur-md sticky bottom-0">
        <div className="max-w-4xl mx-auto flex gap-3">
          <input
            value={input}
            onChange={(e) => setInput(e.target.value)}
            placeholder="Ask your Second Brain anything..."
            disabled={isLoading}
            className="flex-1 px-5 py-3 rounded-xl bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent placeholder:text-white/40 disabled:opacity-50 disabled:cursor-not-allowed transition-all"
          />
          <button
            type="submit"
            disabled={isLoading || !input.trim()}
            className="px-6 py-3 rounded-xl bg-gradient-to-r from-blue-600 to-blue-500 hover:from-blue-500 hover:to-blue-400 transition-all disabled:opacity-50 disabled:cursor-not-allowed font-medium shadow-lg shadow-blue-500/20 hover:shadow-blue-500/40"
          >
            {isLoading ? (
              <span className="flex items-center gap-2">
                <svg className="animate-spin h-4 w-4" viewBox="0 0 24 24">
                  <circle className="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" strokeWidth="4" fill="none" />
                  <path className="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z" />
                </svg>
                Thinking...
              </span>
            ) : (
              "Send"
            )}
          </button>
        </div>
      </form>
    </div>
  );
}