/
main.go
134 lines (116 loc) · 3.25 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package main
import (
"context"
"encoding/json"
"fmt"
"log"
"os"
"os/exec"
"path"
"runtime"
"strings"
"github.com/jmorganca/ollama/api"
"github.com/pterm/pterm"
)
const (
ollamaHost = "127.0.0.1:11435"
ollamaModel = "codellama"
)
func getLogPath() (string, error) {
xdgStateHome, ok := os.LookupEnv("XDG_STATE_HOME")
if !ok {
if homeDir, err := os.UserHomeDir(); err != nil {
return "", err
} else {
xdgStateHome = path.Join(homeDir, "state")
}
}
logPath := path.Join(xdgStateHome, "revman")
if err := os.MkdirAll(logPath, os.ModePerm); err != nil {
return "", err
}
return path.Join(logPath, "general.log"), nil
}
func selectOption(keyValueOptions map[string]any) string {
options := []string{}
for compl, desc := range keyValueOptions {
if desc, ok := desc.(string); ok {
options = append(options, fmt.Sprintf("%s (%s)", compl, pterm.Green(desc)))
}
}
selectedOption, _ := pterm.DefaultInteractiveSelect.WithOptions(options).Show()
return selectedOption
}
func main() {
logPath, _ := getLogPath()
f, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
if err != nil {
log.Fatalf("error opening file: %v", err)
}
defer f.Close()
log.SetOutput(f)
if envErr := os.Setenv("OLLAMA_HOST", ollamaHost); envErr != nil {
log.Fatalf("failed to set the `OLLAMA_HOST` env var: %s", envErr)
}
cmd := exec.Command("ollama", "serve")
if startCmdErr := cmd.Start(); startCmdErr != nil {
log.Fatalf("Failed to start cmd: %v", startCmdErr)
}
log.Println("Ollama server has started")
defer cmd.Wait()
ollamaClient, err := api.ClientFromEnvironment()
if err != nil {
log.Fatal(err)
}
ctx := context.Background()
if pullErr := ollamaClient.Pull(
ctx,
&api.PullRequest{Model: ollamaModel, Name: ollamaModel},
func(prog api.ProgressResponse) error {
log.Printf("%#v\n", prog)
return nil
},
); pullErr != nil {
log.Fatalf("could not pull the code llama model: %s", err)
}
command := strings.Join(os.Args[1:], " ")
systemPrompt := []string{
"You are a CLI completion tool from the UNIX man pages and you show all possible subcommands and flags as key value pairs,",
"where the key is the subcommand/flag and the value is the description. The commands should be available on the",
runtime.GOOS,
"operating system.",
"Please respond using JSON.",
}
if commandManPage, err := exec.Command("man", os.Args[1]).Output(); err == nil {
fullDesc := []string{
"The manual pages for the",
os.Args[1],
"command is",
string(commandManPage),
}
systemPrompt = append(systemPrompt, fullDesc...)
}
genResp := ""
if genErr := ollamaClient.Generate(
ctx,
&api.GenerateRequest{
Model: ollamaModel,
System: strings.Join(systemPrompt, " "),
Prompt: "Please complete the following command: " + command,
Format: "json",
},
func(resp api.GenerateResponse) error {
genResp += resp.Response
log.Printf("%+v\n", resp)
return nil
},
); genErr != nil {
log.Fatalf("could not get a response: %v", genErr)
}
var respJSON map[string]any
if unmarshErr := json.Unmarshal([]byte(genResp), &respJSON); unmarshErr != nil {
log.Fatalf("could not parse Ollama response: %s", unmarshErr)
}
selected := selectOption(respJSON)
fmt.Print(command + " " + strings.Split(selected, "(")[0])
}