magnitude-core
Version:
Magnitude e2e testing agent
123 lines (122 loc) • 5.43 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 = options;
//this.llm = options.llm
this.collector = new Collector("macro");
this.cr = new ClientRegistry();
const client = this.options.llm;
let bamlClientOptions = convertToBamlClientOptions(this.options.llm);
this.cr.addLlmClient('Macro', client.provider, bamlClientOptions, 'DefaultRetryPolicy');
this.cr.setPrimary('Macro');
this.baml = b.withOptions({ collector: this.collector, clientRegistry: this.cr });
this.logger = logger.child({ name: 'llm' });
}
// getInfo(): ModelUsage {
// return {
// provider: this.options.client.provider,
// model: 'model' in this.options.client.options ?
// this.options.client.options.model : 'unknown',
// inputTokens: this.collector.usage.inputTokens ?? 0,
// outputTokens: this.collector.usage.outputTokens ?? 0,
// numCalls: this.collector.logs.length
// }
// }
reportUsage() {
// Get tokens used since last call to reportUsage
const inputTokens = (this.collector.usage.inputTokens ?? 0) - this.prevTotalInputTokens;
const outputTokens = (this.collector.usage.outputTokens ?? 0) - this.prevTotalOutputTokens;
const usage = {
llm: {
provider: this.options.llm.provider,
model: this.options.llm.options.model ?? 'unknown'
}, //this.options.llm,
inputTokens: inputTokens,
outputTokens: outputTokens
};
this.events.emit('tokensUsed', usage);
this.prevTotalInputTokens = inputTokens;
this.prevTotalOutputTokens = outputTokens;
}
async createPartialRecipe(context, // Changed to ModularMemoryContext
task, actionVocabulary) {
const tb = new TypeBuilder();
tb.PartialRecipe.addProperty('actions', tb.list(convertActionDefinitionsToBaml(tb, actionVocabulary)));
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, { 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, { 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, { tb });
this.reportUsage();
if (schema instanceof z.ZodObject) {
return resp;
}
else {
return resp.data;
}
}
}