Ollama tools

[RicardoBassete]

Open Ollama tools in Script Kit

export const metadata: Metadata = {
  name: 'Ollama tools',
  description: 'A script for managing downloading, stopping, removing and listing models with ollama',
  author: 'Ricardo Gonçalves Bassete',
}

import '@johnlindquist/kit'
import { load } from 'cheerio'

async function getRunningModels() {
  const { stdout } = await exec('ollama ps')

  const models = stdout
    .split('\n')
    .filter(line => line.trim() !== '')
    .map(line => line.split(' ')[0])

  models.splice(0, 1)

  return models
}

async function getDownloadedModels() {
  const { stdout } = await exec('ollama list')

  const models = stdout
    .split('\n')
    .filter(line => line.trim() !== '')
    .map(line => line.split(' ')[0])

  models.splice(0, 1)

  return models
}

async function getAvaliableModels() {
  const html = await fetch('https://ollama.com/library').then(res => res.text())
  const $ = load(html)

  const models = $('li')
    .get()
    .map(li => {
      const name = $(li).find('h2').text().trim()
      const description = $(li).find('p.max-w-lg.break-words.text-neutral-800.text-md').text().trim()
      const tags = $(li)
        .find('div.flex.flex-wrap.space-x-2')
        .find('span')
        .get()
        .map(span => $(span).text())
      return {
        name,
        description,
        tags,
      }
    })

  return models
}

async function getAllModelsOf(model: string) {
  const html = await fetch(`https://ollama.com/library/${model}/tags`).then(res => res.text())
  const $ = load(html)

  const models = $('div.group.px-4.py-3')
    .get()
    .map(div => {
      const name = $(div).find('span.flex.items-center.font-medium.col-span-6.group.text-sm').text().trim().split(' ')[0]
      const [size, context] = $(div)
        .find('p')
        .get()
        .map(p => $(p).text())
      const input = $(div).find('div.col-span-2').text()

      return { name, size, context, input }
    })

  return models
}

const action = await arg({
  placeholder: 'What do you want to do?',
  choices: [
    { name: '[L]ist', value: 'List' },
    { name: '[D]ownload', value: 'Download' },
    { name: '[S]top', value: 'Stop' },
    { name: '[R]emove', value: 'Remove' },
  ],
})

if (action == 'List') {
  const downloadedModels = await getDownloadedModels()
  await div(`
      <div class="p-5 prose prose-sm">
        <h1>Downloaded Models</h1>
        <ul class="pl-5">
        ${downloadedModels.map(model => `<li>${model}</li>`).join('')}
        </ul>
      </div>
    `)
}

if (action == 'Download') {
  const avaliableModels = await getAvaliableModels()

  const targetModelFamily = await arg(
    'Select model family',
    avaliableModels.map(model => {
      return {
        name: model.name,
        value: model.name,
        description: model.description,
        tag: model.tags.join('-'),
      }
    })
  )

  const models = await getAllModelsOf(targetModelFamily)

  const modelToDownload = await arg(
    'Select model to download',
    models.map(m => {
      return {
        name: m.name,
        value: m.name,
        tag: `Size: ${m.size} - Context: ${m.context} - Input: ${m.input}`,
      }
    })
  )

  await term(`ollama pull ${modelToDownload}`)
}

if (action == 'Stop') {
  const runningModels = await getRunningModels()

  const modelsToStop = await select<string[]>('Select the models to stop', runningModels)

  for (const model of modelsToStop) {
    await exec(`ollama stop ${model}`)
  }
}

if (action == 'Remove') {
  const models = await getDownloadedModels()

  const modelsToRemove = await select<string[]>('Select the models you wish to remove', models)

  for (const model of modelsToRemove) {
    await exec(`ollama rm ${model}`)
  }
}