dtamind-components
Version:
Apps integration for Dtamind. Contain Nodes and Credentials.
154 lines • 6.86 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.EvaluationRunTracer = void 0;
const run_collector_1 = require("@langchain/core/tracers/run_collector");
const EvaluationRunner_1 = require("./EvaluationRunner");
const tiktoken_1 = require("@dqbd/tiktoken");
class EvaluationRunTracer extends run_collector_1.RunCollectorCallbackHandler {
constructor(id) {
super();
this.countPromptTokens = (encoding, run) => {
let promptTokenCount = 0;
if (encoding) {
if (run.inputs?.messages?.length > 0 && run.inputs?.messages[0]?.length > 0) {
run.inputs.messages[0].map((message) => {
let content = message.content
? message.content
: message.SystemMessage?.content
? message.SystemMessage.content
: message.HumanMessage?.content
? message.HumanMessage.content
: message.AIMessage?.content
? message.AIMessage.content
: undefined;
promptTokenCount += content ? encoding.encode(content).length : 0;
});
}
if (run.inputs?.prompts?.length > 0) {
const content = run.inputs.prompts[0];
promptTokenCount += content ? encoding.encode(content).length : 0;
}
}
return promptTokenCount;
};
this.countCompletionTokens = (encoding, run) => {
let completionTokenCount = 0;
if (encoding) {
if (run.outputs?.generations?.length > 0 && run.outputs?.generations[0]?.length > 0) {
run.outputs?.generations[0].map((chunk) => {
let content = chunk.text ? chunk.text : chunk.message?.content ? chunk.message?.content : undefined;
completionTokenCount += content ? encoding.encode(content).length : 0;
});
}
}
return completionTokenCount;
};
this.extractModelName = (run) => {
return (run?.serialized?.kwargs?.model ||
run?.serialized?.kwargs?.model_name ||
run?.extra?.metadata?.ls_model_name ||
run?.extra?.metadata?.fw_model_name);
};
this.evaluationRunId = id;
}
async persistRun(run) {
return super.persistRun(run);
}
onLLMEnd(run) {
if (run.name) {
let provider = run.name;
if (provider === 'BedrockChat') {
provider = 'awsChatBedrock';
}
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, JSON.stringify({
provider: provider
}));
}
let model = this.extractModelName(run);
if (run.outputs?.llmOutput?.tokenUsage) {
const tokenUsage = run.outputs?.llmOutput?.tokenUsage;
if (tokenUsage) {
const metric = {
completionTokens: tokenUsage.completionTokens,
promptTokens: tokenUsage.promptTokens,
model: model,
totalTokens: tokenUsage.totalTokens
};
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, JSON.stringify(metric));
}
}
else if (run.outputs?.generations?.length > 0 &&
run.outputs?.generations[0].length > 0 &&
run.outputs?.generations[0][0]?.message?.usage_metadata?.total_tokens) {
const usage_metadata = run.outputs?.generations[0][0]?.message?.usage_metadata;
if (usage_metadata) {
const metric = {
completionTokens: usage_metadata.output_tokens,
promptTokens: usage_metadata.input_tokens,
model: model || this.model,
totalTokens: usage_metadata.total_tokens
};
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, JSON.stringify(metric));
}
}
else {
let encoding = undefined;
let promptInputTokens = 0;
let completionTokenCount = 0;
try {
encoding = (0, tiktoken_1.encoding_for_model)(model);
promptInputTokens = this.countPromptTokens(encoding, run);
completionTokenCount = this.countCompletionTokens(encoding, run);
}
catch (e) {
try {
// as tiktoken will fail for non openai models, assume that is 'cl100k_base'
encoding = (0, tiktoken_1.get_encoding)('cl100k_base');
promptInputTokens = this.countPromptTokens(encoding, run);
completionTokenCount = this.countCompletionTokens(encoding, run);
}
catch (e) {
// stay silent
}
}
const metric = {
completionTokens: completionTokenCount,
promptTokens: promptInputTokens,
model: model,
totalTokens: promptInputTokens + completionTokenCount
};
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, JSON.stringify(metric));
//cleanup
this.model = '';
}
}
async onRunUpdate(run) {
const json = {
[run.run_type]: elapsed(run)
};
let metric = JSON.stringify(json);
if (metric) {
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, metric);
}
if (run.run_type === 'llm') {
let model = this.extractModelName(run);
if (model) {
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, JSON.stringify({ model: model }));
this.model = model;
}
// OpenAI non streaming models
const estimatedTokenUsage = run.outputs?.llmOutput?.estimatedTokenUsage;
if (estimatedTokenUsage && typeof estimatedTokenUsage === 'object' && Object.keys(estimatedTokenUsage).length > 0) {
EvaluationRunner_1.EvaluationRunner.addMetrics(this.evaluationRunId, estimatedTokenUsage);
}
}
}
}
exports.EvaluationRunTracer = EvaluationRunTracer;
function elapsed(run) {
if (!run.end_time)
return '';
const elapsed = run.end_time - run.start_time;
return `${elapsed.toFixed(2)}`;
}
//# sourceMappingURL=EvaluationRunTracer.js.map