parser.go 5.11 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

Michael Yang's avatar
Michael Yang committed
34
func Format(cmds []Command) string {
35
	var sb strings.Builder
Michael Yang's avatar
Michael Yang committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
	for _, cmd := range cmds {
		name := cmd.Name
		args := cmd.Args

		switch cmd.Name {
		case "model":
			name = "from"
			args = cmd.Args
		case "license", "template", "system", "adapter":
			args = quote(args)
		case "message":
			role, message, _ := strings.Cut(cmd.Args, ": ")
			args = role + " " + quote(message)
		default:
			name = "parameter"
51
			args = cmd.Name + " " + quote(cmd.Args)
Michael Yang's avatar
Michael Yang committed
52
53
		}

54
		fmt.Fprintln(&sb, strings.ToUpper(name), args)
Michael Yang's avatar
Michael Yang committed
55
56
	}

57
	return sb.String()
Michael Yang's avatar
Michael Yang committed
58
59
}

Michael Yang's avatar
Michael Yang committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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
		}
74

Michael Yang's avatar
Michael Yang committed
75
76
77
78
79
		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
80
81
		}

Michael Yang's avatar
Michael Yang committed
82
		// process the state transition, some transitions need to be intercepted and redirected
Michael Yang's avatar
Michael Yang committed
83
84
85
		if next != curr {
			switch curr {
			case stateName, stateParameter:
Michael Yang's avatar
Michael Yang committed
86
				// next state sometimes depends on the current buffer value
Michael Yang's avatar
Michael Yang committed
87
88
89
90
				switch s := strings.ToLower(b.String()); s {
				case "from":
					cmd.Name = "model"
				case "parameter":
Michael Yang's avatar
Michael Yang committed
91
					// transition to stateParameter which sets command name
Michael Yang's avatar
Michael Yang committed
92
93
					next = stateParameter
				case "message":
Michael Yang's avatar
Michael Yang committed
94
					// transition to stateMessage which validates the message role
Michael Yang's avatar
Michael Yang committed
95
96
97
98
99
100
					next = stateMessage
					fallthrough
				default:
					cmd.Name = s
				}
			case stateMessage:
Michael Yang's avatar
Michael Yang committed
101
				if !isValidMessageRole(b.String()) {
Michael Yang's avatar
Michael Yang committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
					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
125
126
			}

Michael Yang's avatar
Michael Yang committed
127
128
129
130
131
132
133
			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
134
			}
Michael Yang's avatar
Michael Yang committed
135
136
137
138
139
140
141
142
		}
	}

	// flush the buffer
	switch curr {
	case stateComment, stateNil:
		// pass; nothing to flush
	case stateValue:
Michael Yang's avatar
Michael Yang committed
143
144
		s, ok := unquote(b.String())
		if !ok {
Michael Yang's avatar
Michael Yang committed
145
			return nil, io.ErrUnexpectedEOF
146
		}
147

Michael Yang's avatar
Michael Yang committed
148
149
150
151
152
		if role != "" {
			s = role + ": " + s
		}

		cmd.Args = s
Michael Yang's avatar
Michael Yang committed
153
154
155
		cmds = append(cmds, cmd)
	default:
		return nil, io.ErrUnexpectedEOF
156
157
	}

Michael Yang's avatar
Michael Yang committed
158
159
160
161
	for _, cmd := range cmds {
		if cmd.Name == "model" {
			return cmds, nil
		}
162
163
	}

Michael Yang's avatar
tests  
Michael Yang committed
164
	return nil, errMissingFrom
165
}
166

Michael Yang's avatar
Michael Yang committed
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
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
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("")
223
	}
Michael Yang's avatar
Michael Yang committed
224
}
225

Michael Yang's avatar
Michael Yang committed
226
func quote(s string) string {
227
	if strings.Contains(s, "\n") || strings.HasPrefix(s, " ") || strings.HasSuffix(s, " ") {
Michael Yang's avatar
Michael Yang committed
228
229
230
231
		if strings.Contains(s, "\"") {
			return `"""` + s + `"""`
		}

232
		return `"` + s + `"`
Michael Yang's avatar
Michael Yang committed
233
234
235
236
237
	}

	return s
}

Michael Yang's avatar
Michael Yang committed
238
239
240
func unquote(s string) (string, bool) {
	if len(s) == 0 {
		return "", false
241
	}
242

Michael Yang's avatar
Michael Yang committed
243
244
245
246
247
248
249
	// 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
250
251
	}

Michael Yang's avatar
Michael Yang committed
252
253
254
255
256
257
	if len(s) >= 1 && s[0] == '"' {
		if len(s) >= 2 && s[len(s)-1] == '"' {
			return s[1 : len(s)-1], true
		}

		return "", false
258
259
	}

Michael Yang's avatar
Michael Yang committed
260
	return s, true
261
262
}

Michael Yang's avatar
Michael Yang committed
263
264
265
func isAlpha(r rune) bool {
	return r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z'
}
266

Michael Yang's avatar
Michael Yang committed
267
268
269
func isNumber(r rune) bool {
	return r >= '0' && r <= '9'
}
270

Michael Yang's avatar
Michael Yang committed
271
272
273
func isSpace(r rune) bool {
	return r == ' ' || r == '\t'
}
Michael Yang's avatar
Michael Yang committed
274

Michael Yang's avatar
Michael Yang committed
275
276
277
func isNewline(r rune) bool {
	return r == '\r' || r == '\n'
}
278

Michael Yang's avatar
Michael Yang committed
279
func isValidMessageRole(role string) bool {
Michael Yang's avatar
Michael Yang committed
280
	return role == "system" || role == "user" || role == "assistant"
281
}