magnitude-core
Version:
Magnitude e2e testing agent
190 lines (189 loc) • 10 kB
JavaScript
import { convertToBamlClientOptions } from "./util";
// Import ModularMemoryContext instead of old MemoryContext
import { b } from "@/ai/baml_client";
import { Collector, ClientRegistry } from "@boundaryml/baml";
import logger from "@/logger";
import TypeBuilder from "./baml_client/type_builder";
import { z } from 'zod';
import { convertActionDefinitionsToBaml, convertZodToBaml } from "@/actions/util";
import EventEmitter from "eventemitter3";
export class ModelHarness {
/**
* Strong reasoning agent for high level strategy and planning.
*/
events = new EventEmitter();
options;
collector;
cr;
baml;
logger;
prevTotalInputTokens = 0;
prevTotalOutputTokens = 0;
constructor(options) {
this.options = {
llm: options.llm,
//promptCaching: options.promptCaching ?? false
};
this.logger = logger.child({ name: 'llm' });
}
async setup() {
// Must be called after constructor
this.collector = new Collector("macro");
this.cr = new ClientRegistry();
let bamlClientOptions = await convertToBamlClientOptions(this.options.llm);
this.cr.addLlmClient('Magnus', this.options.llm.provider === 'claude-code' ? 'anthropic' : this.options.llm.provider, bamlClientOptions, 'DefaultRetryPolicy');
this.cr.setPrimary('Magnus');
this.baml = b.withOptions({ collector: this.collector, clientRegistry: this.cr });
}
describeModel() {
return `${this.options.llm.provider}:${'model' in this.options.llm.options ? this.options.llm.options.model : 'unknown'}`;
}
_reportUsage() {
// console.log('this.collector.last', this.collector.last)
// if (this.collector.last) console.log("calls:", this.collector.last.calls)//console.log("Response: ", this.collector.last.calls[-1].httpResponse);
//console.log('last call:', this.collector.last?.calls.at(-1)?.httpResponse?.body.json());
// Get tokens used since last call to reportUsage
//console.log(this.collector.usage);
let inputTokens = 0;
let outputTokens = 0;
let cacheWriteInputTokens = 0;
let cacheReadInputTokens = 0;
if (this.options.llm.provider === 'anthropic' || this.options.llm.provider === 'claude-code') {
const usage = this.collector.last?.calls.at(-1)?.httpResponse?.body.json().usage;
//console.log("Usage from Anthropic:", usage);
if (!usage) {
// Sometimes apparently this happens? Happened once after extract for example
logger.warn("No usage returned from Anthropic provider, cached cost may be inaccurate");
inputTokens = (this.collector.usage.inputTokens ?? 0) - this.prevTotalInputTokens;
outputTokens = (this.collector.usage.outputTokens ?? 0) - this.prevTotalOutputTokens;
}
else {
inputTokens = usage.input_tokens;
outputTokens = usage.output_tokens;
cacheWriteInputTokens = usage.cache_creation_input_tokens;
cacheReadInputTokens = usage.cache_read_input_tokens;
}
}
else {
inputTokens = (this.collector.usage.inputTokens ?? 0) - this.prevTotalInputTokens;
outputTokens = (this.collector.usage.outputTokens ?? 0) - this.prevTotalOutputTokens;
}
const model = this.options.llm.options.model ?? 'unknown';
// Get cost if known
const knownCostMap = {
// TODO: track cached savings on Gemini
'gemini-2.5-pro': { inputTokens: 1.25, outputTokens: 10.0 },
'gemini-2.5-flash': { inputTokens: 0.30, outputTokens: 2.50 },
'gemini-2.5-flash-lite': { inputTokens: 0.10, outputTokens: 0.40 },
'claude-3.5-sonnet': { inputTokens: 3.00, outputTokens: 15.00, cacheWriteInputTokens: 3.75, cacheReadInputTokens: 0.30 },
'claude-3.7-sonnet': { inputTokens: 3.00, outputTokens: 15.00, cacheWriteInputTokens: 3.75, cacheReadInputTokens: 0.30 },
'claude-sonnet-4': { inputTokens: 3.00, outputTokens: 15.00, cacheWriteInputTokens: 3.75, cacheReadInputTokens: 0.30 },
'claude-haiku-4-5': { inputTokens: 1.00, outputTokens: 5.00, cacheWriteInputTokens: 1.25, cacheReadInputTokens: 0.10 },
'claude-opus-4': { inputTokens: 15.00, outputTokens: 75.00, cacheWriteInputTokens: 18.75, cacheReadInputTokens: 1.50 },
'gpt-4.1': { inputTokens: 2.00, outputTokens: 8.00 },
'gpt-4.1-mini': { inputTokens: 0.40, outputTokens: 1.60 },
'gpt-4.1-nano': { inputTokens: 0.10, outputTokens: 0.40 },
'gpt-4o': { inputTokens: 3.75, outputTokens: 15.00 },
// Assuming Nebius prices, may be higher
'qwen2.5-vl-72b': { inputTokens: 0.25, outputTokens: 0.75 }
};
let inputTokenCost;
let outputTokenCost;
let cacheWriteInputTokenCost;
let cacheReadInputTokenCost;
for (const [name, costs] of Object.entries(knownCostMap)) {
if (model.includes(name)) {
inputTokenCost = costs.inputTokens / 1_000_000;
outputTokenCost = costs.outputTokens / 1_000_000;
cacheReadInputTokenCost = costs.cacheReadInputTokens ? costs.cacheReadInputTokens / 1_000_000 : undefined;
cacheWriteInputTokenCost = costs.cacheWriteInputTokens ? costs.cacheWriteInputTokens / 1_000_000 : undefined;
}
}
// console.log("cacheWriteInputTokenCost:", cacheWriteInputTokenCost);
// console.log("cacheWriteInputTokens:", cacheWriteInputTokens);
const usage = {
llm: {
provider: this.options.llm.provider,
model: model
}, //this.options.llm,
inputTokens: inputTokens,
outputTokens: outputTokens,
...(cacheWriteInputTokens ? { cacheWriteInputTokens } : {}),
...(cacheReadInputTokens ? { cacheReadInputTokens } : {}),
...(inputTokenCost ? {
inputCost: inputTokens * inputTokenCost +
(cacheWriteInputTokenCost ? cacheWriteInputTokenCost * cacheWriteInputTokens : 0.0) +
(cacheReadInputTokenCost ? cacheReadInputTokenCost * cacheReadInputTokens : 0.0)
} : {}),
...(outputTokenCost ? { outputCost: outputTokens * outputTokenCost } : {}),
// ...(cacheWriteInputTokenCost ? { : inputTokens * inputTokenCost } : {}),
// ...(cacheReadInputTokenCost ? { outputCost: outputTokens * outputTokenCost } : {})
};
this.events.emit('tokensUsed', usage);
//console.log("Usage:", usage);
this.prevTotalInputTokens += inputTokens;
this.prevTotalOutputTokens += outputTokens;
}
async partialAct(context, // Changed to ModularMemoryContext
task, data, actionVocabulary) {
const tb = new TypeBuilder();
tb.PartialRecipe.addProperty('actions', tb.list(convertActionDefinitionsToBaml(tb, actionVocabulary))).description('Always provide at least one action');
const start = Date.now();
// Assuming this.baml.CreatePartialRecipe is now typed to accept ModularMemoryContext
// after BAML generation picked up changes in planner.baml
const response = await this.baml.CreatePartialRecipe(context, task, data, this.options.llm.provider === 'claude-code', { tb });
this.logger.trace(`createPartialRecipe took ${Date.now() - start}ms`);
// BAML does not carry over action type to @@dynamic of PartialRecipe, so forced cast necssary
//return response as unknown as { actions: z.infer<ActionDefinition<T>['schema']>[] };//, finished: boolean };
this._reportUsage();
return {
reasoning: response.reasoning, //(response.observations ? response.observations + " " : "") + response.meta_reasoning + " " + response.reasoning,
actions: response.actions // as z.infer<ActionDefinition<T>['schema']>[]
};
}
async extract(instructions, schema, screenshot, domContent) {
const tb = new TypeBuilder();
if (schema instanceof z.ZodObject) {
// populate ExtractedData with schema KVs instead of wrapping in data key unnecessarily
for (const [key, fieldSchema] of Object.entries(schema.shape)) {
tb.ExtractedData.addProperty(key, convertZodToBaml(tb, fieldSchema));
}
}
else {
// for array or primitive have to wrap data key
tb.ExtractedData.addProperty('data', convertZodToBaml(tb, schema));
}
// } else if (schema instanceof z.ZodArray) {
// }
const resp = await this.baml.ExtractData(instructions, await screenshot.toBaml(), domContent, this.options.llm.provider === 'claude-code', { tb });
this._reportUsage();
if (schema instanceof z.ZodObject) {
return resp;
}
else {
return resp.data;
}
}
// ^ extract could prob be a subset of query w trimmed mem
async query(context, query, schema) {
const tb = new TypeBuilder();
if (schema instanceof z.ZodObject) {
// populate ExtractedData with schema KVs instead of wrapping in data key unnecessarily
for (const [key, fieldSchema] of Object.entries(schema.shape)) {
tb.QueryResponse.addProperty(key, convertZodToBaml(tb, fieldSchema));
}
}
else {
// for array or primitive have to wrap data key
tb.QueryResponse.addProperty('data', convertZodToBaml(tb, schema));
}
const resp = await this.baml.QueryMemory(context, query, this.options.llm.provider === 'claude-code', { tb });
this._reportUsage();
if (schema instanceof z.ZodObject) {
return resp;
}
else {
return resp.data;
}
}
}