parser.go 2.38 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
	"log"
10
11
12
13
)

type Command struct {
	Name string
14
15
16
17
18
19
	Args string
}

func (c *Command) Reset() {
	c.Name = ""
	c.Args = ""
20
21
22
23
}

func Parse(reader io.Reader) ([]Command, error) {
	var commands []Command
24
	var command, modelCommand Command
25
26

	scanner := bufio.NewScanner(reader)
27
	scanner.Split(scanModelfile)
28
	for scanner.Scan() {
29
30
31
		line := scanner.Bytes()

		fields := bytes.SplitN(line, []byte(" "), 2)
Michael Yang's avatar
Michael Yang committed
32
		if len(fields) == 0 || len(fields[0]) == 0 {
33
34
35
			continue
		}

36
		switch string(bytes.ToUpper(fields[0])) {
37
38
		case "FROM":
			command.Name = "model"
39
40
41
			command.Args = string(fields[1])
			// copy command for validation
			modelCommand = command
Jeffrey Morgan's avatar
Jeffrey Morgan committed
42
		case "LICENSE", "TEMPLATE", "SYSTEM", "PROMPT":
43
			command.Name = string(bytes.ToLower(fields[0]))
44
			command.Args = string(fields[1])
45
		case "PARAMETER":
46
47
48
			fields = bytes.SplitN(fields[1], []byte(" "), 2)
			command.Name = string(fields[0])
			command.Args = string(fields[1])
49
		default:
50
			// log a warning for unknown commands
Michael Yang's avatar
Michael Yang committed
51
			log.Printf("WARNING: Unknown command: %s", fields[0])
52
			continue
53
		}
54
55
56

		commands = append(commands, command)
		command.Reset()
57
58
	}

59
	if modelCommand.Args == "" {
60
		return nil, errors.New("no FROM line for the model was specified")
61
62
63
64
	}

	return commands, scanner.Err()
}
65
66

func scanModelfile(data []byte, atEOF bool) (advance int, token []byte, err error) {
67
68
69
70
71
72
73
74
	advance, token, err = scan([]byte(`"""`), []byte(`"""`), data, atEOF)
	if err != nil {
		return 0, nil, err
	}

	if advance > 0 && token != nil {
		return advance, token, nil
	}
75

76
77
78
79
80
81
82
83
84
85
86
87
88
	advance, token, err = scan([]byte(`"`), []byte(`"`), data, atEOF)
	if err != nil {
		return 0, nil, err
	}

	if advance > 0 && token != nil {
		return advance, token, nil
	}

	return bufio.ScanLines(data, atEOF)
}

func scan(openBytes, closeBytes, data []byte, atEOF bool) (advance int, token []byte, err error) {
89
90
	newline := bytes.IndexByte(data, '\n')

91
92
	if start := bytes.Index(data, openBytes); start >= 0 && start < newline {
		end := bytes.Index(data[start+len(openBytes):], closeBytes)
93
		if end < 0 {
94
			if atEOF {
95
				return 0, nil, fmt.Errorf("unterminated %s: expecting %s", openBytes, closeBytes)
96
97
98
			} else {
				return 0, nil, nil
			}
99
100
		}

101
		n := start + len(openBytes) + end + len(closeBytes)
Michael Yang's avatar
Michael Yang committed
102
103

		newData := data[:start]
104
		newData = append(newData, data[start+len(openBytes):n-len(closeBytes)]...)
Michael Yang's avatar
Michael Yang committed
105
		return n, newData, nil
106
107
	}

108
	return 0, nil, nil
109
}