Skip to content

Commit

Permalink
fix parser name
Browse files Browse the repository at this point in the history
  • Loading branch information
mxyng committed May 1, 2024
1 parent 9cf0f2e commit bd8eed5
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 6 deletions.
26 changes: 21 additions & 5 deletions parser/parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,9 @@ const (
)

var (
errMissingFrom = errors.New("no FROM line")
errInvalidRole = errors.New("role must be one of \"system\", \"user\", or \"assistant\"")
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\"")
)

func Format(cmds []Command) string {
Expand Down Expand Up @@ -82,7 +83,11 @@ func Parse(r io.Reader) (cmds []Command, err error) {
// process the state transition, some transitions need to be intercepted and redirected
if next != curr {
switch curr {
case stateName, stateParameter:
case stateName:
if !isValidCommand(b.String()) {
return nil, errInvalidCommand
}

// next state sometimes depends on the current buffer value
switch s := strings.ToLower(b.String()); s {
case "from":
Expand All @@ -97,9 +102,11 @@ func Parse(r io.Reader) (cmds []Command, err error) {
default:
cmd.Name = s
}
case stateParameter:
cmd.Name = b.String()
case stateMessage:
if !isValidMessageRole(b.String()) {
return nil, errInvalidRole
return nil, errInvalidMessageRole
}

role = b.String()
Expand Down Expand Up @@ -182,7 +189,7 @@ func parseRuneForState(r rune, cs state) (state, rune, error) {
case isSpace(r):
return stateValue, 0, nil
default:
return stateNil, 0, errors.New("invalid")
return stateNil, 0, errInvalidCommand
}
case stateValue:
switch {
Expand Down Expand Up @@ -279,3 +286,12 @@ func isNewline(r rune) bool {
func isValidMessageRole(role string) bool {
return role == "system" || role == "user" || role == "assistant"
}

func isValidCommand(cmd string) bool {
switch strings.ToLower(cmd) {
case "from", "license", "template", "system", "adapter", "parameter", "message":
return true
default:
return false
}
}
12 changes: 11 additions & 1 deletion parser/parser_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,16 @@ PARAMETER param1
assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
}

func TestParserBadCommand(t *testing.T) {
input := `
FROM foo
BADCOMMAND param1 value1
`
_, err := Parse(strings.NewReader(input))
assert.ErrorIs(t, err, errInvalidCommand)

}

func TestParserMessages(t *testing.T) {
var cases = []struct {
input string
Expand Down Expand Up @@ -165,7 +175,7 @@ FROM foo
MESSAGE badguy I'm a bad guy!
`,
nil,
errInvalidRole,
errInvalidMessageRole,
},
{
`
Expand Down

0 comments on commit bd8eed5

Please sign in to comment.