sk/Perplexica
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(UI): Enhance model statistics tracking and citation handling in chat components

Browse source
This commit is contained in:
Willie Zutz 2025-08-04 00:41:31 -06:00
parent 3e238303b0
commit 71120c997a
11 changed files with 354 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

92
src/lib/search/simplifiedAgent.ts
View file

@ -20,6 +20,43 @@ import { formatDateForLLM } from '../utils';
import { getModelName } from '../utils/modelUtils';
import { removeThinkingBlocks } from '../utils/contentUtils';
/**
* Normalize usage metadata from different LLM providers
*/
function normalizeUsageMetadata(usageData: any): {
input_tokens: number;
output_tokens: number;
total_tokens: number;
} {
if (!usageData) return { input_tokens: 0, output_tokens: 0, total_tokens: 0 };
// Handle different provider formats
const inputTokens =
usageData.input_tokens ||
usageData.prompt_tokens ||
usageData.promptTokens ||
usageData.usedTokens ||
0;
const outputTokens =
usageData.output_tokens ||
usageData.completion_tokens ||
usageData.completionTokens ||
0;
const totalTokens =
usageData.total_tokens ||
usageData.totalTokens ||
usageData.usedTokens ||
inputTokens + outputTokens;
return {
input_tokens: inputTokens,
output_tokens: outputTokens,
total_tokens: totalTokens,
};
}
/**
* Simplified Agent using createReactAgent
*
@ -487,6 +524,11 @@ Use all available tools strategically to provide comprehensive, well-researched,
let finalResult: any = null;
let collectedDocuments: any[] = [];
let currentResponseBuffer = '';
let totalUsage = {
input_tokens: 0,
output_tokens: 0,
total_tokens: 0,
};
// Process the event stream
for await (const event of eventStream) {
@ -540,6 +582,31 @@ Use all available tools strategically to provide comprehensive, well-researched,
// Handle streaming tool calls (for thought messages)
if (event.event === 'on_chat_model_end' && event.data.output) {
const output = event.data.output;
// Collect token usage from chat model end events
if (output.usage_metadata) {
const normalized = normalizeUsageMetadata(output.usage_metadata);
totalUsage.input_tokens += normalized.input_tokens;
totalUsage.output_tokens += normalized.output_tokens;
totalUsage.total_tokens += normalized.total_tokens;
console.log(
'SimplifiedAgent: Collected usage from usage_metadata:',
normalized,
);
} else if (output.response_metadata?.usage) {
// Fallback to response_metadata for different model providers
const normalized = normalizeUsageMetadata(
output.response_metadata.usage,
);
totalUsage.input_tokens += normalized.input_tokens;
totalUsage.output_tokens += normalized.output_tokens;
totalUsage.total_tokens += normalized.total_tokens;
console.log(
'SimplifiedAgent: Collected usage from response_metadata:',
normalized,
);
}
if (
output._getType() === 'ai' &&
output.tool_calls &&
@ -589,6 +656,25 @@ Use all available tools strategically to provide comprehensive, well-researched,
}
}
// Handle LLM end events for token usage tracking
if (event.event === 'on_llm_end' && event.data.output) {
const output = event.data.output;
// Collect token usage from LLM end events
if (output.llmOutput?.tokenUsage) {
const normalized = normalizeUsageMetadata(
output.llmOutput.tokenUsage,
);
totalUsage.input_tokens += normalized.input_tokens;
totalUsage.output_tokens += normalized.output_tokens;
totalUsage.total_tokens += normalized.total_tokens;
console.log(
'SimplifiedAgent: Collected usage from llmOutput:',
normalized,
);
}
}
// Handle token-level streaming for the final response
if (event.event === 'on_chat_model_stream' && event.data.chunk) {
const chunk = event.data.chunk;
@ -676,11 +762,15 @@ Use all available tools strategically to provide comprehensive, well-researched,
// Emit model stats and end signal after streaming is complete
const modelName = getModelName(this.llm);
console.log('SimplifiedAgent: Total usage collected:', totalUsage);
this.emitter.emit(
'stats',
JSON.stringify({
type: 'modelStats',
data: { modelName },
data: {
modelName,
usage: totalUsage,
},
}),
);