RagChat

"use client";

import { useEffect, useRef, useState } from "react";

export interface Source {
  content: string;
  source?: string;
  page?: number;
  score?: number;
}

interface Message {
  role: "user" | "assistant";
  content: string;
  sources?: Source[];
}

interface RagChatProps {
  /** Local route handler that proxies your pipeline. Default: "/api/chat". */
  route?: string;
  placeholder?: string;
  title?: string;
  className?: string;
}

function patchLast(messages: Message[], patch: Partial<Message>): Message[] {
  const copy = [...messages];
  const i = copy.length - 1;
  if (i >= 0 && copy[i].role === "assistant") copy[i] = { ...copy[i], ...patch };
  return copy;
}

/**
 * RagChat (Next.js) — same UI as the React variant, but it talks to a local
 * Route Handler (app/api/chat/route.ts) that proxies your pipeline. Your
 * RAG_ENDPOINT + RAG_API_KEY stay server-side, and the answer streams in over
 * SSE token-by-token.
 */
export function RagChat({ route = "/api/chat", placeholder = "Ask a question…", title = "Ask the docs", className = "" }: RagChatProps) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState("");
  const [loading, setLoading] = useState(false);
  const scroller = useRef<HTMLDivElement>(null);

  useEffect(() => {
    scroller.current?.scrollTo({ top: scroller.current.scrollHeight, behavior: "smooth" });
  }, [messages, loading]);

  async function send() {
    const q = input.trim();
    if (!q || loading) return;
    setInput("");
    setMessages((m) => [...m, { role: "user", content: q }, { role: "assistant", content: "" }]);
    setLoading(true);

    try {
      const res = await fetch(route, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ query: q }),
      });
      if (!res.body) throw new Error("no stream");

      const reader = res.body.getReader();
      const decoder = new TextDecoder();
      let buffer = "";
      let acc = "";

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });
        const parts = buffer.split("\n\n");
        buffer = parts.pop() ?? "";
        for (const part of parts) {
          const line = part.replace(/^data: ?/, "").trim();
          if (!line) continue;
          let evt: { token?: string; sources?: Source[]; done?: boolean };
          try {
            evt = JSON.parse(line);
          } catch {
            continue;
          }
          if (evt.token) {
            acc += evt.token;
            setMessages((m) => patchLast(m, { content: acc }));
          }
          if (evt.done) {
            setMessages((m) => patchLast(m, { sources: evt.sources ?? [] }));
          }
        }
      }
    } catch {
      setMessages((m) => patchLast(m, { content: "Something went wrong reaching the endpoint." }));
    } finally {
      setLoading(false);
    }
  }

  return (
    <div className={`flex h-[460px] flex-col overflow-hidden rounded-xl border border-zinc-200 bg-white dark:border-zinc-800 dark:bg-zinc-950 ${className}`}>
      <div className="border-b border-zinc-200 px-4 py-3 text-sm font-semibold text-zinc-700 dark:border-zinc-800 dark:text-zinc-200">
        {title}
      </div>

      <div ref={scroller} className="flex-1 space-y-4 overflow-y-auto p-4">
        {messages.length === 0 && (
          <p className="mt-12 text-center text-sm text-zinc-400">Ask anything about your documents.</p>
        )}
        {messages.map((m, i) => {
          const streaming = m.role === "assistant" && loading && i === messages.length - 1;
          return (
            <div key={i} className={`flex ${m.role === "user" ? "justify-end" : "justify-start"}`}>
              <div
                className={`max-w-[85%] rounded-2xl px-3.5 py-2 text-[13.5px] leading-relaxed ${
                  m.role === "user"
                    ? "bg-zinc-900 text-white dark:bg-white dark:text-zinc-900"
                    : "bg-zinc-100 text-zinc-800 dark:bg-zinc-800 dark:text-zinc-100"
                }`}
              >
                {m.content === "" && streaming ? (
                  <span className="flex gap-1 py-0.5">
                    <Dot /><Dot delay={0.15} /><Dot delay={0.3} />
                  </span>
                ) : (
                  <p className="whitespace-pre-wrap">
                    {m.content}
                    {streaming && <span className="ml-0.5 inline-block w-[2px] animate-pulse bg-current align-text-bottom" style={{ height: "1em" }} />}
                  </p>
                )}
                {m.sources && m.sources.length > 0 && (
                  <div className="mt-2 flex flex-wrap gap-1">
                    {m.sources.map((s, j) => (
                      <span key={j} title={s.content} className="rounded-md bg-white/70 px-1.5 py-0.5 text-[11px] text-zinc-500 dark:bg-black/30 dark:text-zinc-400">
                        [{j + 1}] {s.source || "source"}
                      </span>
                    ))}
                  </div>
                )}
              </div>
            </div>
          );
        })}
      </div>

      <div className="border-t border-zinc-200 p-3 dark:border-zinc-800">
        <div className="flex items-center gap-2">
          <input
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={(e) => e.key === "Enter" && send()}
            placeholder={placeholder}
            className="flex-1 rounded-lg border border-zinc-200 bg-white px-3 py-2 text-sm outline-none focus:border-zinc-400 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100"
          />
          <button
            onClick={send}
            disabled={loading || !input.trim()}
            className="rounded-lg bg-zinc-900 px-3.5 py-2 text-sm font-medium text-white transition disabled:opacity-40 dark:bg-white dark:text-zinc-900"
          >
            Send
          </button>
        </div>
      </div>
    </div>
  );
}

function Dot({ delay = 0 }: { delay?: number }) {
  return <span className="h-1.5 w-1.5 animate-bounce rounded-full bg-zinc-400" style={{ animationDelay: `${delay}s` }} />;
}

// app/api/chat/route.ts
//
// Proxies a Ragiment pipeline and streams the answer back to the browser over
// SSE. RAG_ENDPOINT and RAG_API_KEY stay on the server — they're never exposed
// to the client. Set them in .env.local.

export const dynamic = "force-dynamic";

export async function POST(req: Request) {
  const { query } = (await req.json()) as { query: string };

  const endpoint = process.env.RAG_ENDPOINT;
  if (!endpoint) {
    return new Response("RAG_ENDPOINT is not set", { status: 500 });
  }

  // Call the pipeline server-side (the key never reaches the browser).
  const upstream = await fetch(endpoint, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      ...(process.env.RAG_API_KEY ? { "X-API-Key": process.env.RAG_API_KEY } : {}),
    },
    body: JSON.stringify({ query }),
  });

  if (!upstream.ok) {
    return new Response(`Pipeline error: ${upstream.status}`, { status: 502 });
  }

  const data = await upstream.json();
  const answer: string = data.answer ?? "";
  const sources = data.sources ?? [];

  // Stream the answer token-by-token, then a final event with the sources.
  // If your pipeline exposes a native token stream, pipe `upstream.body`
  // through here instead of re-chunking the finished answer.
  const encoder = new TextEncoder();
  const stream = new ReadableStream<Uint8Array>({
    async start(controller) {
      for (const token of answer.split(/(\s+)/)) {
        controller.enqueue(encoder.encode(`data: ${JSON.stringify({ token })}\n\n`));
        await new Promise((r) => setTimeout(r, 18));
      }
      controller.enqueue(encoder.encode(`data: ${JSON.stringify({ sources, done: true })}\n\n`));
      controller.close();
    },
  });

  return new Response(stream, {
    headers: {
      "Content-Type": "text/event-stream",
      "Cache-Control": "no-cache, no-transform",
      Connection: "keep-alive",
    },
  });
}