ollama

1
package readline
2

3
import (
4
	"bufio"
5
	"fmt"
6
	"io"
7
	"os"
8
	"syscall"
9
)
10

11
type Prompt struct {
12
	Prompt         string
13
	AltPrompt      string
14
	Placeholder    string
15
	AltPlaceholder string
16
	UseAlt         bool
17
}
18

19
func (p *Prompt) prompt() string {
20
	if p.UseAlt {
21
		return p.AltPrompt
22
	}
23
	return p.Prompt
24
}
25

26
func (p *Prompt) placeholder() string {
27
	if p.UseAlt {
28
		return p.AltPlaceholder
29
	}
30
	return p.Placeholder
31
}
32

33
type Terminal struct {
34
	outchan chan rune
35
	rawmode bool
36
	termios any
37
}
38

39
type Instance struct {
40
	Prompt   *Prompt
41
	Terminal *Terminal
42
	History  *History
43
	Pasting  bool
44
}
45

46
func New(prompt Prompt) (*Instance, error) {
47
	term, err := NewTerminal()
48
	if err != nil {
49
		return nil, err
50
	}
51

52
	history, err := NewHistory()
53
	if err != nil {
54
		return nil, err
55
	}
56

57
	return &Instance{
58
		Prompt:   &prompt,
59
		Terminal: term,
60
		History:  history,
61
	}, nil
62
}
63

