AI Streaming Response

1"use client";
2 
3import { useEffect, useMemo, useRef, useState } from "react";
4import { cn } from "@/lib/utils";
5import AIMessage from "./ai-message";
6 
7const STREAMING_CONFIG = {
8  TOKEN_DELAY: 30,
9  PAUSE_AFTER_PERIOD: 200,
10  PAUSE_AFTER_COMMA: 100,
11  PAUSE_AFTER_PARAGRAPH: 300,
12} as const;
13 
14interface AIStreamingResponseProps {
15  content: string;
16  onComplete?: () => void;
17  autoStart?: boolean;
18  className?: string;
19}
20 
21function useTokens(content: string): string[] {
22  const tokens = useMemo(() => {
23    const tokenRegex = /(\S+|\s+)/g;
24    return content.match(tokenRegex) || [];
25  }, [content]);
26 
27  return tokens;
28}
29 
30function useStreaming(
31  tokens: string[],
32  autoStart: boolean,
33  onComplete?: () => void
34) {
35  const [displayedTokens, setDisplayedTokens] = useState<string[]>([]);
36  const [isPaused, setIsPaused] = useState(!autoStart);
37  const [isComplete, setIsComplete] = useState(false);
38  const timeoutRef = useRef<NodeJS.Timeout | null>(null);
39  const currentIndexRef = useRef(0);
40 
41  useEffect(() => {
42    if (tokens.length === 0) return;
43 
44    if (timeoutRef.current) {
45      clearTimeout(timeoutRef.current);
46      timeoutRef.current = null;
47    }
48 
49    const timeoutId = setTimeout(() => {
50      setDisplayedTokens([]);
51      setIsComplete(false);
52      setIsPaused(!autoStart);
53    }, 0);
54    currentIndexRef.current = 0;
55 
56    return () => clearTimeout(timeoutId);
57  }, [tokens, autoStart]);
58 
59  useEffect(() => {
60    if (tokens.length === 0) return;
61    if (isPaused || isComplete || currentIndexRef.current >= tokens.length) {
62      return;
63    }
64 
65    const streamNext = () => {
66      if (currentIndexRef.current >= tokens.length) {
67        setIsComplete(true);
68        if (onComplete) {
69          onComplete();
70        }
71        return;
72      }
73 
74      const token = tokens[currentIndexRef.current];
75      setDisplayedTokens((prev) => [...prev, token]);
76      currentIndexRef.current += 1;
77 
78      const trimmedToken = token.trim();
79 
80      let delay = STREAMING_CONFIG.TOKEN_DELAY;
81 
82      if (
83        trimmedToken.endsWith(".") ||
84        trimmedToken.endsWith("!") ||
85        trimmedToken.endsWith("?")
86      ) {
87        delay += STREAMING_CONFIG.PAUSE_AFTER_PERIOD;
88      } else if (
89        trimmedToken.endsWith(",") ||
90        trimmedToken.endsWith(";") ||
91        trimmedToken.endsWith(":")
92      ) {
93        delay += STREAMING_CONFIG.PAUSE_AFTER_COMMA;
94      } else if (token.includes("\n\n")) {
95        delay += STREAMING_CONFIG.PAUSE_AFTER_PARAGRAPH;
96      }
97 
98// … truncated