136 lines
2.7 KiB
Go
136 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"github.com/ollama/ollama/api"
|
|
"log"
|
|
"os"
|
|
)
|
|
|
|
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 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 sendPromptInput(input string, handler func(response api.ChatResponse) error) {
|
|
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
|
|
err := handler(resp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !resp.Done {
|
|
return nil
|
|
}
|
|
return SaveMessage(MT_ASSISTANT, fullResponse)
|
|
}
|
|
|
|
err = ollama.Chat(ctx, req, respFunc)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
conversation = append(conversation, api.Message{
|
|
Role: convertRole(MT_ASSISTANT),
|
|
Content: fullResponse,
|
|
})
|
|
}
|
|
|
|
const usage = `Usage:
|
|
llamachat [OPTION]... [CONVERSATION]
|
|
|
|
Options:
|
|
-w, --web Access via a web page.
|
|
|
|
If CONVERSATION is left out, the conversation will be kept in-memory. If it is
|
|
present, the conversation will be stored in a file called CONVERSATION.chat.
|
|
Alternatively, it can also be passed the literal string '-' (a singel dash) to
|
|
ensure an in-memory database is used.
|
|
`
|
|
|
|
func main() {
|
|
flag.Usage = func() { fmt.Fprintf(os.Stderr, "%s\n", usage) }
|
|
|
|
var err error
|
|
var useWeb bool
|
|
db := ""
|
|
flag.BoolVar(&useWeb, "web", false, "run in web mode")
|
|
flag.BoolVar(&useWeb, "w", false, "run in web mode")
|
|
|
|
flag.Parse()
|
|
|
|
db = flag.Arg(0)
|
|
if db == "-" {
|
|
db = ""
|
|
}
|
|
OpenDb(db)
|
|
conversation = loadMessageFromDb()
|
|
|
|
ollama, err = api.ClientFromEnvironment()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
if useWeb {
|
|
runAsWeb()
|
|
} else {
|
|
runAsCommandLine()
|
|
}
|
|
}
|