64
func (i *Instance) Readline() (string, error) {
65
	if !i.Terminal.rawmode {
66
		fd := int(syscall.Stdin)
67
		termios, err := SetRawMode(fd)
68
		if err != nil {
69
			return "", err
70
		}
71
		i.Terminal.rawmode = true
72
		i.Terminal.termios = termios
73
	}
74

75
	prompt := i.Prompt.prompt()
76
	if i.Pasting {
77
		// force alt prompt when pasting
78
		prompt = i.Prompt.AltPrompt
79
	}
80
	fmt.Print(prompt)
81

82
	defer func() {
83
		fd := int(syscall.Stdin)
84
		// nolint: errcheck
85
		UnsetRawMode(fd, i.Terminal.termios)
86
		i.Terminal.rawmode = false
87
	}()
88

89
	buf, _ := NewBuffer(i.Prompt)
90

91
	var esc bool
92
	var escex bool
93
	var metaDel bool
94

95
	var currentLineBuf []rune
96

97
	for {
98
		// don't show placeholder when pasting unless we're in multiline mode
99
		showPlaceholder := !i.Pasting || i.Prompt.UseAlt
100
		if buf.IsEmpty() && showPlaceholder {
101
			ph := i.Prompt.placeholder()
102
			fmt.Printf(ColorGrey + ph + fmt.Sprintf(CursorLeftN, len(ph)) + ColorDefault)
103
		}
104

105
		r, err := i.Terminal.Read()
106

107
		if buf.IsEmpty() {
108
			fmt.Print(ClearToEOL)
109
		}
110

111
		if err != nil {
112
			return "", io.EOF
113
		}
114

115
		if escex {
116
			escex = false
117

118
			switch r {
119
			case KeyUp:
120
				if i.History.Pos > 0 {
121
					if i.History.Pos == i.History.Size() {
122
						currentLineBuf = []rune(buf.String())
123
					}
124
					buf.Replace(i.History.Prev())
125
				}
126
			case KeyDown:
127
				if i.History.Pos < i.History.Size() {
128
					buf.Replace(i.History.Next())
129
					if i.History.Pos == i.History.Size() {
130
						buf.Replace(currentLineBuf)
131
					}
132
				}
133
			case KeyLeft:
134
				buf.MoveLeft()
135
			case KeyRight:
136
				buf.MoveRight()
137
			case CharBracketedPaste:
138
				var code string
139
				for cnt := 0; cnt < 3; cnt++ {
140
					r, err = i.Terminal.Read()
141
					if err != nil {
142
						return "", io.EOF
143
					}
144

145
					code += string(r)
146
				}
147
				if code == CharBracketedPasteStart {
148
					i.Pasting = true
149
				} else if code == CharBracketedPasteEnd {
150
					i.Pasting = false
151
				}
152
			case KeyDel:
153
				if buf.Size() > 0 {
154
					buf.Delete()
155
				}
156
				metaDel = true
157
			case MetaStart:
158
				buf.MoveToStart()
159
			case MetaEnd:
160
				buf.MoveToEnd()
161
			default:
162
				// skip any keys we don't know about
163
				continue
164
			}
165
			continue
166
		} else if esc {
167
			esc = false
168

169
			switch r {
170
			case 'b':
171
				buf.MoveLeftWord()
172
			case 'f':
173
				buf.MoveRightWord()
174
			case CharBackspace:
175
				buf.DeleteWord()
176
			case CharEscapeEx:
177
				escex = true
178
			}
179
			continue
180
		}
181

182
		switch r {
183
		case CharNull:
184
			continue
185
		case CharEsc:
186
			esc = true
187
		case CharInterrupt:
188
			return "", ErrInterrupt
189
		case CharLineStart:
190
			buf.MoveToStart()
191
		case CharLineEnd:
192
			buf.MoveToEnd()
193
		case CharBackward:
194
			buf.MoveLeft()
195
		case CharForward:
196
			buf.MoveRight()
197
		case CharBackspace, CharCtrlH:
198
			buf.Remove()
199
		case CharTab:
200
			// todo: convert back to real tabs
201
			for cnt := 0; cnt < 8; cnt++ {
202
				buf.Add(' ')
203
			}
204
		case CharDelete:
205
			if buf.Size() > 0 {
206
				buf.Delete()
207
			} else {
208
				return "", io.EOF
209
			}
210
		case CharKill:
211
			buf.DeleteRemaining()
212
		case CharCtrlU:
213
			buf.DeleteBefore()
214
		case CharCtrlL:
215
			buf.ClearScreen()
216
		case CharCtrlW:
217
			buf.DeleteWord()
218
		case CharCtrlZ:
219
			fd := int(syscall.Stdin)
220
			return handleCharCtrlZ(fd, i.Terminal.termios)
221
		case CharEnter:
222
			output := buf.String()
223
			if output != "" {
224
				i.History.Add([]rune(output))
225
			}
226
			buf.MoveToEnd()
227
			fmt.Println()
228

229
			return output, nil
230
		default:
231
			if metaDel {
232
				metaDel = false
233
				continue
234
			}
235
			if r >= CharSpace || r == CharEnter {
236
				buf.Add(r)
237
			}
238
		}
239
	}
240
}
241

242
func (i *Instance) HistoryEnable() {
243
	i.History.Enabled = true
244
}
245

246
func (i *Instance) HistoryDisable() {
247
	i.History.Enabled = false
248
}
249

250
func NewTerminal() (*Terminal, error) {
251
	fd := int(syscall.Stdin)
252
	termios, err := SetRawMode(fd)
253
	if err != nil {
254
		return nil, err
255
	}
256

257
	t := &Terminal{
258
		outchan: make(chan rune),
259
		rawmode: true,
260
		termios: termios,
261
	}
262

263
	go t.ioloop()
264

265
	return t, nil
266
}
267

268
func (t *Terminal) ioloop() {
269
	buf := bufio.NewReader(os.Stdin)
270

271
	for {
272
		r, _, err := buf.ReadRune()
273
		if err != nil {
274
			close(t.outchan)
275
			break
276
		}
277
		t.outchan <- r
278
	}
279
}
280

281
func (t *Terminal) Read() (rune, error) {
282
	r, ok := <-t.outchan
283
	if !ok {
284
		return 0, io.EOF
285
	}
286

287
	return r, nil
288
}
289