parser.go 4.08 KB
Newer Older
1
2
3
4
package parser

import (
	"bufio"
5
	"bytes"
6
	"errors"
7
	"fmt"
8
	"io"
Michael Yang's avatar
Michael Yang committed
9
10
	"strconv"
	"strings"
11
12
13
14
)

type Command struct {
	Name string
15
16
17
	Args string
}

Michael Yang's avatar
Michael Yang committed
18
type state int
19

Michael Yang's avatar
Michael Yang committed
20
21
22
23
24
25
26
27
const (
	stateNil state = iota
	stateName
	stateValue
	stateParameter
	stateMessage
	stateComment
)
28

Michael Yang's avatar
tests  
Michael Yang committed
29
30
31
32
var (
	errMissingFrom = errors.New("no FROM line")
	errInvalidRole = errors.New("role must be one of \"system\", \"user\", or \"assistant\"")
)
Michael Yang's avatar
Michael Yang committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

func Parse(r io.Reader) (cmds []Command, err error) {
	var cmd Command
	var curr state
	var b bytes.Buffer
	var role string

	br := bufio.NewReader(r)
	for {
		r, _, err := br.ReadRune()
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return nil, err
		}
48

Michael Yang's avatar
Michael Yang committed
49
50
51
52
53
		next, r, err := parseRuneForState(r, curr)
		if errors.Is(err, io.ErrUnexpectedEOF) {
			return nil, fmt.Errorf("%w: %s", err, b.String())
		} else if err != nil {
			return nil, err
54
55
		}

Michael Yang's avatar
Michael Yang committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
		if next != curr {
			switch curr {
			case stateName, stateParameter:
				switch s := strings.ToLower(b.String()); s {
				case "from":
					cmd.Name = "model"
				case "parameter":
					next = stateParameter
				case "message":
					next = stateMessage
					fallthrough
				default:
					cmd.Name = s
				}
			case stateMessage:
Michael Yang's avatar
linting  
Michael Yang committed
71
				if !isValidRole(b.String()) {
Michael Yang's avatar
Michael Yang committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
					return nil, errInvalidRole
				}

				role = b.String()
			case stateComment, stateNil:
				// pass
			case stateValue:
				s, ok := unquote(b.String())
				if !ok || isSpace(r) {
					if _, err := b.WriteRune(r); err != nil {
						return nil, err
					}

					continue
				}

				if role != "" {
					s = role + ": " + s
					role = ""
				}

				cmd.Args = s
				cmds = append(cmds, cmd)
Michael Yang's avatar
Michael Yang committed
95
96
			}

Michael Yang's avatar
Michael Yang committed
97
98
99
100
101
102
103
			b.Reset()
			curr = next
		}

		if strconv.IsPrint(r) {
			if _, err := b.WriteRune(r); err != nil {
				return nil, err
Michael Yang's avatar
Michael Yang committed
104
			}
Michael Yang's avatar
Michael Yang committed
105
106
107
108
109
110
111
112
		}
	}

	// flush the buffer
	switch curr {
	case stateComment, stateNil:
		// pass; nothing to flush
	case stateValue:
Michael Yang's avatar
Michael Yang committed
113
114
		s, ok := unquote(b.String())
		if !ok {
Michael Yang's avatar
Michael Yang committed
115
			return nil, io.ErrUnexpectedEOF
116
		}
117

Michael Yang's avatar
Michael Yang committed
118
119
120
121
122
		if role != "" {
			s = role + ": " + s
		}

		cmd.Args = s
Michael Yang's avatar
Michael Yang committed
123
124
125
		cmds = append(cmds, cmd)
	default:
		return nil, io.ErrUnexpectedEOF
126
127
	}

Michael Yang's avatar
Michael Yang committed
128
129
130
131
	for _, cmd := range cmds {
		if cmd.Name == "model" {
			return cmds, nil
		}
132
133
	}

Michael Yang's avatar
tests  
Michael Yang committed
134
	return nil, errMissingFrom
135
}
136

Michael Yang's avatar
Michael Yang committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
func parseRuneForState(r rune, cs state) (state, rune, error) {
	switch cs {
	case stateNil:
		switch {
		case r == '#':
			return stateComment, 0, nil
		case isSpace(r), isNewline(r):
			return stateNil, 0, nil
		default:
			return stateName, r, nil
		}
	case stateName:
		switch {
		case isAlpha(r):
			return stateName, r, nil
		case isSpace(r):
			return stateValue, 0, nil
		default:
			return stateNil, 0, errors.New("invalid")
		}
	case stateValue:
		switch {
		case isNewline(r):
			return stateNil, r, nil
		case isSpace(r):
			return stateNil, r, nil
		default:
			return stateValue, r, nil
		}
	case stateParameter:
		switch {
		case isAlpha(r), isNumber(r), r == '_':
			return stateParameter, r, nil
		case isSpace(r):
			return stateValue, 0, nil
		default:
			return stateNil, 0, io.ErrUnexpectedEOF
		}
	case stateMessage:
		switch {
		case isAlpha(r):
			return stateMessage, r, nil
		case isSpace(r):
			return stateValue, 0, nil
		default:
			return stateNil, 0, io.ErrUnexpectedEOF
		}
	case stateComment:
		switch {
		case isNewline(r):
			return stateNil, 0, nil
		default:
			return stateComment, 0, nil
		}
	default:
		return stateNil, 0, errors.New("")
193
	}
Michael Yang's avatar
Michael Yang committed
194
}
195

Michael Yang's avatar
Michael Yang committed
196
197
198
func unquote(s string) (string, bool) {
	if len(s) == 0 {
		return "", false
199
	}
200

Michael Yang's avatar
Michael Yang committed
201
202
203
204
205
206
207
	// TODO: single quotes
	if len(s) >= 3 && s[:3] == `"""` {
		if len(s) >= 6 && s[len(s)-3:] == `"""` {
			return s[3 : len(s)-3], true
		}

		return "", false
208
209
	}

Michael Yang's avatar
Michael Yang committed
210
211
212
213
214
215
	if len(s) >= 1 && s[0] == '"' {
		if len(s) >= 2 && s[len(s)-1] == '"' {
			return s[1 : len(s)-1], true
		}

		return "", false
216
217
	}

Michael Yang's avatar
Michael Yang committed
218
	return s, true
219
220
}

Michael Yang's avatar
Michael Yang committed
221
222
223
func isAlpha(r rune) bool {
	return r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z'
}
224

Michael Yang's avatar
Michael Yang committed
225
226
227
func isNumber(r rune) bool {
	return r >= '0' && r <= '9'
}
228

Michael Yang's avatar
Michael Yang committed
229
230
231
func isSpace(r rune) bool {
	return r == ' ' || r == '\t'
}
Michael Yang's avatar
Michael Yang committed
232

Michael Yang's avatar
Michael Yang committed
233
234
235
func isNewline(r rune) bool {
	return r == '\r' || r == '\n'
}
236

Michael Yang's avatar
Michael Yang committed
237
238
func isValidRole(role string) bool {
	return role == "system" || role == "user" || role == "assistant"
239
}