Comment simuler une réponse streamée avec le SDK Vercel AI

Chez Premier Octet, nous développons le SaaS Postula, une application permettant d'importer ses flux RSS et de générer facilement des posts LinkedIn à partir d'articles.

L'un des principaux composants est un widget permettant de générer des posts LinkedIn à partir d'un flux RSS. Nous voulions montrer ce widget (visible ici) dès la homepage, pour des raisons de démonstration : un exemple concret vaut mieux que des mots.

Le problème : celui-ci fait des appels API vers OpenAI, appels qui nous sont facturés.

Nous avons donc eu l'idée de simuler la réponse de manière progressive, pour donner l'illusion d'un vrai appel à un modèle LLM tel que GPT sans le coût associé.

L'existant

Dans notre composant, nous utilisons le hook useCompletion de Vercel AI SDK pour générer le contenu des posts LinkedIn :

import { useCompletion } from 'ai'

const { complete } = useCompletion({
  api: '/api/completion/post',
})

Du côté de l'API, nous avons donc un endpoint /api/completion/post qui fait un appel à OpenAI :

import { openai } from '@ai-sdk/openai'
import { streamText } from 'ai'

export async function POST() {
  const result = await streamText({
    model: openai('gpt-4o'),
    prompt: `Write an Haiku about React.js`,
  })

  return result.toDataStreamResponse()
}

Afin de voir le format de la réponse, nous allons faire un appel à cet endpoint puis inspecter la réponse via l'onglet "Network" de notre navigateur :

HTTP/1.1 200 OK
vary: RSC, Next-Router-State-Tree, Next-Router-Prefetch
content-type: text/plain; charset=utf-8
x-vercel-ai-data-stream: v1
Date: Tue, 22 Oct 2024 15:41:15 GMT
Connection: keep-alive
Keep-Alive: timeout=5
Transfer-Encoding: chunked

0:"Java"
0:"Script"
0:" dances"
0:","
0:"  \n"
0:"Components"
0:" weave"
0:" through"
0:" the"
0:" DOM"
0:","
0:"  \n"
0:"React"
0:" breath"
0:"es"
0:" new"
0:" life"
0:"."
e:{"finishReason":"stop","usage":{"promptTokens":14,"completionTokens":18},"isContinued":false}
d:{"finishReason":"stop","usage":{"promptTokens":14,"completionTokens":18}}

Les morceaux de texte sont envoyés au fur et à mesure, ligne par ligne, avec "0:" en préfixe puis à la fin, des informations indiquant que la réponse est terminée.

Nous allons donc simuler cette réponse.

Commençons par créer la fonction qui va découper le texte en chunks.

Création des chunks

Maintenant que nous connaissons la structure de la réponse, nous pouvons créer une fonction qui va découper un texte en chunks :

export const chunkText = (inputText: string): string[] => {
  // Découpe le texte en un tableau de mots et espaces
  const chunks = inputText.match(/\S+|\s/g) || []

  return chunks.map((chunk) => {
    // Gère les saut de lignes
    if (chunk === '\n') {
      return '0:"\\n"\n'
      // Gère les espaces
    } else if (chunk.trim() === '') {
      return chunk
        .split('')
        .map((char) => (char === '\n' ? '0:"\\n"\n' : `0:"${char}"\n`))
        .join('')
      // Gère les mots
    } else {
      return `0:"${chunk}"\n`
    }
  })
}

Cette fonction découpe le texte en morceaux (mots, espaces) et les prépare à être envoyés sous forme de réponse incrémentielle. Le préfixe 0: est ajouté pour imiter le format de réponse d'API que nous avons observé précédemment.

Création du endpoint pour streamer la réponse

La prochaine étape consiste à créer une route d'API qui va streamer notre texte chunk par chunk.

export async function POST() {
  const text = `Lorem ipsum dolor sit amet...` // Texte à streamer
  const predefinedChunks = chunkText(text) // Découpe du texte en chunks

  const encoder = new TextEncoder()

  const stream = new ReadableStream({
    async start(controller) {
      // Pause de 1s pour simuler l'attente
      await new Promise((resolve) => setTimeout(resolve, 1000))

      let index = 0

      function push() {
        if (index < predefinedChunks.length) {
          const chunk = predefinedChunks[index]
          controller.enqueue(encoder.encode(chunk)) // Envoi du chunk encodé
          index++

          // Intervalle de 10ms entre chaque envoi
          // pour simuler un stream avec un peu de délai
          setTimeout(push, 10)
        } else {
          controller.close() // Ferme le stream à la fin
        }
      }

      // Démarre le stream
      push()
    },
  })

  const headers = new Headers({
    'Content-Type': 'text/plain; charset=utf-8',
    'Transfer-Encoding': 'chunked',
  })

  return new Response(stream, { headers })
}

Cette route envoie chaque morceau de texte progressivement avec un délai de 10ms (avec un délai initial de 1s) donnant ainsi l’impression d’une réponse en direct.

Test du endpoint avec notre composant

Nous pouvons maintenant utiliser notre endpoint de démo pour simuler une réponse d'un modèle LLM :

useCompletion({
  api: isDemo ? '/api/completion/demo/post' : '/api/completion/post',
})

Cette approche permet de basculer facilement entre la démo et le vrai endpoint, en ajoutant une props qui active le mode démo.

Bien sûr, n'oubliez pas de sécuriser votre endpoint /api/completion/post en vérifiant si l'utilisateur est authentifié et dispose d'assez de crédits.

La réponse est bien streamée, et le composant fonctionne comme prévu !

À noter qu'il existe depuis la version 3.4 de Vercel AI SDK une fonction streamText qui permet de mock une réponse de stream :

import { streamText } from 'ai'
import { convertArrayToReadableStream, MockLanguageModelV1 } from 'ai/test'

const result = await streamText({
  model: new MockLanguageModelV1({
    doStream: async () => ({
      stream: convertArrayToReadableStream([
        { type: 'text-delta', textDelta: 'Hello' },
        { type: 'text-delta', textDelta: ', ' },
        { type: 'text-delta', textDelta: `world!` },
      ]),
    }),
  }),
})

Mais celle-ci est destinée à des tests unitaires ou E2E, et ne simule donc pas de délai entre chaque chunk. Ce qui rend le résultat moins réaliste.

👋