Skip to content

Commit a512a6e

Browse files
roomoteroomote
committed
refactor: roll out stream helpers to friendli opencode mimo specs
1 parent 4497c6f commit a512a6e

4 files changed

Lines changed: 279 additions & 475 deletions

File tree

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

Lines changed: 43 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { friendliDefaultModelId, friendliModels } from "@roo-code/types"
88
import { buildApiHandler } from "../../index"
99
import { getModelMaxOutputTokens } from "../../../shared/api"
1010
import { FriendliHandler } from "../friendli"
11+
import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
1112

1213
// Create mock functions
1314
const mockCreate = vi.fn()
@@ -31,18 +32,18 @@ describe("FriendliHandler", () => {
3132
beforeEach(() => {
3233
vi.clearAllMocks()
3334
// Set up default mock implementation
34-
mockCreate.mockImplementation(async () => ({
35-
[Symbol.asyncIterator]: async function* () {
36-
yield {
35+
mockCreate.mockImplementation(async () =>
36+
asyncStreamFrom([
37+
{
3738
choices: [
3839
{
3940
delta: { content: "Test response" },
4041
index: 0,
4142
},
4243
],
4344
usage: null,
44-
}
45-
yield {
45+
},
46+
{
4647
choices: [
4748
{
4849
delta: {},
@@ -54,9 +55,9 @@ describe("FriendliHandler", () => {
5455
completion_tokens: 5,
5556
total_tokens: 15,
5657
},
57-
}
58-
},
59-
}))
58+
},
59+
]),
60+
)
6061
handler = new FriendliHandler({ friendliApiKey: "test-key" })
6162
})
6263

@@ -189,19 +190,7 @@ describe("FriendliHandler", () => {
189190
it("createMessage should yield text content from stream", async () => {
190191
const testContent = "This is test content from Friendli stream"
191192

192-
mockCreate.mockImplementationOnce(() => {
193-
return {
194-
[Symbol.asyncIterator]: () => ({
195-
next: vi
196-
.fn()
197-
.mockResolvedValueOnce({
198-
done: false,
199-
value: { choices: [{ delta: { content: testContent } }] },
200-
})
201-
.mockResolvedValueOnce({ done: true }),
202-
}),
203-
}
204-
})
193+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([{ choices: [{ delta: { content: testContent } }] }]))
205194

206195
const stream = handler.createMessage("system prompt", [])
207196
const firstChunk = await stream.next()
@@ -211,19 +200,9 @@ describe("FriendliHandler", () => {
211200
})
212201

213202
it("createMessage should yield usage data from stream", async () => {
214-
mockCreate.mockImplementationOnce(() => {
215-
return {
216-
[Symbol.asyncIterator]: () => ({
217-
next: vi
218-
.fn()
219-
.mockResolvedValueOnce({
220-
done: false,
221-
value: { choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 20 } },
222-
})
223-
.mockResolvedValueOnce({ done: true }),
224-
}),
225-
}
226-
})
203+
mockCreate.mockImplementationOnce(() =>
204+
asyncStreamFrom([{ choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 20 } }]),
205+
)
227206

228207
const stream = handler.createMessage("system prompt", [])
229208
const firstChunk = await stream.next()
@@ -240,15 +219,7 @@ describe("FriendliHandler", () => {
240219
friendliApiKey: "test-friendli-api-key",
241220
})
242221

243-
mockCreate.mockImplementationOnce(() => {
244-
return {
245-
[Symbol.asyncIterator]: () => ({
246-
async next() {
247-
return { done: true }
248-
},
249-
}),
250-
}
251-
})
222+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
252223

253224
const systemPrompt = "Test system prompt for Friendli"
254225
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Friendli" }]
@@ -276,13 +247,7 @@ describe("FriendliHandler", () => {
276247
modelTemperature: 0.3,
277248
})
278249

279-
mockCreate.mockImplementationOnce(() => ({
280-
[Symbol.asyncIterator]: () => ({
281-
async next() {
282-
return { done: true }
283-
},
284-
}),
285-
}))
250+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
286251

287252
const messageGenerator = handlerWithModel.createMessage("system", [])
288253
await messageGenerator.next()
@@ -308,27 +273,27 @@ describe("FriendliHandler", () => {
308273
})
309274

310275
it("createMessage should handle stream with multiple chunks", async () => {
311-
mockCreate.mockImplementationOnce(async () => ({
312-
[Symbol.asyncIterator]: async function* () {
313-
yield {
276+
mockCreate.mockImplementationOnce(async () =>
277+
asyncStreamFrom([
278+
{
314279
choices: [
315280
{
316281
delta: { content: "Hello" },
317282
index: 0,
318283
},
319284
],
320285
usage: null,
321-
}
322-
yield {
286+
},
287+
{
323288
choices: [
324289
{
325290
delta: { content: " world" },
326291
index: 0,
327292
},
328293
],
329294
usage: null,
330-
}
331-
yield {
295+
},
296+
{
332297
choices: [
333298
{
334299
delta: {},
@@ -340,18 +305,15 @@ describe("FriendliHandler", () => {
340305
completion_tokens: 10,
341306
total_tokens: 15,
342307
},
343-
}
344-
},
345-
}))
308+
},
309+
]),
310+
)
346311

347312
const systemPrompt = "You are a helpful assistant."
348313
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hi" }]
349314

350315
const stream = handler.createMessage(systemPrompt, messages)
351-
const chunks = []
352-
for await (const chunk of stream) {
353-
chunks.push(chunk)
354-
}
316+
const chunks = await collectStream(stream)
355317

356318
expect(chunks[0]).toEqual({ type: "text", text: "Hello" })
357319
expect(chunks[1]).toEqual({ type: "text", text: " world" })
@@ -417,13 +379,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
417379
reasoningEffort: "high",
418380
})
419381

420-
mockCreate.mockImplementationOnce(() => ({
421-
[Symbol.asyncIterator]: () => ({
422-
async next() {
423-
return { done: true }
424-
},
425-
}),
426-
}))
382+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
427383

428384
await handler.createMessage("system", []).next()
429385

@@ -446,13 +402,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
446402
enableReasoningEffort: false,
447403
})
448404

449-
mockCreate.mockImplementationOnce(() => ({
450-
[Symbol.asyncIterator]: () => ({
451-
async next() {
452-
return { done: true }
453-
},
454-
}),
455-
}))
405+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
456406

457407
await handler.createMessage("system", []).next()
458408

@@ -471,13 +421,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
471421
reasoningEffort: "none",
472422
})
473423

474-
mockCreate.mockImplementationOnce(() => ({
475-
[Symbol.asyncIterator]: () => ({
476-
async next() {
477-
return { done: true }
478-
},
479-
}),
480-
}))
424+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
481425

482426
await handler.createMessage("system", []).next()
483427

@@ -496,13 +440,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
496440
reasoningEffort: "disable",
497441
})
498442

499-
mockCreate.mockImplementationOnce(() => ({
500-
[Symbol.asyncIterator]: () => ({
501-
async next() {
502-
return { done: true }
503-
},
504-
}),
505-
}))
443+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
506444

507445
await handler.createMessage("system", []).next()
508446

@@ -520,13 +458,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
520458
// No enableReasoningEffort or reasoningEffort — model default "high" kicks in
521459
})
522460

523-
mockCreate.mockImplementationOnce(() => ({
524-
[Symbol.asyncIterator]: () => ({
525-
async next() {
526-
return { done: true }
527-
},
528-
}),
529-
}))
461+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
530462

531463
await handler.createMessage("system", []).next()
532464

@@ -545,13 +477,7 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
545477
reasoningEffort: "high",
546478
})
547479

548-
mockCreate.mockImplementationOnce(() => ({
549-
[Symbol.asyncIterator]: () => ({
550-
async next() {
551-
return { done: true }
552-
},
553-
}),
554-
}))
480+
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
555481

556482
await handler.createMessage("system", []).next()
557483

@@ -569,28 +495,25 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => {
569495
reasoningEffort: "high",
570496
})
571497

572-
mockCreate.mockImplementationOnce(async () => ({
573-
[Symbol.asyncIterator]: async function* () {
574-
yield {
498+
mockCreate.mockImplementationOnce(async () =>
499+
asyncStreamFrom([
500+
{
575501
choices: [{ delta: { reasoning_content: "Let me think..." } }],
576502
usage: null,
577-
}
578-
yield {
503+
},
504+
{
579505
choices: [{ delta: { content: "The answer is 42" } }],
580506
usage: null,
581-
}
582-
yield {
507+
},
508+
{
583509
choices: [{ delta: {} }],
584510
usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 },
585-
}
586-
},
587-
}))
511+
},
512+
]),
513+
)
588514

589515
const stream = handler.createMessage("system", [])
590-
const chunks = []
591-
for await (const chunk of stream) {
592-
chunks.push(chunk)
593-
}
516+
const chunks = await collectStream(stream)
594517

595518
expect(chunks).toContainEqual({ type: "reasoning", text: "Let me think..." })
596519
expect(chunks).toContainEqual({ type: "text", text: "The answer is 42" })

0 commit comments

Comments
 (0)