Browse Source

Revert "parser: remove role validation from Modelfile parser" (#9917)

This reverts commit ffbfe833da387f9b6806fe887b85992c11d26eaa.
Parth Sareen 1 month ago
parent
commit
00ebda8cc4
2 changed files with 16 additions and 12 deletions
  1. 11 7
      parser/parser.go
  2. 5 5
      parser/parser_test.go

+ 11 - 7
parser/parser.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"log/slog"
 	"net/http"
 	"os"
 	"os/user"
@@ -301,8 +300,9 @@ const (
 )
 
 var (
-	errMissingFrom    = errors.New("no FROM line")
-	errInvalidCommand = errors.New("command must be one of \"from\", \"license\", \"template\", \"system\", \"adapter\", \"parameter\", or \"message\"")
+	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\"")
 )
 
 type ParserError struct {
@@ -379,10 +379,14 @@ func ParseFile(r io.Reader) (*Modelfile, error) {
 			case stateParameter:
 				cmd.Name = b.String()
 			case stateMessage:
-				role = b.String()
-				if !isKnownMessageRole(b.String()) {
-					slog.Warn("received non-standard role", "role", role)
+				if !isValidMessageRole(b.String()) {
+					return nil, &ParserError{
+						LineNumber: currLine,
+						Msg:        errInvalidMessageRole.Error(),
+					}
 				}
+
+				role = b.String()
 			case stateComment, stateNil:
 				// pass
 			case stateValue:
@@ -552,7 +556,7 @@ func isNewline(r rune) bool {
 	return r == '\r' || r == '\n'
 }
 
-func isKnownMessageRole(role string) bool {
+func isValidMessageRole(role string) bool {
 	return role == "system" || role == "user" || role == "assistant"
 }
 

+ 5 - 5
parser/parser_test.go

@@ -256,13 +256,13 @@ You are a multiline file parser. Always parse things.
 		{
 			`
 FROM foo
-MESSAGE somerandomrole I'm ok with you adding any role message now!
+MESSAGE badguy I'm a bad guy!
 `,
-			[]Command{
-				{Name: "model", Args: "foo"},
-				{Name: "message", Args: "somerandomrole: I'm ok with you adding any role message now!"},
-			},
 			nil,
+			&ParserError{
+				LineNumber: 3,
+				Msg:        errInvalidMessageRole.Error(),
+			},
 		},
 		{
 			`