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

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

	"golang.org/x/text/encoding/unicode"
	"golang.org/x/text/transform"
14
15
)

Michael Yang's avatar
Michael Yang committed
16
17
18
19
20
21
22
23
24
25
26
27
28
type File struct {
	Commands []Command
}

func (f File) String() string {
	var sb strings.Builder
	for _, cmd := range f.Commands {
		fmt.Fprintln(&sb, cmd.String())
	}

	return sb.String()
}

29
30
type Command struct {
	Name string
31
32
33
	Args string
}

Michael Yang's avatar
Michael Yang committed
34
func (c Command) String() string {
Michael Yang's avatar
Michael Yang committed
35
	var sb strings.Builder
Michael Yang's avatar
Michael Yang committed
36
37
	switch c.Name {
	case "model":
Michael Yang's avatar
Michael Yang committed
38
		fmt.Fprintf(&sb, "FROM %s", c.Args)
Michael Yang's avatar
Michael Yang committed
39
	case "license", "template", "system", "adapter":
Michael Yang's avatar
Michael Yang committed
40
		fmt.Fprintf(&sb, "%s %s", strings.ToUpper(c.Name), quote(c.Args))
Michael Yang's avatar
Michael Yang committed
41
42
	case "message":
		role, message, _ := strings.Cut(c.Args, ": ")
Michael Yang's avatar
Michael Yang committed
43
		fmt.Fprintf(&sb, "MESSAGE %s %s", role, quote(message))
Michael Yang's avatar
Michael Yang committed
44
	default:
Michael Yang's avatar
Michael Yang committed
45
		fmt.Fprintf(&sb, "PARAMETER %s %s", c.Name, quote(c.Args))
Michael Yang's avatar
Michael Yang committed
46
47
	}

Michael Yang's avatar
Michael Yang committed
48
	return sb.String()
Michael Yang's avatar
Michael Yang committed
49
50
}

Michael Yang's avatar
Michael Yang committed
51
type state int
52

Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
58
59
60
const (
	stateNil state = iota
	stateName
	stateValue
	stateParameter
	stateMessage
	stateComment
)
61

Michael Yang's avatar
tests  
Michael Yang committed
62
var (
Michael Yang's avatar
Michael Yang committed
63
64
65
	errMissingFrom        = errors.New("no FROM line")
	errInvalidMessageRole = errors.New("message role must be one of \"system\", \"user\", or \"assistant\"")
	errInvalidCommand     = errors.New("command must be one of \"from\", \"license\", \"template\", \"system\", \"adapter\", \"parameter\", or \"message\"")
Michael Yang's avatar
tests  
Michael Yang committed
66
)
Michael Yang's avatar
Michael Yang committed
67

68
69
70
71
72
73
74
75
76
77
78
79
type ParserError struct {
	LineNumber int
	Msg        string
}

func (e *ParserError) Error() string {
	if e.LineNumber > 0 {
		return fmt.Sprintf("(line %d): %s", e.LineNumber, e.Msg)
	}
	return e.Msg
}

