history.go 2.18 KB
Newer Older
Patrick Devine's avatar
Patrick Devine committed
1
2
3
4
5
package readline

import (
	"bufio"
	"errors"
Michael Yang's avatar
Michael Yang committed
6
	"fmt"
Patrick Devine's avatar
Patrick Devine committed
7
8
9
10
11
	"io"
	"os"
	"path/filepath"
	"strings"

Michael Yang's avatar
Michael Yang committed
12
	"github.com/emirpasic/gods/v2/lists/arraylist"
Patrick Devine's avatar
Patrick Devine committed
13
14
15
)

type History struct {
Michael Yang's avatar
Michael Yang committed
16
	Buf      *arraylist.List[string]
Patrick Devine's avatar
Patrick Devine committed
17
18
19
20
21
22
23
24
25
	Autosave bool
	Pos      int
	Limit    int
	Filename string
	Enabled  bool
}

func NewHistory() (*History, error) {
	h := &History{
Michael Yang's avatar
Michael Yang committed
26
		Buf:      arraylist.New[string](),
Michael Yang's avatar
Michael Yang committed
27
		Limit:    100, // resizeme
Patrick Devine's avatar
Patrick Devine committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
		Autosave: true,
		Enabled:  true,
	}

	err := h.Init()
	if err != nil {
		return nil, err
	}

	return h, nil
}

func (h *History) Init() error {
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	path := filepath.Join(home, ".ollama", "history")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
47
48
49
50
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
51
52
	h.Filename = path

Michael Yang's avatar
Michael Yang committed
53
	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDONLY, 0o600)
Patrick Devine's avatar
Patrick Devine committed
54
55
56
57
58
59
60
61
62
63
64
65
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return err
	}
	defer f.Close()

	r := bufio.NewReader(f)
	for {
		line, err := r.ReadString('\n')
		if err != nil {
66
			if errors.Is(err, io.EOF) {
Patrick Devine's avatar
Patrick Devine committed
67
68
69
70
71
72
73
74
75
76
				break
			}
			return err
		}

		line = strings.TrimSpace(line)
		if len(line) == 0 {
			continue
		}

Michael Yang's avatar
Michael Yang committed
77
		h.Add(line)
Patrick Devine's avatar
Patrick Devine committed
78
79
80
81
82
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
83
84
func (h *History) Add(s string) {
	h.Buf.Add(s)
Patrick Devine's avatar
Patrick Devine committed
85
	h.Compact()
86
	h.Pos = h.Size()
Patrick Devine's avatar
Patrick Devine committed
87
	if h.Autosave {
Michael Yang's avatar
Michael Yang committed
88
		_ = h.Save()
Patrick Devine's avatar
Patrick Devine committed
89
90
91
92
93
94
	}
}

func (h *History) Compact() {
	s := h.Buf.Size()
	if s > h.Limit {
Michael Yang's avatar
lint  
Michael Yang committed
95
		for range s - h.Limit {
Patrick Devine's avatar
Patrick Devine committed
96
97
98
99
100
101
102
103
104
			h.Buf.Remove(0)
		}
	}
}

func (h *History) Clear() {
	h.Buf.Clear()
}

Michael Yang's avatar
Michael Yang committed
105
func (h *History) Prev() (line string) {
Patrick Devine's avatar
Patrick Devine committed
106
107
108
	if h.Pos > 0 {
		h.Pos -= 1
	}
Michael Yang's avatar
Michael Yang committed
109
	line, _ = h.Buf.Get(h.Pos)
Patrick Devine's avatar
Patrick Devine committed
110
111
112
	return line
}

Michael Yang's avatar
Michael Yang committed
113
func (h *History) Next() (line string) {
Patrick Devine's avatar
Patrick Devine committed
114
115
	if h.Pos < h.Buf.Size() {
		h.Pos += 1
Michael Yang's avatar
Michael Yang committed
116
		line, _ = h.Buf.Get(h.Pos)
Patrick Devine's avatar
Patrick Devine committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
	}
	return line
}

func (h *History) Size() int {
	return h.Buf.Size()
}

func (h *History) Save() error {
	if !h.Enabled {
		return nil
	}

	tmpFile := h.Filename + ".tmp"

132
	f, err := os.OpenFile(tmpFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC|os.O_APPEND, 0o600)
Patrick Devine's avatar
Patrick Devine committed
133
134
135
136
137
138
	if err != nil {
		return err
	}
	defer f.Close()

	buf := bufio.NewWriter(f)
Michael Yang's avatar
lint  
Michael Yang committed
139
	for cnt := range h.Size() {
Michael Yang's avatar
Michael Yang committed
140
141
		line, _ := h.Buf.Get(cnt)
		fmt.Fprintln(buf, line)
Patrick Devine's avatar
Patrick Devine committed
142
143
144
145
146
147
148
149
150
151
	}
	buf.Flush()
	f.Close()

	if err = os.Rename(tmpFile, h.Filename); err != nil {
		return err
	}

	return nil
}