konnosDashboard
ai-stream.ts133 lines · main
1import { env } from '../env.js';
2import { log } from '../lib/logger.js';
3
4import { AiNotConfiguredError } from './ai-schema-gen.js';
5
6/**
7 * Streaming variant of the ai-* services. The dashboard surfaces can
8 * now render tokens as they arrive instead of waiting up to 60s for
9 * the full response. The Ollama backend supports streaming natively
10 * via `stream: true` — the response body is a newline-delimited JSON
11 * stream where each line carries a `{response: "<chunk>", done:
12 * false}` event, with a final `{done: true, …}` marker.
13 *
14 * We re-emit those events as Server-Sent Events (SSE) so the browser
15 * can consume them via EventSource. SSE is supported natively on every
16 * modern browser, no extra library on the dashboard.
17 *
18 * Auth is the same as the non-streaming services — X-API-Key when
19 * BRIVEN_OLLAMA_API_KEY is set, no header otherwise.
20 */
21
22export interface StreamInput {
23 system: string;
24 prompt: string;
25 model: string;
26 temperature: number;
27 timeoutMs?: number;
28}
29
30export async function streamAiResponse(input: StreamInput): Promise<Response> {
31 if (!env.BRIVEN_OLLAMA_URL) {
32 throw new AiNotConfiguredError();
33 }
34
35 const headers: Record<string, string> = { 'content-type': 'application/json' };
36 if (env.BRIVEN_OLLAMA_API_KEY) {
37 headers['x-api-key'] = env.BRIVEN_OLLAMA_API_KEY;
38 }
39
40 const upstreamRes = await fetch(`${env.BRIVEN_OLLAMA_URL.replace(/\/$/, '')}/api/generate`, {
41 method: 'POST',
42 headers,
43 body: JSON.stringify({
44 model: input.model,
45 system: input.system,
46 prompt: input.prompt,
47 options: { temperature: input.temperature },
48 stream: true,
49 }),
50 signal: AbortSignal.timeout(input.timeoutMs ?? 120_000),
51 });
52
53 if (!upstreamRes.ok || !upstreamRes.body) {
54 const body = await upstreamRes.text().catch(() => '');
55 log.warn('ai_stream_upstream_error', {
56 status: upstreamRes.status,
57 bodyPreview: body.slice(0, 240),
58 });
59 throw new Error(`Ollama returned ${upstreamRes.status}`);
60 }
61
62 // Translate the ollama ndjson stream into SSE. Each ollama line:
63 // {"response":"foo","done":false}
64 // becomes an SSE event:
65 // event: token
66 // data: foo
67 //
68 // The final ollama line is `{done:true, ...}` — we emit a final
69 // `event: done` so the client can dispose the EventSource cleanly.
70 const ollamaStream = upstreamRes.body;
71 const sseStream = new ReadableStream({
72 async start(controller) {
73 const enc = new TextEncoder();
74 const reader = ollamaStream.getReader();
75 const dec = new TextDecoder();
76 let buffer = '';
77 try {
78 while (true) {
79 const { done, value } = await reader.read();
80 if (done) break;
81 buffer += dec.decode(value, { stream: true });
82 let idx;
83 while ((idx = buffer.indexOf('\n')) >= 0) {
84 const line = buffer.slice(0, idx).trim();
85 buffer = buffer.slice(idx + 1);
86 if (!line) continue;
87 let parsed: { response?: string; done?: boolean };
88 try {
89 parsed = JSON.parse(line) as { response?: string; done?: boolean };
90 } catch {
91 continue;
92 }
93 if (parsed.response) {
94 controller.enqueue(
95 enc.encode(`event: token\ndata: ${jsonEscape(parsed.response)}\n\n`),
96 );
97 }
98 if (parsed.done) {
99 controller.enqueue(enc.encode(`event: done\ndata: {}\n\n`));
100 }
101 }
102 }
103 } catch (err) {
104 controller.enqueue(
105 enc.encode(
106 `event: error\ndata: ${jsonEscape(err instanceof Error ? err.message : 'stream error')}\n\n`,
107 ),
108 );
109 } finally {
110 controller.close();
111 }
112 },
113 });
114
115 return new Response(sseStream, {
116 status: 200,
117 headers: {
118 'content-type': 'text/event-stream; charset=utf-8',
119 'cache-control': 'no-cache, no-transform',
120 // Disable buffering on proxies that respect this header (nginx).
121 'x-accel-buffering': 'no',
122 },
123 });
124}
125
126/**
127 * SSE data lines can't contain literal newlines without escaping. JSON-
128 * stringify handles that (and quotes everything), then strip the outer
129 * quotes since SSE doesn't need them.
130 */
131function jsonEscape(s: string): string {
132 return JSON.stringify(s).slice(1, -1);
133}