wehub-resource-sync
115 lines
3.7 KiB
TypeScript
115 lines
3.7 KiB
TypeScript
import test from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import fs from "node:fs";
|
|
import os from "node:os";
|
|
import path from "node:path";
|
|
|
|
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-stream-zero-prompt-"));
|
|
process.env.DATA_DIR = TEST_DATA_DIR;
|
|
const core = await import("../../src/lib/db/core.ts");
|
|
|
|
const { createSSEStream } = await import("../../open-sse/utils/stream.ts");
|
|
const { FORMATS } = await import("../../open-sse/translator/formats.ts");
|
|
|
|
const textEncoder = new TextEncoder();
|
|
|
|
async function readTransformed(chunks, options) {
|
|
const source = new ReadableStream({
|
|
start(controller) {
|
|
for (const chunk of chunks) {
|
|
controller.enqueue(textEncoder.encode(chunk));
|
|
}
|
|
controller.close();
|
|
},
|
|
});
|
|
|
|
return new Response(source.pipeThrough(createSSEStream(options))).text();
|
|
}
|
|
|
|
test("createSSEStream passthrough estimates input tokens when upstream reports prompt_tokens=0 (Ollama Cloud pattern)", async () => {
|
|
let onCompletePayload = null;
|
|
const body = { messages: [{ role: "user", content: "Write a 500-line research report" }] };
|
|
const text = await readTransformed(
|
|
[
|
|
`data: ${JSON.stringify({
|
|
id: "chatcmpl_ollama",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "minimax-m3",
|
|
system_fingerprint: "fp_ollama",
|
|
choices: [{ index: 0, delta: { content: "Here" } }],
|
|
})}\n\n`,
|
|
`data: ${JSON.stringify({
|
|
id: "chatcmpl_ollama",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "minimax-m3",
|
|
system_fingerprint: "fp_ollama",
|
|
choices: [{ index: 0, delta: { content: " is" } }],
|
|
})}\n\n`,
|
|
`data: ${JSON.stringify({
|
|
id: "chatcmpl_ollama",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "minimax-m3",
|
|
system_fingerprint: "fp_ollama",
|
|
choices: [{ index: 0, delta: { content: " the" } }],
|
|
})}\n\n`,
|
|
`data: ${JSON.stringify({
|
|
id: "chatcmpl_ollama",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "minimax-m3",
|
|
system_fingerprint: "fp_ollama",
|
|
choices: [{ index: 0, delta: { content: " response" } }],
|
|
})}\n\n`,
|
|
// Ollama Cloud returns usage with prompt_tokens: 0 — it doesn't count input tokens
|
|
`data: ${JSON.stringify({
|
|
id: "chatcmpl_ollama",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "minimax-m3",
|
|
system_fingerprint: "fp_ollama",
|
|
choices: [],
|
|
usage: { prompt_tokens: 0, completion_tokens: 10, total_tokens: 0 },
|
|
})}\n\n`,
|
|
`data: [DONE]\n\n`,
|
|
],
|
|
{
|
|
mode: "passthrough",
|
|
sourceFormat: FORMATS.OPENAI,
|
|
provider: "ollamacloud",
|
|
model: "minimax-m3",
|
|
body,
|
|
onComplete(payload) {
|
|
onCompletePayload = payload;
|
|
},
|
|
}
|
|
);
|
|
|
|
// The fix should estimate input tokens from the request body instead of showing 0
|
|
assert.equal(
|
|
onCompletePayload.responseBody.usage.completion_tokens,
|
|
10,
|
|
"completion tokens should be unchanged"
|
|
);
|
|
assert.ok(
|
|
onCompletePayload.responseBody.usage.prompt_tokens > 0,
|
|
`prompt_tokens should be estimated (> 0), got ${onCompletePayload.responseBody.usage.prompt_tokens}`
|
|
);
|
|
assert.equal(
|
|
onCompletePayload.responseBody.usage.total_tokens,
|
|
onCompletePayload.responseBody.usage.prompt_tokens + 10,
|
|
"total_tokens should equal prompt_tokens + completion_tokens"
|
|
);
|
|
});
|
|
|
|
test.after(() => {
|
|
core.resetDbInstance();
|
|
if (fs.existsSync(TEST_DATA_DIR)) {
|
|
for (const entry of fs.readdirSync(TEST_DATA_DIR)) {
|
|
fs.rmSync(path.join(TEST_DATA_DIR, entry), { recursive: true, force: true });
|
|
}
|
|
}
|
|
});
|