1
import { z } from "zod";
2
 
3
export const AiUsage = z.object({
4
  promptTokens: z.number().optional(),
5
  completionTokens: z.number().optional(),
6
  totalTokens: z.number().optional(),
7
});
8
export type AiUsage = z.infer<typeof AiUsage>;
9
 
10
export const AiErrorCode = z.enum([
11
  "ai_misconfigured",
12
  "ai_not_entitled",
13
  "ai_chat_failed",
14
  "ai_chat_no_stream",
15
  "ai_summarize_failed",
16
  "ai_summarize_empty",
17
  "ai_backend_failed",
18
  "ai_failed",
19
  "page_empty",
20
  "rate_limited",
21
  "unauthorized",
22
  "not_found",
23
  "request_failed",
24
  "validation_error",
25
]);
26
export type AiErrorCode = z.infer<typeof AiErrorCode>;
27
 
28
export type AiSseFrame =
29
  | { event: "chunk"; data: { text: string } }
30
  | { event: "done"; data: { usage?: AiUsage } }
31
  | { event: "error"; data: { message: string; code?: AiErrorCode } };
32
 
33
export function encodeAiSseFrame(frame: AiSseFrame): string {
34
  return `event: ${frame.event}\ndata: ${JSON.stringify(frame.data)}\n\n`;
35
}
36
 
37
export function encodeAiSseChunk(text: string): Uint8Array {
38
  return new TextEncoder().encode(encodeAiSseFrame({ event: "chunk", data: { text } }));
39
}
40
 
41
export function encodeAiSseDone(usage?: AiUsage): Uint8Array {
42
  const data = usage ? { usage } : {};
43
  return new TextEncoder().encode(encodeAiSseFrame({ event: "done", data }));
44
}
45
 
46
export function encodeAiSseError(message: string, code?: AiErrorCode): Uint8Array {
47
  return new TextEncoder().encode(encodeAiSseFrame({ event: "error", data: { message, code } }));
48
}
49
 
50
export async function* parseAiSseStream(body: ReadableStream<Uint8Array>): AsyncGenerator<AiSseFrame> {
51
  const reader = body.getReader();
52
  const decoder = new TextDecoder();
53
  let buffer = "";
54
 
55
  try {
56
    while (true) {
57
      const { value, done } = await reader.read();
58
      if (done) break;
59
      buffer += decoder.decode(value, { stream: true });
60
 
61
      let delimiterIndex = buffer.indexOf("\n\n");
62
      while (delimiterIndex !== -1) {
63
        const rawFrame = buffer.slice(0, delimiterIndex);
64
        buffer = buffer.slice(delimiterIndex + 2);
65
        const parsed = parseFrame(rawFrame);
66
        if (parsed) yield parsed;
67
        delimiterIndex = buffer.indexOf("\n\n");
68
      }
69
    }
70
  } finally {
71
    reader.releaseLock();
72
  }
73
}
74
 
75
function parseFrame(raw: string): AiSseFrame | null {
76
  let event: string | null = null;
77
  let data: string | null = null;
78
  for (const line of raw.split("\n")) {
79
    if (line.startsWith("event: ")) event = line.slice(7).trim();
80
    else if (line.startsWith("data: ")) data = line.slice(6);
81
  }
82
  if (!event || !data) return null;
83
  let parsedData: unknown;
84
  try {
85
    parsedData = JSON.parse(data);
86
  } catch {
87
    return null;
88
  }
89
  if (event === "chunk" && isChunkData(parsedData)) return { event: "chunk", data: parsedData };
90
  if (event === "done") return { event: "done", data: parseDoneData(parsedData) };
91
  if (event === "error" && isErrorData(parsedData)) return { event: "error", data: parsedData };
92
  return null;
93
}
94
 
95
function isChunkData(value: unknown): value is { text: string } {
96
  return typeof value === "object" && value !== null && typeof (value as { text?: unknown }).text === "string";
97
}
98
 
99
function isErrorData(value: unknown): value is { message: string; code?: AiErrorCode } {
100
  if (typeof value !== "object" || value === null) return false;
101
  const v = value as { message?: unknown; code?: unknown };
102
  if (typeof v.message !== "string") return false;
103
  if (v.code === undefined) return true;
104
  return typeof v.code === "string" && AiErrorCode.safeParse(v.code).success;
105
}
106
 
107
function parseDoneData(value: unknown): { usage?: AiUsage } {
108
  if (typeof value !== "object" || value === null) return {};
109
  const usage = (value as { usage?: unknown }).usage;
110
  if (typeof usage !== "object" || usage === null) return {};
111
  const u = usage as { promptTokens?: unknown; completionTokens?: unknown; totalTokens?: unknown };
112
  const parsed: AiUsage = {};
113
  if (typeof u.promptTokens === "number") parsed.promptTokens = u.promptTokens;
114
  if (typeof u.completionTokens === "number") parsed.completionTokens = u.completionTokens;
115
  if (typeof u.totalTokens === "number") parsed.totalTokens = u.totalTokens;
116
  return Object.keys(parsed).length > 0 ? { usage: parsed } : {};
117
}