parser.go 4 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
Michael Yang committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
var errInvalidRole = errors.New("role must be one of \"system\", \"user\", or \"assistant\"")

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
		}
45

Michael Yang's avatar
Michael Yang committed
46
47
48
49
50
		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
51
52
		}

Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
		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
68
				if !isValidRole(b.String()) {
Michael Yang's avatar
Michael Yang committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
					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
92
93
			}

Michael Yang's avatar
Michael Yang committed
94
95
96
97
98
99
100
			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
101
			}
Michael Yang's avatar
Michael Yang committed
102
103
104
105
106
107
108
109
110
111
		}
	}

	// flush the buffer
	switch curr {
	case stateComment, stateNil:
		// pass; nothing to flush
	case stateValue:
		if _, ok := unquote(b.String()); !ok {
			return nil, io.ErrUnexpectedEOF
112
		}
113

Michael Yang's avatar
Michael Yang committed
114
115
116
117
		cmd.Args = b.String()
		cmds = append(cmds, cmd)
	default:
		return nil, io.ErrUnexpectedEOF
118
119
	}

Michael Yang's avatar
Michael Yang committed
120
121
122
123
	for _, cmd := range cmds {
		if cmd.Name == "model" {
			return cmds, nil
		}
124
125
	}

Michael Yang's avatar
Michael Yang committed
126
	return nil, errors.New("no FROM line")
127
}
128

Michael Yang's avatar
Michael Yang committed
129
130
131
132
133
134
135
136
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
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("")
185
	}
Michael Yang's avatar
Michael Yang committed
186
}
187

Michael Yang's avatar
Michael Yang committed
188
189
190
func unquote(s string) (string, bool) {
	if len(s) == 0 {
		return "", false
191
	}
192

Michael Yang's avatar
Michael Yang committed
193
194
195
196
197
198
199
	// 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
200
201
	}

Michael Yang's avatar
Michael Yang committed
202
203
204
205
206
207
	if len(s) >= 1 && s[0] == '"' {
		if len(s) >= 2 && s[len(s)-1] == '"' {
			return s[1 : len(s)-1], true
		}

		return "", false
208
209
	}

Michael Yang's avatar
Michael Yang committed
210
	return s, true
211
212
}

Michael Yang's avatar
Michael Yang committed
213
214
215
func isAlpha(r rune) bool {
	return r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z'
}
216

Michael Yang's avatar
Michael Yang committed
217
218
219
func isNumber(r rune) bool {
	return r >= '0' && r <= '9'
}
220

Michael Yang's avatar
Michael Yang committed
221
222
223
func isSpace(r rune) bool {
	return r == ' ' || r == '\t'
}
Michael Yang's avatar
Michael Yang committed
224

Michael Yang's avatar
Michael Yang committed
225
226
227
func isNewline(r rune) bool {
	return r == '\r' || r == '\n'
}
228

Michael Yang's avatar
Michael Yang committed
229
230
func isValidRole(role string) bool {
	return role == "system" || role == "user" || role == "assistant"
231
}