AI SDK + AI Gateway

import { convertToModelMessages, createUIMessageStreamResponse, streamText, toUIMessageStream, validateUIMessages } from "ai";
import { labModel } from "@/lib/ai-model";
import { checkLimit, demoLimit, retryIn } from "@/lib/redis";

const MAX_MESSAGES = 20;
const MAX_INPUT_CHARS = 500;

// The chat endpoint useChat talks to. Even with the simulated model, it's
// built as if every call cost money: rate-limited, with capped input and output.
export async function POST(request: Request) {
  const limit = await checkLimit(demoLimit);
  if (limit && !limit.success) {
    return Response.json({ error: `Slow down: try again in ${retryIn(limit.reset)} s.` }, { status: 429 });
  }

  let messages;
  try {
    messages = await validateUIMessages({ messages: ((await request.json()) as { messages?: unknown }).messages });
  } catch {
    return Response.json({ error: "Invalid messages" }, { status: 400 });
  }
  const recent = messages.slice(-MAX_MESSAGES);
  const lastText = recent.at(-1)?.parts.map((p) => (p.type === "text" ? p.text : "")).join("") ?? "";
  if (lastText.length > MAX_INPUT_CHARS) {
    return Response.json({ error: `Messages are limited to ${MAX_INPUT_CHARS} characters.` }, { status: 413 });
  }

  const result = streamText({
    model: labModel(),
    system: "You are the Vercel Feature Lab's assistant. Answer briefly and plainly.",
    messages: await convertToModelMessages(recent),
    maxOutputTokens: 300,
  });

  return createUIMessageStreamResponse({ stream: toUIMessageStream({ stream: result.stream }) });
}

"use client";

import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";
import { useState } from "react";

// useChat keeps the conversation, posts it to the route and renders the reply
// as it streams in. Messages are lists of parts (text, tool calls, files…);
// this chat only uses text.
export function Chat() {
  const [input, setInput] = useState("");
  const { messages, sendMessage, status, error, stop } = useChat({
    // A fixed id: otherwise useChat makes a random one while the page is
    // prerendered, which Cache Components refuses.
    id: "lab-chat",
    transport: new DefaultChatTransport({ api: "/api/ai/chat" }),
  });
  const busy = status === "submitted" || status === "streaming";

  return (
    <>
      <div className="chat-log" aria-live="polite">
        {messages.length === 0 && <p className="muted small">Ask anything. The reply streams in word by word.</p>}
        {messages.map((m) => (
          <p key={m.id}>
            <strong>{m.role === "user" ? "You" : "Lab"}:</strong>{" "}
            {m.parts.map((part, i) => (part.type === "text" ? <span key={i}>{part.text}</span> : null))}
          </p>
        ))}
      </div>
      {error && <p className="error" role="alert">{error.message.includes("429") || error.message.includes("Slow down") ? "Slow down a little and try again." : "Something went wrong. Try again."}</p>}
      <form
        className="form-row"
        onSubmit={(e) => {
          e.preventDefault();
          const text = input.trim();
          if (!text || busy) return;
          sendMessage({ text });
          setInput("");
        }}
      >
        <label className="grow">Message<input value={input} onChange={(e) => setInput(e.target.value)} maxLength={500} placeholder="What does streamText do?" /></label>
        {busy ? <button type="button" className="ghost" onClick={stop}>Stop</button> : <button type="submit" disabled={!input.trim()}>Send</button>}
      </form>
    </>
  );
}

import { type LanguageModel, simulateReadableStream } from "ai";
import { MockLanguageModelV4 } from "ai/test";

// The model behind /ai. By default it's SIMULATED: the AI SDK's mock model,
// which streams a canned reply without calling any provider, so nothing can
// ever be billed. Set AI_MODEL to a gateway model id (for example
// "anthropic/claude-haiku-4.5") to use a real model through Vercel AI Gateway;
// the rest of the code doesn't change. (AI Gateway needs a card on file, even
// for its free tier.)

export const isSimulated = () => !process.env.AI_MODEL;

export function labModel(): LanguageModel {
  return process.env.AI_MODEL || simulatedModel();
}

type PromptMessage = { role: string; content: unknown };

// The text of the newest user message in a model prompt.
export function lastUserText(prompt: PromptMessage[]) {
  const last = [...prompt].reverse().find((m) => m.role === "user");
  if (!last || !Array.isArray(last.content)) return "";
  return last.content
    .map((part: { type?: string; text?: string }) => (part.type === "text" ? part.text ?? "" : ""))
    .join(" ")
    .trim();
}

export function simulatedAnswer(question: string) {
  const asked = question.length > 120 ? `${question.slice(0, 117)}…` : question;
  return [
    `This is a simulated reply: no AI model is connected, so nothing is billed.`,
    `You asked: "${asked}".`,
    `Here's what just happened. useChat sent your message to /api/ai/chat. The route checked the rate limit, converted the chat`,
    `into model messages and called streamText. The model, a mock from the AI SDK's test tools, emitted this text word by word,`,
    `and toUIMessageStream turned it into the stream your browser is reading now. With AI_MODEL set, the same route would`,
    `stream a real answer through Vercel AI Gateway instead.`,
  ].join(" ");
}

function simulatedModel() {
  return new MockLanguageModelV4({
    provider: "lab-simulated",
    modelId: "simulated",
    doStream: async ({ prompt }) => {
      const words = simulatedAnswer(lastUserText(prompt)).split(/(?<=\s)/);
      return {
        stream: simulateReadableStream({
          initialDelayInMs: 300,
          chunkDelayInMs: 35,
          chunks: [
            { type: "text-start", id: "text-1" },
            ...words.map((delta) => ({ type: "text-delta" as const, id: "text-1", delta })),
            { type: "text-end", id: "text-1" },
            {
              type: "finish",
              finishReason: { unified: "stop", raw: undefined },
              logprobs: undefined,
              usage: {
                inputTokens: { total: 0, noCache: 0, cacheRead: undefined, cacheWrite: undefined },
                outputTokens: { total: words.length, text: words.length, reasoning: undefined },
              },
            },
          ],
        }),
      };
    },
  });
}