refactor: upgradeBrains

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--}

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')

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`)
     }