magnitude-core
Version:
Magnitude e2e testing agent
53 lines (52 loc) • 2.06 kB
JavaScript
import { ModelHarness } from "./modelHarness";
import { allBrowserAgentRoles } from "./types";
import EventEmitter from "eventemitter3";
export class MultiModelHarness {
/**
* Delegates model responsibilites to different LLMs and consolidates their usage
*/
// Roles may reference the same harness
roles = {};
uniqueModels = [];
events = new EventEmitter();
constructor(clients) {
// Sort by specificity (from least specific to most specific)
const sortedClients = clients.toSorted((a, b) => (b.roles ? b.roles.length : 9999) - (a.roles ? a.roles.length : 9999));
for (const client of sortedClients) {
const harness = new ModelHarness({ llm: client });
this.uniqueModels.push(harness);
if (client.roles) {
for (const role of client.roles) {
this.roles[role] = harness;
}
}
else {
for (const role of allBrowserAgentRoles) {
this.roles[role] = harness;
}
}
// Forward token usage events upward
harness.events.on('tokensUsed', (usage) => { this.events.emit('tokensUsed', usage); }, this);
}
}
async setup() {
await Promise.all(this.uniqueModels.map(model => model.setup()));
}
describe() {
// for now - describe least specific model
return this.uniqueModels[0].describeModel();
}
// TODO: generalize responsibility delegation
async partialAct(context, task, data, actionVocabulary) {
return await this.roles['act'].partialAct(context, task, data, actionVocabulary);
}
async extract(instructions, schema, screenshot, domContent) {
return await this.roles['extract'].extract(instructions, schema, screenshot, domContent);
}
async query(context, query, schema) {
return await this.roles['query'].query(context, query, schema);
}
get numUniqueModels() {
return this.uniqueModels.length;
}
}