feat: enhance token usage tracking and history management
- Updated HTML assets for improved loading. - Integrated token usage tracking in chat processing, appending usage details to transcripts. - Enhanced OpenAIProvider to include usage data in chat completion responses. - Implemented asynchronous retrieval of recent token usage history. - Added utility functions for managing transcript files and parsing usage data. - Updated UI components to reflect changes in usage status handling. - Ensured consistent usage status definitions across the application.
This commit is contained in:
@@ -6,6 +6,7 @@ import logManager from '@electron/service/logger';
|
||||
import type { RawMessage } from '@src/pages/home/model/ChatModel';
|
||||
import { sessionStore } from '../session-store';
|
||||
import type { GatewayEvent, GatewayRpcParams, GatewayRpcReturns } from '../types';
|
||||
import { appendTranscriptLine } from '@electron/utils/token-usage-writer';
|
||||
|
||||
export interface GatewayChatMessage {
|
||||
role: 'system' | 'user' | 'assistant' | 'tool';
|
||||
@@ -34,11 +35,13 @@ async function processChatStream(
|
||||
runId: string,
|
||||
provider: BaseProvider,
|
||||
model: string,
|
||||
providerName: string,
|
||||
messages: GatewayChatMessage[],
|
||||
signal: AbortSignal,
|
||||
broadcast: (event: GatewayEvent) => void
|
||||
) {
|
||||
let assistantContent = '';
|
||||
let finalUsage: any = undefined;
|
||||
|
||||
try {
|
||||
const chunks = await provider.chat(messages, model, { signal });
|
||||
@@ -56,9 +59,12 @@ async function processChatStream(
|
||||
});
|
||||
}
|
||||
|
||||
if (chunk.isEnd) {
|
||||
break;
|
||||
if (chunk.usage !== undefined) {
|
||||
finalUsage = chunk.usage;
|
||||
}
|
||||
|
||||
// Do not break on isEnd; the iterable may still yield a trailing usage chunk.
|
||||
// The loop will finish naturally when the generator is done.
|
||||
}
|
||||
|
||||
if (!signal.aborted) {
|
||||
@@ -70,6 +76,18 @@ async function processChatStream(
|
||||
sessionStore.appendMessage(sessionKey, finalMessage);
|
||||
sessionStore.clearActiveRun(sessionKey);
|
||||
|
||||
appendTranscriptLine(sessionKey, {
|
||||
type: 'message',
|
||||
timestamp: new Date().toISOString(),
|
||||
message: {
|
||||
role: 'assistant',
|
||||
content: assistantContent,
|
||||
model,
|
||||
provider: providerName,
|
||||
usage: finalUsage,
|
||||
},
|
||||
});
|
||||
|
||||
broadcast({
|
||||
type: 'chat:final',
|
||||
sessionKey,
|
||||
@@ -96,9 +114,19 @@ export function handleChatSend(
|
||||
const runId = randomUUID();
|
||||
|
||||
// 1. Append user message
|
||||
sessionStore.appendMessage(sessionKey, {
|
||||
const userMessage: RawMessage = {
|
||||
...message,
|
||||
timestamp: message.timestamp || Date.now(),
|
||||
};
|
||||
sessionStore.appendMessage(sessionKey, userMessage);
|
||||
|
||||
appendTranscriptLine(sessionKey, {
|
||||
type: 'message',
|
||||
timestamp: new Date().toISOString(),
|
||||
message: {
|
||||
role: 'user',
|
||||
content: typeof userMessage.content === 'string' ? userMessage.content : '',
|
||||
},
|
||||
});
|
||||
|
||||
// 2. Resolve provider account
|
||||
@@ -127,7 +155,8 @@ export function handleChatSend(
|
||||
|
||||
// Run async stream processing in background
|
||||
const provider = createProvider(accountId);
|
||||
processChatStream(sessionKey, runId, provider, model, messages, abortController.signal, broadcast).catch(
|
||||
const providerName = account.vendorId || account.label || account.model || 'unknown';
|
||||
processChatStream(sessionKey, runId, provider, model, providerName, messages, abortController.signal, broadcast).catch(
|
||||
(err) => {
|
||||
logManager.error('Unexpected error in processChatStream:', err);
|
||||
sessionStore.clearActiveRun(sessionKey);
|
||||
|
||||
Reference in New Issue
Block a user