1
import { AiErrorCode, parseAiSseStream, type AiUsage } from "@/shared/ai";
2
import type { AiAskRequest, AiGenerateRequest, AiRewriteRequest, AiSummarizeResponse, ApiError } from "@/shared/types";
3
import { sendApiRequest } from "@/client/lib/api";
4
 
5
export interface AiStreamChunk {
6
  text: string;
7
}
8
 
9
export interface AiStreamFinal {
10
  usage?: AiUsage;
11
}
12
 
13
export class AiStreamError extends Error {
14
  readonly code: AiErrorCode;
15
  constructor(message: string, code: AiErrorCode) {
16
    super(message);
17
    this.name = "AiStreamError";
18
    this.code = code;
19
  }
20
}
21
 
22
function coerceErrorCode(raw: unknown): AiErrorCode {
23
  if (typeof raw !== "string") return "ai_failed";
24
  const parsed = AiErrorCode.safeParse(raw);
25
  return parsed.success ? parsed.data : "ai_failed";
26
}
27
 
28
function wrapThrownError(err: unknown): AiStreamError {
29
  if (err instanceof AiStreamError) return err;
30
  if (typeof err === "object" && err !== null && "message" in err && "error" in err) {
31
    const payload = err as ApiError;
32
    return new AiStreamError(payload.message ?? "AI request failed", coerceErrorCode(payload.error));
33
  }
34
  const message = err instanceof Error ? err.message : "AI request failed";
35
  return new AiStreamError(message, "ai_failed");
36
}
37
 
38
export function streamRewrite(
39
  workspaceId: string,
40
  pageId: string,
41
  body: AiRewriteRequest,
42
  signal?: AbortSignal,
43
): AsyncGenerator<AiStreamChunk, AiStreamFinal, void> {
44
  return streamAiPost(`/workspaces/${workspaceId}/pages/${pageId}/rewrite`, body, signal);
45
}
46
 
47
export function streamGenerate(
48
  workspaceId: string,
49
  pageId: string,
50
  body: AiGenerateRequest,
51
  signal?: AbortSignal,
52
): AsyncGenerator<AiStreamChunk, AiStreamFinal, void> {
53
  return streamAiPost(`/workspaces/${workspaceId}/pages/${pageId}/generate`, body, signal);
54
}
55
 
56
export function streamAskPage(
57
  workspaceId: string,
58
  pageId: string,
59
  body: AiAskRequest,
60
  signal?: AbortSignal,
61
): AsyncGenerator<AiStreamChunk, AiStreamFinal, void> {
62
  return streamAiPost(`/workspaces/${workspaceId}/pages/${pageId}/ask`, body, signal);
63
}
64
 
65
export async function summarizePage(workspaceId: string, pageId: string): Promise<AiSummarizeResponse> {
66
  try {
67
    const response = await sendApiRequest(`/workspaces/${workspaceId}/pages/${pageId}/summarize`, { method: "POST" });
68
    return (await response.json()) as AiSummarizeResponse;
69
  } catch (err) {
70
    throw wrapThrownError(err);
71
  }
72
}
73
 
74
async function* streamAiPost(
75
  path: string,
76
  body: unknown,
77
  signal?: AbortSignal,
78
): AsyncGenerator<AiStreamChunk, AiStreamFinal, void> {
79
  let response: Response;
80
  try {
81
    response = await sendApiRequest(path, {
82
      method: "POST",
83
      body: JSON.stringify(body),
84
      headers: { accept: "text/event-stream" },
85
      signal,
86
    });
87
  } catch (err) {
88
    throw wrapThrownError(err);
89
  }
90
 
91
  if (!response.body) {
92
    throw new AiStreamError("Empty response body", "ai_failed");
93
  }
94
 
95
  for await (const frame of parseAiSseStream(response.body)) {
96
    if (frame.event === "chunk") {
97
      yield { text: frame.data.text };
98
      continue;
99
    }
100
    if (frame.event === "error") {
101
      throw new AiStreamError(frame.data.message, frame.data.code ?? "ai_failed");
102
    }
103
    if (frame.event === "done") {
104
      return frame.data.usage ? { usage: frame.data.usage } : {};
105
    }
106
  }
107
  return {};
108
}