Skip to content

Commit 4876302

Browse files
IzumiSyIzumiSy
committed
Refactor AI chat client around completion events
1 parent d4a7609 commit 4876302

5 files changed

Lines changed: 123 additions & 64 deletions

File tree

‎packages/core/src/ai/client.test.ts‎

Lines changed: 45 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
import { describe, expect, it, vi } from "vitest";
22
import type { AuthClient } from "@tailor-platform/auth-public-client";
3-
import { createAIGatewayClient, type AIGatewayChatRequest, type AIGatewayClient } from "./client";
3+
import {
4+
createAIGatewayClient,
5+
type AIGatewayChatRequest,
6+
type AIGatewayClient,
7+
type AIChatCompletionEvent,
8+
} from "./client";
49

510
function createMockAuthClient(response: Response | Promise<Response>) {
611
return {
@@ -29,24 +34,25 @@ function createSSEStream(chunks: string[]): ReadableStream<Uint8Array> {
2934
});
3035
}
3136

32-
async function collectStream(client: AIGatewayClient, request: AIGatewayChatRequest) {
33-
const deltas: string[] = [];
37+
async function collectEvents(client: AIGatewayClient, request: AIGatewayChatRequest) {
38+
const events: AIChatCompletionEvent[] = [];
3439

35-
for await (const delta of client.chatCompletionStream(request)) {
36-
deltas.push(delta);
40+
for await (const event of client.streamChatCompletion(request)) {
41+
events.push(event);
3742
}
3843

39-
return deltas;
44+
return events;
4045
}
4146

4247
describe("createAIGatewayClient", () => {
43-
it("streams OpenAI-compatible text deltas through authClient.fetch", async () => {
48+
it("streams OpenAI-compatible events through authClient.fetch", async () => {
4449
const authClient = createMockAuthClient(
4550
new Response(
4651
createSSEStream([
4752
'data: {"choices":[{"delta":{"content":"Hel',
4853
'lo"}}]}\n\n',
4954
'data: {"choices":[{"delta":{"content":" world"}}]}\n\n',
55+
'data: {"choices":[{"finish_reason":"stop"}]}\n\n',
5056
"data: [DONE]\n\n",
5157
]),
5258
{
@@ -61,9 +67,13 @@ describe("createAIGatewayClient", () => {
6167
authClient,
6268
});
6369

64-
const deltas = await collectStream(client, createRequest({ signal }));
70+
const events = await collectEvents(client, createRequest({ signal }));
6571

66-
expect(deltas).toEqual(["Hello", " world"]);
72+
expect(events).toEqual([
73+
{ type: "text-delta", text: "Hello" },
74+
{ type: "text-delta", text: " world" },
75+
{ type: "done", finishReason: "stop" },
76+
]);
6777
expect(authClient.fetch).toHaveBeenCalledWith(
6878
"https://gateway.example.com/v1/chat/completions",
6979
expect.objectContaining({
@@ -82,7 +92,7 @@ describe("createAIGatewayClient", () => {
8292
});
8393
});
8494

85-
it("yields a single text chunk for json responses", async () => {
95+
it("yields text and done events for json responses", async () => {
8696
const authClient = createMockAuthClient(
8797
new Response(
8898
JSON.stringify({
@@ -91,6 +101,7 @@ describe("createAIGatewayClient", () => {
91101
message: {
92102
content: [{ type: "text", text: "Grounded answer" }],
93103
},
104+
finish_reason: "stop",
94105
},
95106
],
96107
}),
@@ -105,12 +116,15 @@ describe("createAIGatewayClient", () => {
105116
authClient,
106117
});
107118

108-
const deltas = await collectStream(
119+
const events = await collectEvents(
109120
client,
110121
createRequest({ model: "gemini-2.5-flash", stream: false }),
111122
);
112123

113-
expect(deltas).toEqual(["Grounded answer"]);
124+
expect(events).toEqual([
125+
{ type: "text-delta", text: "Grounded answer" },
126+
{ type: "done", finishReason: "stop" },
127+
]);
114128
expect(authClient.fetch).toHaveBeenCalledWith(
115129
"https://gateway.example.com/base/v1/chat/completions",
116130
expect.objectContaining({
@@ -127,43 +141,46 @@ describe("createAIGatewayClient", () => {
127141
});
128142
});
129143

130-
it("throws on non-ok responses", async () => {
144+
it("emits done even when json responses do not include assistant text", async () => {
131145
const authClient = createMockAuthClient(
132-
new Response("nope", { status: 401, statusText: "Nope" }),
146+
new Response(
147+
JSON.stringify({ choices: [{ message: { content: [] }, finish_reason: "tool_calls" }] }),
148+
{ status: 200 },
149+
),
133150
);
134151
const client = createAIGatewayClient({
135152
gatewayUri: "https://gateway.example.com",
136153
authClient,
137154
});
138155

139-
await expect(collectStream(client, createRequest())).rejects.toThrow(
140-
"AI Gateway streaming request failed (401 Nope): nope",
141-
);
156+
await expect(
157+
collectEvents(client, createRequest({ model: "gemini-2.5-flash", stream: false })),
158+
).resolves.toEqual([{ type: "done", finishReason: "tool_calls" }]);
142159
});
143160

144-
it("throws when a streaming response body is missing", async () => {
145-
const authClient = createMockAuthClient(new Response(null, { status: 200 }));
161+
it("throws on non-ok responses", async () => {
162+
const authClient = createMockAuthClient(
163+
new Response("nope", { status: 401, statusText: "Nope" }),
164+
);
146165
const client = createAIGatewayClient({
147166
gatewayUri: "https://gateway.example.com",
148167
authClient,
149168
});
150169

151-
await expect(collectStream(client, createRequest())).rejects.toThrow(
152-
"AI Gateway streaming response did not include a body.",
170+
await expect(collectEvents(client, createRequest())).rejects.toThrow(
171+
"AI Gateway streaming request failed (401 Nope): nope",
153172
);
154173
});
155174

156-
it("throws when json responses do not include assistant text", async () => {
157-
const authClient = createMockAuthClient(
158-
new Response(JSON.stringify({ choices: [{ message: { content: [] } }] }), { status: 200 }),
159-
);
175+
it("throws when a streaming response body is missing", async () => {
176+
const authClient = createMockAuthClient(new Response(null, { status: 200 }));
160177
const client = createAIGatewayClient({
161178
gatewayUri: "https://gateway.example.com",
162179
authClient,
163180
});
164181

165-
await expect(
166-
collectStream(client, createRequest({ model: "gemini-2.5-flash", stream: false })),
167-
).rejects.toThrow("AI Gateway JSON response did not include assistant text.");
182+
await expect(collectEvents(client, createRequest())).rejects.toThrow(
183+
"AI Gateway streaming response did not include a body.",
184+
);
168185
});
169186
});

‎packages/core/src/ai/client.ts‎

Lines changed: 47 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
11
import type { AuthClient } from "@tailor-platform/auth-public-client";
22

3-
export interface AIGatewayChatMessage {
4-
role: "system" | "user" | "assistant";
5-
content: string;
6-
}
3+
export type AIGatewayChatMessage =
4+
| {
5+
role: "system" | "user";
6+
content: string;
7+
}
8+
| {
9+
role: "assistant";
10+
content?: string;
11+
};
712

813
export interface AIGatewayChatRequest {
914
model: string;
@@ -12,18 +17,29 @@ export interface AIGatewayChatRequest {
1217
signal?: AbortSignal;
1318
}
1419

20+
export type AIChatCompletionEvent =
21+
| {
22+
type: "text-delta";
23+
text: string;
24+
}
25+
| {
26+
type: "done";
27+
finishReason?: string;
28+
};
29+
1530
export interface AIGatewayClient {
1631
/**
1732
* Honor request.signal when possible so callers can stop work early.
1833
*/
19-
chatCompletionStream(request: AIGatewayChatRequest): AsyncIterable<string>;
34+
streamChatCompletion(request: AIGatewayChatRequest): AsyncIterable<AIChatCompletionEvent>;
2035
}
2136

2237
interface OpenAIStreamChunk {
2338
choices?: Array<{
2439
delta?: {
2540
content?: unknown;
2641
};
42+
finish_reason?: unknown;
2743
}>;
2844
}
2945

@@ -32,6 +48,7 @@ interface OpenAIFinalResponse {
3248
message?: {
3349
content?: unknown;
3450
};
51+
finish_reason?: unknown;
3552
}>;
3653
}
3754

@@ -42,7 +59,7 @@ export function createAIGatewayClient(config: {
4259
const endpoint = new URL("v1/chat/completions", withTrailingSlash(config.gatewayUri)).toString();
4360

4461
return {
45-
async *chatCompletionStream(request) {
62+
async *streamChatCompletion(request) {
4663
if (request.stream === false) {
4764
yield* streamJSONResponse({
4865
endpoint,
@@ -65,7 +82,7 @@ async function* streamOpenAICompatibleResponse(input: {
6582
endpoint: string;
6683
authClient: AuthClient;
6784
request: AIGatewayChatRequest;
68-
}): AsyncGenerator<string, void, unknown> {
85+
}): AsyncGenerator<AIChatCompletionEvent, void, unknown> {
6986
const response = await input.authClient.fetch(input.endpoint, {
7087
method: "POST",
7188
headers: {
@@ -86,25 +103,32 @@ async function* streamOpenAICompatibleResponse(input: {
86103
throw new Error("AI Gateway streaming response did not include a body.");
87104
}
88105

106+
let finishReason: string | undefined;
107+
89108
for await (const event of iterateSSEDataEvents(response.body, input.request.signal)) {
90109
if (event === "[DONE]") {
110+
yield createDoneEvent(finishReason);
91111
return;
92112
}
93113

94114
const payload = parseJSON<OpenAIStreamChunk>(event, "AI Gateway SSE event");
95-
const delta = extractText(payload.choices?.[0]?.delta?.content);
115+
const choice = payload.choices?.[0];
116+
const delta = extractText(choice?.delta?.content);
117+
finishReason = extractFinishReason(choice?.finish_reason) ?? finishReason;
96118

97119
if (delta) {
98-
yield delta;
120+
yield { type: "text-delta", text: delta };
99121
}
100122
}
123+
124+
yield createDoneEvent(finishReason);
101125
}
102126

103127
async function* streamJSONResponse(input: {
104128
endpoint: string;
105129
authClient: AuthClient;
106130
request: AIGatewayChatRequest;
107-
}): AsyncGenerator<string, void, unknown> {
131+
}): AsyncGenerator<AIChatCompletionEvent, void, unknown> {
108132
const response = await input.authClient.fetch(input.endpoint, {
109133
method: "POST",
110134
headers: {
@@ -122,13 +146,14 @@ async function* streamJSONResponse(input: {
122146
await assertOK(response, "AI Gateway JSON request");
123147

124148
const payload = parseJSON<OpenAIFinalResponse>(await response.text(), "AI Gateway JSON response");
125-
const text = extractText(payload.choices?.[0]?.message?.content);
149+
const choice = payload.choices?.[0];
150+
const text = extractText(choice?.message?.content);
126151

127-
if (!text) {
128-
throw new Error("AI Gateway JSON response did not include assistant text.");
152+
if (text) {
153+
yield { type: "text-delta", text };
129154
}
130155

131-
yield text;
156+
yield createDoneEvent(extractFinishReason(choice?.finish_reason));
132157
}
133158

134159
async function assertOK(response: Response, context: string): Promise<void> {
@@ -263,6 +288,14 @@ function extractText(content: unknown): string {
263288
.join("");
264289
}
265290

291+
function extractFinishReason(value: unknown): string | undefined {
292+
return typeof value === "string" ? value : undefined;
293+
}
294+
295+
function createDoneEvent(finishReason?: string): AIChatCompletionEvent {
296+
return finishReason ? { type: "done", finishReason } : { type: "done" };
297+
}
298+
266299
function parseJSON<T>(value: string, context: string): T {
267300
try {
268301
return JSON.parse(value) as T;

0 commit comments

Comments
 (0)