148 lines
2.9 KiB
Go
148 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"github.com/elk-language/go-prompt"
|
|
"github.com/fatih/color"
|
|
"github.com/ollama/ollama/api"
|
|
"log"
|
|
"os"
|
|
"strings"
|
|
)
|
|
|
|
var conversation []api.Message
|
|
var ollama *api.Client
|
|
|
|
func convertRole(role MessageType) string {
|
|
if role == MT_SYSTEM {
|
|
return "system"
|
|
} else if role == MT_ASSISTANT {
|
|
return "assistant"
|
|
} else if role == MT_USER {
|
|
return "user"
|
|
} else {
|
|
log.Fatalf("Invalid role type %d\n", role)
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func roleToInt(role string) MessageType {
|
|
if role == "system" {
|
|
return MT_SYSTEM
|
|
} else if role == "assistant" {
|
|
return MT_ASSISTANT
|
|
} else if role == "user" {
|
|
return MT_USER
|
|
} else {
|
|
log.Fatalf("Invalid role type %s\n", role)
|
|
return 0
|
|
}
|
|
}
|
|
|
|
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 ""
|
|
}
|
|
}
|
|
|
|
func loadMessageFromDb() []api.Message {
|
|
dbMessages := GetMesages()
|
|
var chatMessages []api.Message
|
|
for _, msg := range dbMessages {
|
|
message := api.Message{
|
|
Role: convertRole(msg.Type),
|
|
Content: msg.Content,
|
|
}
|
|
chatMessages = append(chatMessages, message)
|
|
fmt.Printf("%s: %s\n", colorRole(roleToInt(message.Role)), message.Content)
|
|
}
|
|
return chatMessages
|
|
}
|
|
|
|
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 onUserInput(input string) {
|
|
if strings.HasPrefix(input, "/") {
|
|
executeCommand(input)
|
|
return
|
|
}
|
|
err := SaveMessage(MT_USER, input)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
conversation = append(conversation, api.Message{
|
|
Role: convertRole(MT_USER),
|
|
Content: input,
|
|
})
|
|
|
|
ctx := context.Background()
|
|
req := &api.ChatRequest{
|
|
Model: "llama3.2:1b",
|
|
Messages: conversation,
|
|
}
|
|
|
|
fullResponse := ""
|
|
respFunc := func(resp api.ChatResponse) error {
|
|
fullResponse = fullResponse + resp.Message.Content
|
|
fmt.Print(resp.Message.Content)
|
|
if !resp.Done {
|
|
return nil
|
|
}
|
|
return SaveMessage(MT_ASSISTANT, fullResponse)
|
|
}
|
|
|
|
fmt.Print(colorRole(MT_ASSISTANT), ": ")
|
|
err = ollama.Chat(ctx, req, respFunc)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
conversation = append(conversation, api.Message{
|
|
Role: convertRole(MT_ASSISTANT),
|
|
Content: fullResponse,
|
|
})
|
|
fmt.Println()
|
|
}
|
|
|
|
func main() {
|
|
var err error
|
|
if len(os.Args) <= 1 {
|
|
log.Fatal("Missing command line parameter")
|
|
}
|
|
|
|
ollama, err = api.ClientFromEnvironment()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
dbName := os.Args[1]
|
|
OpenDb(dbName)
|
|
conversation = loadMessageFromDb()
|
|
|
|
runner := prompt.New(
|
|
onUserInput,
|
|
prompt.WithTitle("llamachat"),
|
|
prompt.WithPrefix("user: "),
|
|
prompt.WithCompleter(Completer),
|
|
)
|
|
|
|
runner.Run()
|
|
}
|