history.go 2.32 KB
Newer Older
Patrick Devine's avatar
Patrick Devine committed
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
package readline

import (
	"bufio"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/emirpasic/gods/lists/arraylist"
)

type History struct {
	Buf      *arraylist.List
	Autosave bool
	Pos      int
	Limit    int
	Filename string
	Enabled  bool
}

func NewHistory() (*History, error) {
	h := &History{
		Buf:      arraylist.New(),
Michael Yang's avatar
Michael Yang committed
26
		Limit:    100, // resizeme
Patrick Devine's avatar
Patrick Devine committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
		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
46
47
48
49
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}

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

Michael Yang's avatar
Michael Yang committed
52
	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDONLY, 0o600)
Patrick Devine's avatar
Patrick Devine committed
53
54
55
56
57
58
59
60
61
62
63
64
	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 {
65
			if errors.Is(err, io.EOF) {
Patrick Devine's avatar
Patrick Devine committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
				break
			}
			return err
		}

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

		h.Add([]rune(line))
	}

	return nil
}

func (h *History) Add(l []rune) {
	h.Buf.Add(l)
	h.Compact()
85
	h.Pos = h.Size()
Patrick Devine's avatar
Patrick Devine committed
86
	if h.Autosave {
Michael Yang's avatar
Michael Yang committed
87
		_ = h.Save()
Patrick Devine's avatar
Patrick Devine committed
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
	}
}

func (h *History) Compact() {
	s := h.Buf.Size()
	if s > h.Limit {
		for cnt := 0; cnt < s-h.Limit; cnt++ {
			h.Buf.Remove(0)
		}
	}
}

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

func (h *History) Prev() []rune {
	var line []rune
	if h.Pos > 0 {
		h.Pos -= 1
	}
	v, _ := h.Buf.Get(h.Pos)
	line, _ = v.([]rune)
	return line
}

func (h *History) Next() []rune {
	var line []rune
	if h.Pos < h.Buf.Size() {
		h.Pos += 1
		v, _ := h.Buf.Get(h.Pos)
		line, _ = v.([]rune)
	}
	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"

135
	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
136
137
138
139
140
141
142
143
144
	if err != nil {
		return err
	}
	defer f.Close()

	buf := bufio.NewWriter(f)
	for cnt := 0; cnt < h.Size(); cnt++ {
		v, _ := h.Buf.Get(cnt)
		line, _ := v.([]rune)
145
146
147
		if _, err := buf.WriteString(string(line) + "\n"); err != nil {
			return err
		}
Patrick Devine's avatar
Patrick Devine committed
148
149
150
151
152
153
154
155
156
157
	}
	buf.Flush()
	f.Close()

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

	return nil
}