ollama runを追え

ollamaは各種LLMをローカルで実行できるCLIツールで、以下のコマンドが用意されています。

Available Commands:
  serve       Start ollama
  create      Create a model from a Modelfile
  show        Show information for a model
  run         Run a model
  pull        Pull a model from a registry
  push        Push a model to a registry
  list        List models
  ps          List running models
  cp          Copy a model
  rm          Remove a model
  help        Help about any command

ここでは、LLMモデルを実行するrunコマンドについて、 ollama run llama3.1とした場合の実行パスをトレースしてみようと思います。

概要

ここではollamaのv0.3.5を使います。 ollamaではCLI作成ライブラリとしてcobraを利用しています。cmd/cmd.goに各コマンドが記述されていて、NewCLIメソッドでrunCmdが作成されます。runCmd実行前（PreRunE）には、checkServerHeartbeatメソッドでサーバー起動を確認、起動してなかった場合はエラー終了します。RunHandlerメソッドが処理本体になっていて、入力されたプロンプトに対して、LLMが生成した回答をコンソール出力するまでのパスを見ていきます。

RunHandlerメソッドでプロンプト入力からLLMが回答を生成するまで

RunHandlerはinteractiveかnot interactiveのどちらかで実行されます。ollama run lllama3.1の場合はinteractiveがtrueで、ollama run llama3.1 helloのようにモデル名のあとにプロンプトを与えた場合にはinteractiveがfalseとなります。

RunHandlerの最初にinteractive := trueと設定していて、プロンプトを与えた場合には以下のコードでプロンプトがあるかを判定して、interactive = falseと設定し直しています。args[1:]がモデル名以降の引数で、promptsとして保持しています。

prompts := args[1:]
// prepend stdin to the prompt if provided
if !term.IsTerminal(int(os.Stdin.Fd())) {
    in, err := io.ReadAll(os.Stdin)
    if err != nil {
        return err
    }

    prompts = append([]string{string(in)}, prompts...)
    opts.WordWrap = false
    interactive = false
}
opts.Prompt = strings.Join(prompts, " ")
if len(prompts) > 0 {
    interactive = false
}

指定したモデルがpullされてなかった場合には、次にpullが行われます。nameがモデル名で、clientのShowメソッドでモデルに関する情報を取得します（ollama show llama3.1とした場合と同じ）。モデルがなければse.StatusCodeが404となるので、PullHandlerでモデルをpullする処理が実行されます。

name := args[0]
info, err := func() (*api.ShowResponse, error) {
    showReq := &api.ShowRequest{Name: name}
    info, err := client.Show(cmd.Context(), showReq)
    var se api.StatusError
    if errors.As(err, &se) && se.StatusCode == http.StatusNotFound {
        if err := PullHandler(cmd, []string{name}); err != nil {
            return nil, err
        }
        return client.Show(cmd.Context(), &api.ShowRequest{Name: name})
    }
    return info, err
}()
if err != nil {
    return err
}

RunHandlerの最後で、interactiveの場合にはgenerateInteractive、not interactiveの場合にはgenerateでLLMから生成された回答を出力します。ここでは、interactiveの動作を追いかけるので、generateInteractiveを見ていきます。

generateInteractiveでは、ターミナルからの入力に対して、/listや/loadなどのコマンドかプロンプトかをswitch caseで判定します。その後、プロンプトの場合には、LLMに対して送信するメッセージを作成し、LLMから生成された回答を受け取ってターミナルに出力します。これはchatメソッドの中で実行されます。また、作成したメッセージと回答を履歴として保持したメッセージに更新して、次の入力を待ちます。

if sb.Len() > 0 && multiline == MultilineNone {
    newMessage := api.Message{Role: "user", Content: sb.String()}

    if opts.MultiModal {
        msg, images, err := extractFileData(sb.String())
        if err != nil {
            return err
        }
            // clear all previous images for better responses
        if len(images) > 0 {
            for i := range opts.Messages {
                opts.Messages[i].Images = nil
            }
        }

        newMessage.Content = msg
        newMessage.Images = images
    }

    opts.Messages = append(opts.Messages, newMessage)

    assistant, err := chat(cmd, opts)
    if err != nil {
        return err
    }
    if assistant != nil {
        opts.Messages = append(opts.Messages, *assistant)
    }

    sb.Reset()
}

最終的にollama run llama3.1で以下のようにやりとりが出来ることになります。

$ ollama run llama3.1
>>> hello
Hello! How are you today? Is there something I can help you with, or would you like to chat?