Implement streaming for AI completions

Update system prompt with requirements and context formatting.
Add eventStream async generator for client-side SSE parsing.
Add ctx param to fetch handler for waitUntil caching.

Author: mantou132

packages/webapp/cf/index.ts

@@ -68,7 +68,7 @@ const resInit = {
 };
 
 export default {
-  async fetch(req, env, _ctx): Promise<Response> {
+  async fetch(req, env, ctx): Promise<Response> {
     if (req.method === 'OPTIONS') {
       return new Response(null, {
         status: 204,
@@ -129,26 +129,43 @@ export default {
       }
       case '/completions': {
         const cacheRes = await env.KV.get(req.url);
-        if (cacheRes) return new Response(cacheRes, resInit);
-
+        if (cacheRes) {
+          return new Response(cacheRes, {
+            headers: { ...resInit.headers, 'Content-Type': 'text/event-stream' },
+          });
+        }
         const q = params.get('q') || '';
         const [values] = await embedding(env.AI, [q]);
         const res = await env.GAMES_SEARCH.query(values, { returnMetadata: true });
         const messages = [
           {
             role: 'system',
-            content: `You are an application assistant and need to answer user questions in the target language(${req.headers.get('Accept-Language') ?? 'en'}) according to the following requirements::
-                      ${res.matches.map((e) => (e.metadata as any).text).join('\n\n\n')}`,
+            content: `You are an application assistant.
+Response user input based on the context and your existing knowledge.
+
+Requirements:
+1. Language: respond in (${req.headers.get('Accept-Language') ?? 'en'}).
+2. Maximum 1000 words.
+
+Context:
+${res.matches.map((e) => (e.metadata as any).text).join('\n\n---\n\n')}
+`,
           },
           {
             role: 'user',
             content: params.get('q') || '',
           },
         ];
-        const result = await env.AI.run('@cf/nvidia/nemotron-3-120b-a12b' as any, { messages });
-        const body = { content: result?.choices?.at?.(0)?.message?.content };
-        await env.KV.put(req.url, JSON.stringify(body), { expirationTtl: 60 * 60 });
-        return Response.json(body, resInit);
+        const stream = await env.AI.run('@cf/nvidia/nemotron-3-120b-a12b' as any, { messages, stream: true });
+        const [stream1, stream2] = stream.tee();
+
+        ctx.waitUntil(
+          new Response(stream1).text().then((text) => env.KV.put(req.url, text, { expirationTtl: 60 * 60 })),
+        );
+
+        return new Response(stream2, {
+          headers: { ...resInit.headers, 'Content-Type': 'text/event-stream' },
+        });
       }
       default: {
         return new Response('Hello World!');

packages/webapp/src/services/index.ts

@@ -150,3 +150,29 @@ export function subscribe<Result, InputVar = Record<string, any>>(
     },
   };
 }
+
+export async function* eventStream(input: string | URL | Request, init?: RequestInit) {
+  const res = await fetch(input, init);
+  if (!res.body) throw new Error('invalid method');
+
+  const reader = res.body.getReader();
+  const decoder = new TextDecoder();
+
+  let buffer = '';
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) return;
+
+    buffer += decoder.decode(value, { stream: true });
+    const lines = buffer.split('\n');
+    buffer = lines.pop() || '';
+
+    for (const line of lines) {
+      if (!line.startsWith('data: ')) continue;
+      const data = line.slice(6);
+      if (data === '[DONE]') return;
+      yield JSON.parse(data);
+    }
+  }
+}