feat(assist): сбор usage-токенов и отдача в meta (время этапов уже было)
This commit is contained in:
@@ -272,6 +272,7 @@ function createSemaphore(max) {
|
||||
* model?:string,
|
||||
* temperature?:number,
|
||||
* responseFormat?:{type:string} // напр. {type:'json_object'}
|
||||
* usageRef?:object // мутабельный объект для сбора токенов {prompt,completion,total}
|
||||
* }} params
|
||||
* @returns {Promise<string>} — содержимое choices[0].message.content
|
||||
*/
|
||||
@@ -284,6 +285,7 @@ async function llmChat({
|
||||
model = ENV.LLM_MODEL,
|
||||
temperature,
|
||||
responseFormat,
|
||||
usageRef,
|
||||
}) {
|
||||
const base = String(baseUrl).replace(/\/+$/, '');
|
||||
const url = `${base}/v1/chat/completions`;
|
||||
@@ -349,6 +351,14 @@ async function llmChat({
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
// Собираем токены (usage) для отображения пользователю.
|
||||
if (usageRef && data && data.usage) {
|
||||
usageRef.prompt = (usageRef.prompt || 0) + (data.usage.prompt_tokens || 0);
|
||||
usageRef.completion = (usageRef.completion || 0) + (data.usage.completion_tokens || 0);
|
||||
usageRef.total = (usageRef.total || 0) + (data.usage.total_tokens || 0);
|
||||
}
|
||||
|
||||
const content = data && data.choices && data.choices[0] &&
|
||||
data.choices[0].message && data.choices[0].message.content;
|
||||
|
||||
|
||||
@@ -241,11 +241,13 @@ async function runPipeline(question, stand, vmDocsBaseUrl, controller, start) {
|
||||
// Этап 1 (resource_type).
|
||||
let resourceType = null;
|
||||
let stage1Ok = false;
|
||||
const usage = { prompt: 0, completion: 0, total: 0 };
|
||||
const budget = TOTAL_DEADLINE_MS - (Date.now() - start);
|
||||
if (budget > 0 && resourceTypes.size > 0) {
|
||||
const r1 = await stage1ExtractResourceType(question, {
|
||||
resourceTypesSet: resourceTypes,
|
||||
controller,
|
||||
usageRef: usage,
|
||||
});
|
||||
stage1Ok = r1.ok;
|
||||
if (r1.ok && r1.resourceType) resourceType = r1.resourceType;
|
||||
@@ -267,6 +269,7 @@ async function runPipeline(question, stand, vmDocsBaseUrl, controller, start) {
|
||||
model: ENV.LLM_MODEL,
|
||||
maxTokens: 2048,
|
||||
controller,
|
||||
usageRef: usage,
|
||||
});
|
||||
const tStage2 = Date.now();
|
||||
|
||||
@@ -288,6 +291,9 @@ async function runPipeline(question, stand, vmDocsBaseUrl, controller, start) {
|
||||
stage1_ms: tStage1 - tIndex,
|
||||
search_ms: tSearch - tStage1,
|
||||
stage2_ms: tStage2 - tSearch,
|
||||
tokens: usage.total,
|
||||
prompt_tokens: usage.prompt,
|
||||
completion_tokens: usage.completion,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
@@ -109,6 +109,7 @@ function parseResourceType(content) {
|
||||
* @param {string} [opts.model]
|
||||
* @param {number} [opts.timeoutMs]
|
||||
* @param {AbortController} [opts.controller]
|
||||
* @param {object} [opts.usageRef]
|
||||
* @returns {Promise<{ok:boolean, resourceType?:string|null, error?:string}>}
|
||||
*/
|
||||
async function stage1ExtractResourceType(question, {
|
||||
@@ -118,6 +119,7 @@ async function stage1ExtractResourceType(question, {
|
||||
model = process.env.LLM_MODEL || DEFAULT_MODEL,
|
||||
timeoutMs,
|
||||
controller,
|
||||
usageRef,
|
||||
}) {
|
||||
const types = new Set(
|
||||
[...new Set(resourceTypesSet || [])].map((t) => String(t).toLowerCase().trim())
|
||||
@@ -139,6 +141,7 @@ async function stage1ExtractResourceType(question, {
|
||||
baseUrl,
|
||||
apiKey,
|
||||
model,
|
||||
usageRef,
|
||||
};
|
||||
let timer = null;
|
||||
let ctrl = controller;
|
||||
|
||||
@@ -229,6 +229,7 @@ function hasHeader(text, header) {
|
||||
* @param {string} [params.model]
|
||||
* @param {number} [params.maxTokens]
|
||||
* @param {AbortController} [params.controller]
|
||||
* @param {object} [params.usageRef]
|
||||
* @returns {Promise<{ok:true, answer:string} | {ok:false, error:string}>}
|
||||
*/
|
||||
async function stage2Generate({
|
||||
@@ -239,6 +240,7 @@ async function stage2Generate({
|
||||
model = process.env.LLM_MODEL || DEFAULT_MODEL,
|
||||
maxTokens = 2048,
|
||||
controller,
|
||||
usageRef,
|
||||
}) {
|
||||
const docsBlock = context
|
||||
.map((c, i) => `[${i + 1}] ${c.title}\n${c.text}`)
|
||||
@@ -270,6 +272,7 @@ async function stage2Generate({
|
||||
baseUrl,
|
||||
apiKey,
|
||||
model,
|
||||
usageRef,
|
||||
});
|
||||
} catch (err) {
|
||||
const code = err && err.code ? err.code : 'stage2_failed';
|
||||
|
||||
Reference in New Issue
Block a user