ai-rag-chat

What this skill does

# AI RAG Chat

RAG-powered chat that retrieves relevant document chunks via vector search, injects them as context into the system prompt, and streams answers with source citations. Extends the existing `ai-chat` skill's route and UI.

## Prerequisites

- Next.js app with `src/` directory and App Router
- `ai-chat` skill installed (streaming chat at `/api/ai/chat`, UI components)
- `ai-rag-vectors` skill installed (`searchChunks()` at `@/lib/rag/search`)
- shadcn/ui initialized

## Installation

No additional packages needed — uses existing `ai` and `@ai-sdk/react` from `ai-chat`.

## What Gets Created

```
src/
├── app/
│   └── api/
│       └── rag/
│           └── chat/
│               └── route.ts               # POST — RAG streaming chat
└── components/
    └── rag/
        ├── rag-chat.tsx                    # RAG chat UI with useChat
        └── citation-badge.tsx             # Clickable citation badge component
```

## Setup Steps

### Step 1: Create `src/app/api/rag/chat/route.ts`

This is a dedicated RAG chat endpoint. It performs vector search, builds a context-augmented system prompt, and streams the response.

```typescript
import {
  streamText,
  convertToModelMessages,
  type UIMessage,
  type JSONValue,
} from "ai";
import { getModel } from "@/lib/ai";
import { withAuth } from "@/lib/auth-guard";
import { searchChunks } from "@/lib/rag/search";
import { db } from "@/lib/db";
import { chatSession, chatMessage } from "@/lib/db/schema/chat";
import { eq, and } from "drizzle-orm";

type RagChatBody = {
  messages: UIMessage[];
  sessionId?: string;
  documentIds?: string[];
};

type Citation = {
  documentId: string;
  documentTitle: string;
  pageNumber: number;
  chunkText: string;
  similarity: number;
};

export const POST = withAuth(async (request, { user }) => {
  const { messages, sessionId, documentIds }: RagChatBody = await request.json();

  const userId = user.id;

  // --- Resolve or create session ---
  let activeSessionId = sessionId;

  if (activeSessionId) {
    const existing = await db
      .select({ id: chatSession.id })
      .from(chatSession)
      .where(
        and(eq(chatSession.id, activeSessionId), eq(chatSession.userId, userId))
      )
      .limit(1);

    if (existing.length === 0) {
      return new Response(JSON.stringify({ error: "Session not found" }), {
        status: 404,
        headers: { "Content-Type": "application/json" },
      });
    }
  } else {
    const firstUserMessage = messages.find((m) => m.role === "user");
    const title =
      firstUserMessage?.parts
        .filter(
          (p): p is Extract<typeof p, { type: "text" }> => p.type === "text"
        )
        .map((p) => p.text)
        .join(" ")
        .slice(0, 100) || "New RAG Chat";

    const [created] = await db
      .insert(chatSession)
      .values({ userId, title })
      .returning({ id: chatSession.id });
    activeSessionId = created.id;
  }

  // --- Extract latest user query for RAG retrieval ---
  const lastUserMessage = messages.at(-1);
  const userQuery = lastUserMessage?.parts
    .filter(
      (p): p is Extract<typeof p, { type: "text" }> => p.type === "text"
    )
    .map((p) => p.text)
    .join(" ");

  // --- Perform vector search ---
  let citations: Citation[] = [];
  let contextBlock = "";

  if (userQuery) {
    const searchResults = await searchChunks({
      query: userQuery,
      documentIds,
      userId,
      limit: 8,
    });

    citations = searchResults.map((r) => ({
      documentId: r.documentId,
      documentTitle: r.documentTitle,
      pageNumber: r.pageNumber,
      chunkText: r.textContent,
      similarity: r.similarity,
    }));

    if (citations.length > 0) {
      const contextParts = citations.map(
        (c, i) =>
          `[Source ${i + 1}] "${c.documentTitle}" — Page ${c.pageNumber}:\n${c.chunkText}`
      );
      contextBlock = contextParts.join("\n\n---\n\n");
    }
  }

  // --- Build system prompt with RAG context ---
  const systemParts: string[] = [
    "You are a helpful assistant that answers questions based on the provided document context.",
    "When answering, cite your sources using [Source N] notation matching the context below.",
    "If the context doesn't contain relevant information, say so honestly.",
    "Be concise, accurate, and helpful.",
  ];

  if (contextBlock) {
    systemParts.push(
      "## Document Context\n\n" + contextBlock
    );
  } else {
    systemParts.push(
      "No document context was found for this query. Answer based on your general knowledge and let the user know you couldn't find relevant document sections."
    );
  }

  // --- Persist user message ---
  if (lastUserMessage && lastUserMessage.role === "user") {
    await db.insert(chatMessage).values({
      sessionId: activeSessionId,
      role: "user",
      parts: lastUserMessage.parts,
    });
  }

  // --- Convert and stream ---
  const modelMessages = await convertToModelMessages(messages);

  const result = streamText({
    model: getModel(),
    system: systemParts.join("\n\n"),
    messages: modelMessages,
    async onFinish({ text }) {
      await db.insert(chatMessage).values({
        sessionId: activeSessionId,
        role: "assistant",
        parts: [{ type: "text", text }],
      });

      await db
        .update(chatSession)
        .set({ updatedAt: new Date() })
        .where(eq(chatSession.id, activeSessionId));
    },
  });

  const response = result.toUIMessageStreamResponse();

  // Append metadata headers
  response.headers.set("X-Session-Id", activeSessionId);
  response.headers.set(
    "X-Rag-Citations",
    encodeURIComponent(JSON.stringify(citations))
  );

  return response;
});
```

