100 lines
3.8 KiB
TypeScript
100 lines
3.8 KiB
TypeScript
// LocalLLMService.chatStream 이 @d3ro/core/llm-chat 계약을 따르는지 (red-team r1-6)
|
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
|
import { ErrorCode } from '@d3ro/core/errors'
|
|
import { initInMemoryConfig, resetInMemoryConfig } from '../../../src/main/services/ConfigService'
|
|
import {
|
|
getLocalLLMService,
|
|
resetLocalLLMServiceForTests,
|
|
} from '../../../src/main/services/LocalLLMService'
|
|
|
|
const encoder = new TextEncoder()
|
|
|
|
function ndjsonResponse(parts: string[]): Response {
|
|
return new Response(new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
for (const part of parts) controller.enqueue(encoder.encode(part))
|
|
controller.close()
|
|
},
|
|
}), { status: 200 })
|
|
}
|
|
|
|
async function collect(gen: AsyncGenerator<string, string>): Promise<{ tokens: string[]; result: string }> {
|
|
const tokens: string[] = []
|
|
for (let next = await gen.next(); ; next = await gen.next()) {
|
|
if (next.done) return { tokens, result: next.value }
|
|
tokens.push(next.value)
|
|
}
|
|
}
|
|
|
|
describe('LocalLLMService chat contract', () => {
|
|
beforeEach(() => {
|
|
initInMemoryConfig()
|
|
resetLocalLLMServiceForTests()
|
|
;(getLocalLLMService() as unknown as { _available: boolean })._available = true
|
|
})
|
|
|
|
afterEach(() => {
|
|
resetLocalLLMServiceForTests()
|
|
resetInMemoryConfig()
|
|
vi.unstubAllGlobals()
|
|
})
|
|
|
|
it('sends the system prompt as a single leading system message', async () => {
|
|
let sent: Array<{ role: string; content: string }> = []
|
|
vi.stubGlobal('fetch', vi.fn(async (_input: RequestInfo | URL, init?: RequestInit) => {
|
|
sent = (JSON.parse(String(init?.body)) as { messages: typeof sent }).messages
|
|
return ndjsonResponse([
|
|
`${JSON.stringify({ message: { content: 'hi' }, done: false })}\n`,
|
|
JSON.stringify({ message: { content: '' }, done: true }),
|
|
])
|
|
}))
|
|
|
|
const { tokens, result } = await collect(getLocalLLMService().chatStream([
|
|
{ role: 'system', content: 'persona' },
|
|
{ role: 'user', content: 'hello' },
|
|
]))
|
|
|
|
expect(tokens).toEqual(['hi'])
|
|
expect(result).toBe('hi')
|
|
expect(sent).toEqual([
|
|
{ role: 'system', content: 'persona' },
|
|
{ role: 'user', content: 'hello' },
|
|
])
|
|
})
|
|
|
|
it('rejects a stream that ends without a done frame', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse([
|
|
`${JSON.stringify({ message: { content: 'partial' }, done: false })}\n`,
|
|
])))
|
|
|
|
await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
|
|
.rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed })
|
|
})
|
|
|
|
it('rejects malformed NDJSON frames, including non-object JSON', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse(['{broken\n'])))
|
|
await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
|
|
.rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed, message: 'Ollama returned malformed NDJSON' })
|
|
|
|
vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse(['null\n'])))
|
|
await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
|
|
.rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed, message: 'Ollama returned malformed NDJSON' })
|
|
})
|
|
|
|
it('streamGenerate still completes on a trailing done frame without newline', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse([
|
|
`${JSON.stringify({ model: 'm', response: 'a', done: false })}\n`,
|
|
JSON.stringify({ model: 'm', response: 'b', done: true, eval_count: 2 }),
|
|
])))
|
|
|
|
const gen = getLocalLLMService().streamGenerate('p')
|
|
const tokens: string[] = []
|
|
let next = await gen.next()
|
|
while (!next.done) {
|
|
tokens.push(next.value)
|
|
next = await gen.next()
|
|
}
|
|
expect(tokens).toEqual(['a', 'b'])
|
|
expect(next.value).toMatchObject({ text: 'ab', completionTokens: 2 })
|
|
})
|
|
})
|