63 lines
1.2 KiB
Go
63 lines
1.2 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/elk-language/go-prompt"
|
|
"github.com/fatih/color"
|
|
"github.com/ollama/ollama/api"
|
|
"log"
|
|
"strings"
|
|
)
|
|
|
|
func onUserInput(input string) {
|
|
if strings.HasPrefix(input, "/") {
|
|
executeCommand(input)
|
|
return
|
|
}
|
|
fmt.Print(colorRole(MT_ASSISTANT), ": ")
|
|
sendPromptInput(input, func(r api.ChatResponse) error {
|
|
_, err := fmt.Print(r.Message.Content)
|
|
return err
|
|
})
|
|
fmt.Println()
|
|
}
|
|
|
|
func runAsCommandLine() {
|
|
for _, msg := range conversation {
|
|
fmt.Printf("%s: %s\n", colorRole(roleToInt(msg.Role)), msg.Content)
|
|
}
|
|
|
|
runner := prompt.New(
|
|
onUserInput,
|
|
prompt.WithTitle("llamachat"),
|
|
prompt.WithPrefix("user: "),
|
|
prompt.WithCompleter(Completer),
|
|
)
|
|
|
|
runner.Run()
|
|
}
|
|
|
|
func executeCommand(cli string) {
|
|
args := strings.Split(cli, " ")
|
|
for _, cmd := range Commands {
|
|
if cmd.Name == args[0] {
|
|
cmd.Action(args)
|
|
return
|
|
}
|
|
}
|
|
fmt.Println("Unknown command")
|
|
}
|
|
|
|
func colorRole(role MessageType) string {
|
|
if role == MT_SYSTEM {
|
|
return color.RedString("system")
|
|
} else if role == MT_ASSISTANT {
|
|
return color.GreenString("assistant")
|
|
} else if role == MT_USER {
|
|
return color.BlueString("user")
|
|
} else {
|
|
log.Fatalf("Invalid role type %d\n", role)
|
|
return ""
|
|
}
|
|
}
|