### Step 2: Create `src/components/rag/citation-badge.tsx`

```tsx
"use client";

import { memo } from "react";

type CitationBadgeProps = {
  sourceIndex: number;
  documentTitle: string;
  pageNumber: number;
  onClick?: () => void;
};

export const CitationBadge = memo(function CitationBadge({
  sourceIndex,
  documentTitle,
  pageNumber,
  onClick,
}: CitationBadgeProps) {
  return (
    <button
      type="button"
      onClick={onClick}
      className="inline-flex items-center gap-1 rounded-md bg-blue-50 px-2 py-0.5 text-xs font-medium text-blue-700 ring-1 ring-inset ring-blue-600/20 transition-colors hover:bg-blue-100 dark:bg-blue-950 dark:text-blue-300 dark:ring-blue-400/30 dark:hover:bg-blue-900"
      title={`${documentTitle} — Page ${pageNumber}`}
    >
      <span>[{sourceIndex}]</span>
      <span className="max-w-[120px] truncate">{documentTitle}</span>
      <span className="text-blue-500">p.{pageNumber}</span>
    </button>
  );
});
```

### Step 3: Create `src/components/rag/rag-chat.tsx`

```tsx
"use client";

import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport, type UIMessage } from "ai";
import { useId, useRef, useEffect, useCallback, useState, useMemo, memo } from "react";
import { ChatContainer } from "@/components/ai/chat-container";
import { PromptInput } from "@/components/ai/prompt-input";
import { Loader } from "@/components/ai/loader";
import { Message } from "@/components/ai/message";
import { Markdown } from "@/components/ai/markdown";
import { CitationBadge } from "./citation-badge";

type Citation = {
  documentId: string;
  documentTitle: string;
  pageNumber: number;
  chunkText: string;
  similarity: number;
};

type RagChatProps = {
  documentIds?: string[];
  sessionId: string | null;
  onSessionCreated?: (sessionId: string) => void;
  onCitationClick?: (citation: Citation) => void;
};

export function RagChat({
  documentIds,
  sessionId,
  onSessionCreated,
  onCitationClick,
}: RagChatProps) {
  const messageListId = useId();
  const scrollRef = useRef<HTMLDivElement>(null);
  const [input, setInput] = useState("");
  const [citations, setCitations] = useState<Citation[]>([]);

  // Refs to hold latest callbacks without causing transport recreation
  const onSessionCreatedRef = useRef(onSession