Michael Yang's avatar
Michael Yang committed
80
func ParseFile(r io.Reader) (*File, error) {
Michael Yang's avatar
Michael Yang committed
81
82
	var cmd Command
	var curr state
83
	var currLine int = 1
Michael Yang's avatar
Michael Yang committed
84
85
86
	var b bytes.Buffer
	var role string

Michael Yang's avatar
Michael Yang committed
87
88
	var f File

Michael Yang's avatar
Michael Yang committed
89
90
91
	tr := unicode.BOMOverride(unicode.UTF8.NewDecoder())
	br := bufio.NewReader(transform.NewReader(r, tr))

Michael Yang's avatar
Michael Yang committed
92
93
94
95
96
97
98
	for {
		r, _, err := br.ReadRune()
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return nil, err
		}
99

100
101
102
103
		if isNewline(r) {
			currLine++
		}

Michael Yang's avatar
Michael Yang committed
104
105
106
107
		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 {
108
109
110
111
			return nil, &ParserError{
				LineNumber: currLine,
				Msg:        err.Error(),
			}
112
113
		}

Michael Yang's avatar
Michael Yang committed
114
		// process the state transition, some transitions need to be intercepted and redirected
Michael Yang's avatar
Michael Yang committed
115
116
		if next != curr {
			switch curr {
Michael Yang's avatar
Michael Yang committed
117
118
			case stateName:
				if !isValidCommand(b.String()) {
119
120
121
122
					return nil, &ParserError{
						LineNumber: currLine,
						Msg:        errInvalidCommand.Error(),
					}
Michael Yang's avatar
Michael Yang committed
123
124
				}

Michael Yang's avatar
Michael Yang committed
125
				// next state sometimes depends on the current buffer value
Michael Yang's avatar
Michael Yang committed
126
127
128
129
				switch s := strings.ToLower(b.String()); s {
				case "from":
					cmd.Name = "model"
				case "parameter":
Michael Yang's avatar
Michael Yang committed
130
					// transition to stateParameter which sets command name
Michael Yang's avatar
Michael Yang committed
131
132
					next = stateParameter
				case "message":
Michael Yang's avatar
Michael Yang committed
133
					// transition to stateMessage which validates the message role
Michael Yang's avatar
Michael Yang committed
134
135
136
137
138
					next = stateMessage
					fallthrough
				default:
					cmd.Name = s
				}
Michael Yang's avatar
Michael Yang committed
139
140
			case stateParameter:
				cmd.Name = b.String()
Michael Yang's avatar
Michael Yang committed
141
			case stateMessage:
Michael Yang's avatar
Michael Yang committed
142
				if !isValidMessageRole(b.String()) {
143
144
145
146
					return nil, &ParserError{
						LineNumber: currLine,
						Msg:        errInvalidMessageRole.Error(),
					}
Michael Yang's avatar
Michael Yang committed
147
148
149
150
151
152
				}

				role = b.String()
			case stateComment, stateNil:
				// pass
			case stateValue:
Josh Yan's avatar
Josh Yan committed
153
				s, ok := unquote(strings.TrimSpace(b.String()))
Michael Yang's avatar
Michael Yang committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
				if !ok || isSpace(r) {
					if _, err := b.WriteRune(r); err != nil {
						return nil, err
					}

					continue
				}

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

				cmd.Args = s
Michael Yang's avatar
Michael Yang committed
168
				f.Commands = append(f.Commands, cmd)
Michael Yang's avatar
Michael Yang committed
169
170
			}

Michael Yang's avatar
Michael Yang committed
171
172
173
174
175
176
177
			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
178
			}
Michael Yang's avatar
Michael Yang committed
179
180
181
182
183
184
185
186
		}
	}

	// flush the buffer
	switch curr {
	case stateComment, stateNil:
		// pass; nothing to flush
	case stateValue:
Josh Yan's avatar
Josh Yan committed
187
		s, ok := unquote(strings.TrimSpace(b.String()))
Michael Yang's avatar
Michael Yang committed
188
		if !ok {
Michael Yang's avatar
Michael Yang committed
189
			return nil, io.ErrUnexpectedEOF
190
		}
191

Michael Yang's avatar
Michael Yang committed
192
193
194
195
196
		if role != "" {
			s = role + ": " + s
		}

		cmd.Args = s
Michael Yang's avatar
Michael Yang committed
197
		f.Commands = append(f.Commands, cmd)
Michael Yang's avatar
Michael Yang committed
198
199
	default:
		return nil, io.ErrUnexpectedEOF
200
201
	}

Michael Yang's avatar
Michael Yang committed
202
	for _, cmd := range f.Commands {
Michael Yang's avatar
Michael Yang committed
203
		if cmd.Name == "model" {
Michael Yang's avatar
Michael Yang committed
204
			return &f, nil
Michael Yang's avatar
Michael Yang committed
205
		}
206
207
	}

Michael Yang's avatar
tests  
Michael Yang committed
208
	return nil, errMissingFrom
209
}
210

Michael Yang's avatar
Michael Yang committed
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
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:
Michael Yang's avatar
Michael Yang committed
229
			return stateNil, 0, errInvalidCommand
Michael Yang's avatar
Michael Yang committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
		}
	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("")
267
	}
Michael Yang's avatar
Michael Yang committed
268
}
269

Michael Yang's avatar
Michael Yang committed
270
func quote(s string) string {
271
	if strings.Contains(s, "\n") || strings.HasPrefix(s, " ") || strings.HasSuffix(s, " ") {
Michael Yang's avatar
Michael Yang committed
272
273
274
275
		if strings.Contains(s, "\"") {
			return `"""` + s + `"""`
		}

276
		return `"` + s + `"`
Michael Yang's avatar
Michael Yang committed
277
278
279
280
281
	}

	return s
}

Michael Yang's avatar
Michael Yang committed
282
283
284
285
286
287
288
289
func unquote(s string) (string, bool) {
	// 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
290
291
	}

Michael Yang's avatar
Michael Yang committed
292
293
294
295
296
297
	if len(s) >= 1 && s[0] == '"' {
		if len(s) >= 2 && s[len(s)-1] == '"' {
			return s[1 : len(s)-1], true
		}

		return "", false
298
299
	}

Michael Yang's avatar
Michael Yang committed
300
	return s, true
301
302
}

Michael Yang's avatar
Michael Yang committed
303
304
305
func isAlpha(r rune) bool {
	return r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z'
}
306

Michael Yang's avatar
Michael Yang committed
307
308
309
func isNumber(r rune) bool {
	return r >= '0' && r <= '9'
}
310

Michael Yang's avatar
Michael Yang committed
311
312
313
func isSpace(r rune) bool {
	return r == ' ' || r == '\t'
}
Michael Yang's avatar
Michael Yang committed
314

Michael Yang's avatar
Michael Yang committed
315
316
317
func isNewline(r rune) bool {
	return r == '\r' || r == '\n'
}
318

Michael Yang's avatar
Michael Yang committed
319
func isValidMessageRole(role string) bool {
Michael Yang's avatar
Michael Yang committed
320
	return role == "system" || role == "user" || role == "assistant"
321
}
Michael Yang's avatar
Michael Yang committed
322
323
324
325
326
327
328
329
330

func isValidCommand(cmd string) bool {
	switch strings.ToLower(cmd) {
	case "from", "license", "template", "system", "adapter", "parameter", "message":
		return true
	default:
		return false
	}
}