287 lines
8.2 KiB
JavaScript
Executable File
287 lines
8.2 KiB
JavaScript
Executable File
#!/usr/bin/env node
|
|
import { readFile, writeFile } from "node:fs/promises"
|
|
import path from "node:path"
|
|
|
|
const paths = {
|
|
source: path.resolve("ee/apps/inference/src/models/base.json"),
|
|
models: path.resolve("ee/apps/inference/src/models/openwork-models.json"),
|
|
aliases: path.resolve("packages/types/src/den/inference.ts"),
|
|
}
|
|
|
|
function usage() {
|
|
console.log(`Usage:
|
|
openwork-models.mjs search <query>
|
|
openwork-models.mjs add <query-or-id>
|
|
openwork-models.mjs remove <id> [<id2>]
|
|
openwork-models.mjs discount <usageFactor> <id> [<id2>]
|
|
openwork-models.mjs sync
|
|
openwork-models.mjs validate`)
|
|
}
|
|
|
|
async function readJson(filePath) {
|
|
return JSON.parse(await readFile(filePath, "utf8"))
|
|
}
|
|
|
|
async function writeJson(filePath, value) {
|
|
await writeFile(filePath, `${JSON.stringify(value, null, 2)}\n`)
|
|
}
|
|
|
|
function isRecord(value) {
|
|
return Boolean(value) && typeof value === "object" && !Array.isArray(value)
|
|
}
|
|
|
|
function normalize(value) {
|
|
return String(value).toLowerCase().replace(/[^a-z0-9]+/g, "")
|
|
}
|
|
|
|
function terms(value) {
|
|
return String(value)
|
|
.toLowerCase()
|
|
.split(/[^a-z0-9.]+/)
|
|
.map((term) => normalize(term))
|
|
.filter(Boolean)
|
|
}
|
|
|
|
function modelText(id, model) {
|
|
return `${id} ${model.id ?? ""} ${model.name ?? ""} ${model.family ?? ""}`
|
|
}
|
|
|
|
function scoreModel(query, id, model) {
|
|
const compactQuery = normalize(query)
|
|
const compactText = normalize(modelText(id, model))
|
|
if (!compactQuery) return 0
|
|
if (id === query || model.id === query) return 1_000
|
|
if (compactText.includes(compactQuery)) return 200 + compactQuery.length
|
|
|
|
const queryTerms = terms(query)
|
|
let score = 0
|
|
for (const term of queryTerms) {
|
|
if (compactText.includes(term)) score += 20 + term.length
|
|
}
|
|
return score
|
|
}
|
|
|
|
function displayMatches(matches) {
|
|
for (const match of matches.slice(0, 12)) {
|
|
console.log(`${match.id}\t${match.model.name ?? ""}`)
|
|
}
|
|
}
|
|
|
|
async function readSourceModels() {
|
|
const base = await readJson(paths.source)
|
|
const sourceModels = base?.openrouter?.models
|
|
if (!isRecord(sourceModels)) {
|
|
throw new Error(`${paths.source} did not contain openrouter.models`)
|
|
}
|
|
return sourceModels
|
|
}
|
|
|
|
async function readOpenworkModels() {
|
|
const models = await readJson(paths.models)
|
|
if (!isRecord(models)) {
|
|
throw new Error(`${paths.models} must be a JSON object keyed by model id`)
|
|
}
|
|
return models
|
|
}
|
|
|
|
function findMatches(sourceModels, query) {
|
|
return Object.entries(sourceModels)
|
|
.map(([id, model]) => ({ id, model, score: scoreModel(query, id, model) }))
|
|
.filter((match) => match.score > 0)
|
|
.sort((left, right) => right.score - left.score || left.id.localeCompare(right.id))
|
|
}
|
|
|
|
function resolveSingleSourceModel(sourceModels, query) {
|
|
if (sourceModels[query]) {
|
|
return { id: query, model: sourceModels[query] }
|
|
}
|
|
|
|
const matches = findMatches(sourceModels, query)
|
|
if (matches.length === 0) {
|
|
console.error(`No source model matched "${query}".`)
|
|
process.exitCode = 2
|
|
return null
|
|
}
|
|
|
|
const [first, second] = matches
|
|
if (first && first.score >= 40 && (!second || first.score - second.score >= 10)) {
|
|
return { id: first.id, model: first.model }
|
|
}
|
|
|
|
console.error(`Multiple source models matched "${query}". Please clarify with exact id:`)
|
|
displayMatches(matches)
|
|
process.exitCode = 2
|
|
return null
|
|
}
|
|
|
|
function resolveExactOpenworkIds(openworkModels, ids) {
|
|
const resolved = []
|
|
for (const id of ids) {
|
|
if (openworkModels[id]) {
|
|
resolved.push(id)
|
|
continue
|
|
}
|
|
|
|
const matches = findMatches(openworkModels, id)
|
|
console.error(`OpenWork model "${id}" was not exact. Please clarify with exact id:`)
|
|
displayMatches(matches)
|
|
process.exitCode = 2
|
|
return null
|
|
}
|
|
return resolved
|
|
}
|
|
|
|
function escapeTsString(value) {
|
|
return JSON.stringify(value)
|
|
}
|
|
|
|
async function readUsageFactors() {
|
|
const source = await readFile(paths.aliases, "utf8")
|
|
const factors = new Map()
|
|
const pattern = /"([^"]+)":\s*\{[\s\S]*?usageFactor:\s*([0-9.]+)/g
|
|
for (const match of source.matchAll(pattern)) {
|
|
factors.set(match[1], Number(match[2]))
|
|
}
|
|
return factors
|
|
}
|
|
|
|
async function syncAliases(models, usageOverrides = new Map()) {
|
|
const source = await readFile(paths.aliases, "utf8")
|
|
const existingFactors = await readUsageFactors()
|
|
const entries = Object.entries(models).map(([id, model]) => {
|
|
const usageFactor = usageOverrides.get(id) ?? existingFactors.get(id) ?? 1
|
|
const displayName = `OpenWork: ${model.name}`
|
|
return ` ${escapeTsString(id)}: {
|
|
upstreamModel: ${escapeTsString(id)},
|
|
displayName: ${escapeTsString(displayName)},
|
|
enabled: true,
|
|
usageFactor: ${usageFactor},
|
|
},`
|
|
})
|
|
|
|
const replacement = `export const INFERENCE_MODEL_ALIASES = {\n${entries.join("\n")}\n} as const;`
|
|
const aliasesPattern = /export const INFERENCE_MODEL_ALIASES = \{[\s\S]*?\n\} as const;/
|
|
if (!aliasesPattern.test(source)) {
|
|
throw new Error(`Could not find INFERENCE_MODEL_ALIASES in ${paths.aliases}`)
|
|
}
|
|
|
|
const nextSource = source.replace(aliasesPattern, replacement)
|
|
|
|
await writeFile(paths.aliases, nextSource)
|
|
}
|
|
|
|
async function syncAll(models, usageOverrides = new Map()) {
|
|
await writeJson(paths.models, models)
|
|
await syncAliases(models, usageOverrides)
|
|
await validate()
|
|
}
|
|
|
|
async function validate() {
|
|
const models = await readOpenworkModels()
|
|
for (const [id, model] of Object.entries(models)) {
|
|
if (!isRecord(model)) throw new Error(`${id} must be an object`)
|
|
if (model.id !== id) throw new Error(`${id} must have matching id field`)
|
|
if (typeof model.name !== "string" || model.name.length === 0) {
|
|
throw new Error(`${id} must have a name`)
|
|
}
|
|
if (typeof model.family !== "string" || model.family.length === 0) {
|
|
throw new Error(`${id} must have a family`)
|
|
}
|
|
if (!isRecord(model.cost)) {
|
|
throw new Error(`${id} must have a cost object`)
|
|
}
|
|
}
|
|
|
|
const aliasSource = await readFile(paths.aliases, "utf8")
|
|
for (const id of Object.keys(models)) {
|
|
if (!aliasSource.includes(`upstreamModel: ${escapeTsString(id)}`)) {
|
|
throw new Error(`${id} is missing from INFERENCE_MODEL_ALIASES`)
|
|
}
|
|
}
|
|
|
|
const aliasKeys = [...aliasSource.matchAll(/^\s*"([^"]+)": \{/gm)].map((match) => match[1])
|
|
for (const aliasKey of aliasKeys) {
|
|
if (!models[aliasKey]) {
|
|
throw new Error(`${aliasKey} exists in INFERENCE_MODEL_ALIASES but not ${paths.models}`)
|
|
}
|
|
}
|
|
|
|
console.log(`validated ${Object.keys(models).length} OpenWork model(s)`)
|
|
}
|
|
|
|
async function search(query) {
|
|
const sourceModels = await readSourceModels()
|
|
const matches = findMatches(sourceModels, query)
|
|
if (matches.length === 0) {
|
|
console.log(`No source models matched "${query}".`)
|
|
return
|
|
}
|
|
displayMatches(matches)
|
|
}
|
|
|
|
async function add(query) {
|
|
const sourceModels = await readSourceModels()
|
|
const selected = resolveSingleSourceModel(sourceModels, query)
|
|
if (!selected) return
|
|
|
|
const models = await readOpenworkModels()
|
|
models[selected.id] = selected.model
|
|
await syncAll(models)
|
|
console.log(`added ${selected.id}`)
|
|
}
|
|
|
|
async function remove(ids) {
|
|
const models = await readOpenworkModels()
|
|
const resolved = resolveExactOpenworkIds(models, ids)
|
|
if (!resolved) return
|
|
|
|
for (const id of resolved) {
|
|
delete models[id]
|
|
}
|
|
await syncAll(models)
|
|
console.log(`removed ${resolved.join(", ")}`)
|
|
}
|
|
|
|
async function discount(factorText, ids) {
|
|
const factor = Number(factorText)
|
|
if (!Number.isFinite(factor) || factor <= 0) {
|
|
throw new Error("usageFactor must be a positive number")
|
|
}
|
|
|
|
const models = await readOpenworkModels()
|
|
const resolved = resolveExactOpenworkIds(models, ids)
|
|
if (!resolved) return
|
|
|
|
const currentFactors = await readUsageFactors()
|
|
for (const id of resolved) {
|
|
currentFactors.set(id, factor)
|
|
}
|
|
|
|
await syncAll(models, currentFactors)
|
|
console.log(`set usageFactor=${factor} for ${resolved.join(", ")}`)
|
|
}
|
|
|
|
const [command, ...args] = process.argv.slice(2)
|
|
|
|
if (!command) {
|
|
usage()
|
|
process.exit(1)
|
|
}
|
|
|
|
if (command === "search") {
|
|
await search(args.join(" "))
|
|
} else if (command === "add") {
|
|
await add(args.join(" "))
|
|
} else if (command === "remove") {
|
|
await remove(args)
|
|
} else if (command === "discount") {
|
|
await discount(args[0], args.slice(1))
|
|
} else if (command === "sync") {
|
|
await syncAll(await readOpenworkModels())
|
|
} else if (command === "validate") {
|
|
await validate()
|
|
} else {
|
|
usage()
|
|
process.exit(1)
|
|
}
|