react-streaming.tsx

/**
 * React + streaming fetch + yalt (parseStream) + KaTeX
 *
 * Uses parseStream to render each math/text event as it arrives,
 * without re-parsing the accumulated buffer. This is the right
 * pattern when you want true O(n) streaming - for instance, if
 * messages are long or you need progressive math rendering.
 *
 * Expects the server to return a plain text stream (not SSE or the
 * AI SDK data protocol). With the Vercel AI SDK, that's:
 *
 *   // app/api/chat/route.ts
 *   import { streamText } from 'ai';
 *   import { openai } from '@ai-sdk/openai';
 *
 *   export async function POST(req: Request) {
 *     const { prompt } = await req.json();
 *     const result = streamText({
 *       model: openai('gpt-4o'),
 *       prompt,
 *     });
 *     return result.toTextStreamResponse();
 *   }
 *
 * Prerequisites:
 *   npm install @benjamin_r/yalt katex
 */
import { useCallback, useRef, useState } from 'react';
import katex from 'katex';
// --- yalt: import the streaming parser and renderer helper ---
import { parseStream, toRenderInput, type YaltEvent } from 'yalt';
import 'katex/dist/katex.min.css';

// ── Read a text/plain streaming response as string chunks ───────────

async function* readTextStream(response: Response): AsyncIterable<string> {
  const reader = response.body!
    .pipeThrough(new TextDecoderStream())
    .getReader();
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      yield value;
    }
  } finally {
    reader.releaseLock();
  }
}

// ── Streaming hook ──────────────────────────────────────────────────

function useMathStream(endpoint: string) {
  const [events, setEvents] = useState<YaltEvent[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const controllerRef = useRef<AbortController | null>(null);

  const send = useCallback(
    async (message: string) => {
      controllerRef.current?.abort();
      const controller = new AbortController();
      controllerRef.current = controller;

      setEvents([]);
      setIsStreaming(true);

      const response = await fetch(endpoint, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ message }),
        signal: controller.signal,
      });

      // --- yalt: pipe the text stream through parseStream for O(n) tokenising ---
      const accumulated: YaltEvent[] = [];
      for await (const event of parseStream(readTextStream(response))) {
        accumulated.push(event);
        setEvents([...accumulated]);
      }

      setIsStreaming(false);
    },
    [endpoint],
  );

  return { events, isStreaming, send };
}

// ── Render events ───────────────────────────────────────────────────

function EventList({ events }: { events: YaltEvent[] }) {
  return (
    <div>
      {events.map((event, i) => {
        if (event.type === 'text') return <span key={i}>{event.value}</span>;
        if (event.type !== 'math') return null;
        // --- yalt: toRenderInput gives KaTeX-ready { tex, displayMode } ---
        const { tex, displayMode } = toRenderInput(event);
        const html = katex.renderToString(tex, {
          displayMode,
          throwOnError: false,
        });
        const Tag = displayMode ? 'div' : 'span';
        return <Tag key={i} dangerouslySetInnerHTML={{ __html: html }} />;
      })}
    </div>
  );
}

// ── Chat UI ─────────────────────────────────────────────────────────

export default function StreamingChat() {
  const { events, isStreaming, send } = useMathStream('/api/chat');
  const [input, setInput] = useState('');

  const handleSubmit = (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim()) return;
    send(input.trim());
    setInput('');
  };

  return (
    <div style={{ maxWidth: 640, margin: '0 auto', padding: 16 }}>
      <EventList events={events} />
      {isStreaming && <span style={{ color: '#999' }}>...</span>}

      <form onSubmit={handleSubmit}>
        <input
          value={input}
          onChange={(e) => setInput(e.target.value)}
          placeholder="Ask a math question..."
          disabled={isStreaming}
          style={{ width: '100%', padding: 8 }}
        />
      </form>
    </div>
  );
}