diff --git a/src/lib/brain.ts b/src/lib/brain.ts index 136a10b..b68daa4 100644 --- a/src/lib/brain.ts +++ b/src/lib/brain.ts @@ -13,10 +13,43 @@ import EventEmitter from 'events' EventEmitter.defaultMaxListeners = 1000 // note: need initTools() first -export async function upgradeBrains(url?: string) { +export async function upgradeBrains(flags?: any) { const brains = ToolFunc.get(BRAINS_FUNC_NAME) as LlmModelsFunc - const count = await brains.$refresh({url}) - return count + let _models: AIModelSettings[]|undefined + let shouldBreak: boolean|undefined + + brains.on('brain:refresh', onRefresh) + process.on('SIGINT', interrupted) + + try { + const count = await brains.$refresh(flags) + return count + } finally { + brains.off('brain:refresh', onRefresh) + process.off('SIGINT', interrupted) + if (shouldBreak) {console.log('saved.')} + } + + function onRefresh(tool_name: string, act: string, model: AIModelSettings, models: AIModelSettings[]|string){ + let s = model._id! + if (s.includes(':')) {s = ' ' + s} + if (Array.isArray(models)) { + _models = models + } else if(models && typeof models === 'string') { + s += ' ' + models + } + console.log(act, s) + if (shouldBreak) {this.result = true} + } + + async function interrupted() { + if (_models) { + console.log('wait to exit...') + shouldBreak = true + } else { + process.exit(0) + } + } } export async function listBrains(userConfig: any, flags: any) { @@ -60,6 +93,16 @@ export async function searchBrains(brainDir: string, flags: any) { return result } +export async function verifyBrains(brains: AIModelSettings[]) { + const brainFunc = ToolFunc.get(BRAINS_FUNC_NAME) as LlmModelsFunc + for (const brain of brains) { + const changed = await brainFunc.verifyFileExists(brain) + if (changed) { + brainFunc.put({id: brain._id, val: brain}) + } + } +} + export function printBrains(brains: AIModelSettings[], flags?: {count?: number}) { let maxArrayLength = flags?.count ?? 100 if (maxArrayLength >= 0) {maxArrayLength--} diff --git a/src/oclif/commands/brain/index.ts b/src/oclif/commands/brain/index.ts index 1108d24..5e8c4d6 100644 --- a/src/oclif/commands/brain/index.ts +++ b/src/oclif/commands/brain/index.ts @@ -2,7 +2,7 @@ import { Args, Flags } from '@oclif/core' import { AICommand } from '../../lib/ai-command.js' import { showBanner } from '../../lib/help.js' import { parseJsJson } from '@isdk/ai-tool' -import { listBrains, printBrains, upgradeBrains } from '../../../lib/brain.js' +import { listBrains, printBrains, upgradeBrains, verifyBrains } from '../../../lib/brain.js' export default class Brain extends AICommand { static args = { @@ -34,6 +34,12 @@ export default class Brain extends AICommand { description: 'the hub mirror url', dependsOn: ['refresh'], }), + verifyQuant: Flags.boolean({ + char: 'v', + aliases: ['verify-quant'], + description: 'whether verify quant when refresh', + dependsOn: ['refresh'], + }), } static summary = '🧠 The AI Agent Brains(LLM) Manager.' @@ -61,7 +67,7 @@ export default class Brain extends AICommand { await this.config.runHook('init_tools', {id: 'brain', userConfig}) if (flags.refresh) { - const count = await upgradeBrains(flags.hubUrl) + const count = await upgradeBrains(flags) this.log(`${count} brains updated`) return count } @@ -70,6 +76,7 @@ export default class Brain extends AICommand { flags.name = args.name flags.downloaded = true const result = await listBrains(userConfig, flags) + if (result) { await verifyBrains(result) } if (!isJson) { if (!result || result.length === 0) { this.log('No brains found') diff --git a/src/oclif/commands/brain/list.ts b/src/oclif/commands/brain/list.ts index e2f56c1..92d423f 100644 --- a/src/oclif/commands/brain/list.ts +++ b/src/oclif/commands/brain/list.ts @@ -66,7 +66,7 @@ export default class AIBrainListCommand extends AICommand { await this.config.runHook('init_tools', {id: 'brain', userConfig}) if (flags.refresh) { - const count = await upgradeBrains(flags.hubUrl) + const count = await upgradeBrains(flags) this.log(`${count} brains updated`) }