Skip to content

Commit b47b19a

Browse files
zoomote[bot]roomote
zoomote[bot]
and
roomote
authored
refactor: roll out stream helpers to requesty spec (#1089)
Co-authored-by: Roomote <roomote@roomote.dev>
1 parent 69a6f79 commit b47b19a

3 files changed

Lines changed: 99 additions & 111 deletions

File tree

‎src/api/providers/__tests__/requesty.spec.ts‎

Lines changed: 79 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { RequestyHandler } from "../requesty"
1313
import { ApiHandlerOptions } from "../../../shared/api"
1414
import { Package } from "../../../shared/package"
1515
import { ApiHandlerCreateMessageMetadata } from "../../index"
16+
import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
1617

1718
const mockCreate = vitest.fn()
1819

@@ -182,38 +183,31 @@ describe("RequestyHandler", () => {
182183
it("generates correct stream chunks", async () => {
183184
const handler = new RequestyHandler(mockOptions)
184185

185-
const mockStream = {
186-
async *[Symbol.asyncIterator]() {
187-
yield {
188-
id: mockOptions.requestyModelId,
189-
choices: [{ delta: { content: "test response" } }],
190-
}
191-
yield {
192-
id: "test-id",
193-
choices: [{ delta: {} }],
194-
usage: {
195-
prompt_tokens: 10,
196-
completion_tokens: 20,
197-
prompt_tokens_details: {
198-
caching_tokens: 5,
199-
cached_tokens: 2,
200-
},
186+
const mockStream = asyncStreamFrom([
187+
{
188+
id: mockOptions.requestyModelId,
189+
choices: [{ delta: { content: "test response" } }],
190+
},
191+
{
192+
id: "test-id",
193+
choices: [{ delta: {} }],
194+
usage: {
195+
prompt_tokens: 10,
196+
completion_tokens: 20,
197+
prompt_tokens_details: {
198+
caching_tokens: 5,
199+
cached_tokens: 2,
201200
},
202-
}
201+
},
203202
},
204-
}
203+
])
205204

206205
mockCreate.mockResolvedValue(mockStream)
207206

208207
const systemPrompt = "test system prompt"
209208
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "test message" }]
210209

211-
const generator = handler.createMessage(systemPrompt, messages)
212-
const chunks = []
213-
214-
for await (const chunk of generator) {
215-
chunks.push(chunk)
216-
}
210+
const chunks = await collectStream(handler.createMessage(systemPrompt, messages))
217211

218212
// Verify stream chunks
219213
expect(chunks).toHaveLength(2) // One text chunk and one usage chunk
@@ -257,15 +251,13 @@ describe("RequestyHandler", () => {
257251
modelMaxTokens: 32768,
258252
})
259253

260-
const mockStream = {
261-
async *[Symbol.asyncIterator]() {
262-
yield {
263-
id: "test-id",
264-
choices: [{ delta: {} }],
265-
usage: { prompt_tokens: 10, completion_tokens: 20 },
266-
}
254+
const mockStream = asyncStreamFrom([
255+
{
256+
id: "test-id",
257+
choices: [{ delta: {} }],
258+
usage: { prompt_tokens: 10, completion_tokens: 20 },
267259
},
268-
}
260+
])
269261

270262
mockCreate.mockResolvedValue(mockStream)
271263

@@ -290,15 +282,13 @@ describe("RequestyHandler", () => {
290282
modelMaxTokens: 32768,
291283
})
292284

293-
const mockStream = {
294-
async *[Symbol.asyncIterator]() {
295-
yield {
296-
id: "test-id",
297-
choices: [{ delta: {} }],
298-
usage: { prompt_tokens: 10, completion_tokens: 20 },
299-
}
285+
const mockStream = asyncStreamFrom([
286+
{
287+
id: "test-id",
288+
choices: [{ delta: {} }],
289+
usage: { prompt_tokens: 10, completion_tokens: 20 },
300290
},
301-
}
291+
])
302292

303293
mockCreate.mockResolvedValue(mockStream)
304294

@@ -323,15 +313,13 @@ describe("RequestyHandler", () => {
323313
modelMaxTokens: 32768,
324314
})
325315

326-
const mockStream = {
327-
async *[Symbol.asyncIterator]() {
328-
yield {
329-
id: "test-id",
330-
choices: [{ delta: {} }],
331-
usage: { prompt_tokens: 10, completion_tokens: 20 },
332-
}
316+
const mockStream = asyncStreamFrom([
317+
{
318+
id: "test-id",
319+
choices: [{ delta: {} }],
320+
usage: { prompt_tokens: 10, completion_tokens: 20 },
333321
},
334-
}
322+
])
335323

336324
mockCreate.mockResolvedValue(mockStream)
337325

@@ -359,53 +347,47 @@ describe("RequestyHandler", () => {
359347

360348
it("streams reasoning chunks from delta.reasoning_content", async () => {
361349
const handler = new RequestyHandler(mockOptions)
362-
mockCreate.mockResolvedValue({
363-
async *[Symbol.asyncIterator]() {
364-
yield { id: "1", choices: [{ delta: { reasoning_content: "thinking..." } }] }
365-
yield { id: "1", choices: [{ delta: { content: "answer" } }] }
366-
yield {
350+
mockCreate.mockResolvedValue(
351+
asyncStreamFrom([
352+
{ id: "1", choices: [{ delta: { reasoning_content: "thinking..." } }] },
353+
{ id: "1", choices: [{ delta: { content: "answer" } }] },
354+
{
367355
id: "1",
368356
choices: [{ delta: {} }],
369357
usage: { prompt_tokens: 1, completion_tokens: 1 },
370-
}
371-
},
372-
})
358+
},
359+
]),
360+
)
373361

374-
const chunks: any[] = []
375-
for await (const chunk of handler.createMessage("sys", [{ role: "user", content: "hi" }])) {
376-
chunks.push(chunk)
377-
}
362+
const chunks = await collectStream(handler.createMessage("sys", [{ role: "user", content: "hi" }]))
378363

379364
expect(chunks).toContainEqual({ type: "reasoning", text: "thinking..." })
380365
})
381366

382367
it("falls back to delta.reasoning when reasoning_content is absent", async () => {
383368
const handler = new RequestyHandler(mockOptions)
384-
mockCreate.mockResolvedValue({
385-
async *[Symbol.asyncIterator]() {
386-
yield { id: "1", choices: [{ delta: { reasoning: "router-style thought" } }] }
387-
yield {
369+
mockCreate.mockResolvedValue(
370+
asyncStreamFrom([
371+
{ id: "1", choices: [{ delta: { reasoning: "router-style thought" } }] },
372+
{
388373
id: "1",
389374
choices: [{ delta: {} }],
390375
usage: { prompt_tokens: 1, completion_tokens: 1 },
391-
}
392-
},
393-
})
376+
},
377+
]),
378+
)
394379

395-
const chunks: any[] = []
396-
for await (const chunk of handler.createMessage("sys", [{ role: "user", content: "hi" }])) {
397-
chunks.push(chunk)
398-
}
380+
const chunks = await collectStream(handler.createMessage("sys", [{ role: "user", content: "hi" }]))
399381

400382
expect(chunks).toContainEqual({ type: "reasoning", text: "router-style thought" })
401383
})
402384

403385
it("prefers delta.reasoning_content over delta.reasoning when both are present", async () => {
404386
const handler = new RequestyHandler(mockOptions)
405387

406-
mockCreate.mockResolvedValue({
407-
async *[Symbol.asyncIterator]() {
408-
yield {
388+
mockCreate.mockResolvedValue(
389+
asyncStreamFrom([
390+
{
409391
id: "1",
410392
choices: [
411393
{
@@ -415,20 +397,16 @@ describe("RequestyHandler", () => {
415397
},
416398
},
417399
],
418-
}
419-
yield {
400+
},
401+
{
420402
id: "1",
421403
choices: [{ delta: {} }],
422404
usage: { prompt_tokens: 1, completion_tokens: 1 },
423-
}
424-
},
425-
})
426-
427-
const chunks: any[] = []
405+
},
406+
]),
407+
)
428408

429-
for await (const chunk of handler.createMessage("sys", [{ role: "user", content: "hi" }])) {
430-
chunks.push(chunk)
431-
}
409+
const chunks = await collectStream(handler.createMessage("sys", [{ role: "user", content: "hi" }]))
432410

433411
const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
434412

@@ -459,15 +437,14 @@ describe("RequestyHandler", () => {
459437
]
460438

461439
beforeEach(() => {
462-
const mockStream = {
463-
async *[Symbol.asyncIterator]() {
464-
yield {
440+
mockCreate.mockResolvedValue(
441+
asyncStreamFrom([
442+
{
465443
id: "test-id",
466444
choices: [{ delta: { content: "test response" } }],
467-
}
468-
},
469-
}
470-
mockCreate.mockResolvedValue(mockStream)
445+
},
446+
]),
447+
)
471448
})
472449

473450
it("should include tools in request when tools are provided", async () => {
@@ -498,9 +475,9 @@ describe("RequestyHandler", () => {
498475
})
499476

500477
it("should handle tool_call_partial chunks in streaming response", async () => {
501-
const mockStreamWithToolCalls = {
502-
async *[Symbol.asyncIterator]() {
503-
yield {
478+
mockCreate.mockResolvedValue(
479+
asyncStreamFrom([
480+
{
504481
id: "test-id",
505482
choices: [
506483
{
@@ -518,8 +495,8 @@ describe("RequestyHandler", () => {
518495
},
519496
},
520497
],
521-
}
522-
yield {
498+
},
499+
{
523500
id: "test-id",
524501
choices: [
525502
{
@@ -535,26 +512,22 @@ describe("RequestyHandler", () => {
535512
},
536513
},
537514
],
538-
}
539-
yield {
515+
},
516+
{
540517
id: "test-id",
541518
choices: [{ delta: {} }],
542519
usage: { prompt_tokens: 10, completion_tokens: 20 },
543-
}
544-
},
545-
}
546-
mockCreate.mockResolvedValue(mockStreamWithToolCalls)
520+
},
521+
]),
522+
)
547523

548524
const metadata: ApiHandlerCreateMessageMetadata = {
549525
taskId: "test-task",
550526
tools: mockTools,
551527
}
552528

553529
const handler = new RequestyHandler(mockOptions)
554-
const chunks = []
555-
for await (const chunk of handler.createMessage(systemPrompt, messages, metadata)) {
556-
chunks.push(chunk)
557-
}
530+
const chunks = await collectStream(handler.createMessage(systemPrompt, messages, metadata))
558531

559532
// Expect two tool_call_partial chunks and one usage chunk
560533
expect(chunks).toHaveLength(3)

‎src/eslint-suppressions.json‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -264,11 +264,6 @@
264264
"count": 5
265265
}
266266
},
267-
"api/providers/__tests__/requesty.spec.ts": {
268-
"@typescript-eslint/no-explicit-any": {
269-
"count": 3
270-
}
271-
},
272267
"api/providers/__tests__/sambanova.spec.ts": {
273268
"@typescript-eslint/no-explicit-any": {
274269
"count": 2

‎src/test-utils/__tests__/stream.spec.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import { asyncStreamFrom, collectStream } from "../stream"
2+
3+
describe("stream test utils", () => {
4+
it("collects chunks in order", async () => {
5+
await expect(collectStream(asyncStreamFrom([1, 2, 3]))).resolves.toEqual([1, 2, 3])
6+
})
7+
8+
it("collects empty streams", async () => {
9+
await expect(collectStream(asyncStreamFrom([]))).resolves.toEqual([])
10+
})
11+
12+
it("propagates stream errors", async () => {
13+
async function* failingStream() {
14+
yield 1
15+
throw new Error("boom")
16+
}
17+
18+
await expect(collectStream(failingStream())).rejects.toThrow("boom")
19+
})
20+
})

0 commit comments

Comments
 (0)