lisone/ds2api
1
0
Fork 0
mirror of https://github.com/CJackHwang/ds2api.git synced 2026-05-20 16:07:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support explicit prompt token tracking in SSE parsing and stream handlers

Browse Source
This commit is contained in:
CJACK
2026-04-07 01:39:27 +08:00
parent da778a18fb
commit b79a13efd5
13 changed files with 136 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
internal/js/chat-stream/sse_parse_impl.js
View File

@@ -20,7 +20,9 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
};
}
const outputTokens = extractAccumulatedTokenUsage(chunk);
const usage = extractAccumulatedTokenUsage(chunk);
const promptTokens = usage.prompt;
const outputTokens = usage.output;
if (Object.prototype.hasOwnProperty.call(chunk, 'error')) {
return {
@@ -29,7 +31,8 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: true,
contentFilter: false,
errorMessage: formatErrorMessage(chunk.error),
outputTokens: 0,
promptTokens,
outputTokens,
newType: currentType,
};
}
@@ -43,6 +46,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: true,
contentFilter: true,
errorMessage: '',
promptTokens,
outputTokens,
newType: currentType,
};
@@ -55,6 +59,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType: currentType,
};
@@ -67,6 +72,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: true,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType: currentType,
};
@@ -77,6 +83,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType: currentType,
};
@@ -89,6 +96,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType: currentType,
};
@@ -157,6 +165,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: true,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -168,6 +177,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -182,6 +192,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -196,6 +207,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: true,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -207,6 +219,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -242,6 +255,7 @@ function parseChunkForContent(chunk, thinkingEnabled, currentType, stripReferenc
finished: false,
contentFilter: false,
errorMessage: '',
promptTokens,
outputTokens,
newType,
};
@@ -429,40 +443,54 @@ function hasContentFilterStatusValue(v) {
}
function extractAccumulatedTokenUsage(chunk) {
return findAccumulatedTokenUsage(chunk);
const usage = findAccumulatedTokenUsage(chunk);
return usage || { prompt: 0, output: 0 };
}
function findAccumulatedTokenUsage(v) {
if (Array.isArray(v)) {
for (const item of v) {
const n = findAccumulatedTokenUsage(item);
if (n > 0) {
return n;
}
const u = findAccumulatedTokenUsage(item);
if (u) return u;
}
return 0;
return null;
}
if (!v || typeof v !== 'object') {
return 0;
return null;
}
const pathValue = asString(v.p);
if (pathValue && pathValue.toLowerCase().includes('accumulated_token_usage')) {
const n = toInt(v.v);
if (n > 0) {
return n;
return { prompt: 0, output: n };
}
}
if (pathValue && pathValue.toLowerCase().includes('token_usage')) {
const u = v.v;
if (u && typeof u === 'object') {
const p = toInt(u.prompt_tokens);
const c = toInt(u.completion_tokens);
if (p > 0 || c > 0) {
return { prompt: p, output: c };
}
}
}
const direct = toInt(v.accumulated_token_usage);
if (direct > 0) {
return direct;
return { prompt: 0, output: direct };
}
for (const value of Object.values(v)) {
const n = findAccumulatedTokenUsage(value);
if (n > 0) {
return n;
if (v.token_usage && typeof v.token_usage === 'object') {
const p = toInt(v.token_usage.prompt_tokens);
const c = toInt(v.token_usage.completion_tokens);
if (p > 0 || c > 0) {
return { prompt: p, output: c };
}
}
return 0;
for (const value of Object.values(v)) {
const u = findAccumulatedTokenUsage(value);
if (u) return u;
}
return null;
}
function toInt(v) {