/
optparse.go
101 lines (81 loc) · 2.04 KB
/
optparse.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package util
import (
"bufio"
"context"
"encoding/json"
"fmt"
"io"
"os"
"regexp"
"strings"
"github.com/PaesslerAG/jsonpath"
"github.com/replicate/replicate-go"
)
func ParseInputs(ctx context.Context, r8 *replicate.Client, args []string, stdin string, sep string) (map[string]string, error) {
re := regexp.MustCompile(`{{(.*?)}}`)
inputs := make(map[string]string)
for _, e := range args {
k, v, found := strings.Cut(e, sep)
if !found {
return nil, fmt.Errorf("invalid input: %s", e)
}
var stdinJSON map[string]interface{}
if stdin != "" {
err := json.Unmarshal([]byte(stdin), &stdinJSON)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal stdin: %w", err)
}
}
// Extract data from JSON
matches := re.FindAllStringSubmatch(v, -1)
for _, match := range matches {
if len(match) < 2 {
continue
}
path := strings.TrimSpace(match[1])
if !strings.HasPrefix(path, "$") {
path = "$" + path
}
value, err := jsonpath.Get(path, stdinJSON)
if err != nil {
return nil, fmt.Errorf("failed to extract data from JSON using path '%s': %w", path, err)
}
// Replace the segment with the extracted value
v = strings.Replace(v, match[0], fmt.Sprintf("%v", value), 1)
}
// Read from file
if strings.HasPrefix(v, "@") {
path := strings.TrimSpace(v[1:])
file, err := r8.CreateFileFromPath(ctx, path, nil)
if err != nil {
return nil, fmt.Errorf("failed to create file from path: %w", err)
}
downloadURL := file.URLs["get"]
if downloadURL == "" {
return nil, fmt.Errorf("failed to get download URL for file")
}
v = downloadURL
}
inputs[k] = v
}
return inputs, nil
}
func GetPipedArgs() (string, error) {
info, err := os.Stdin.Stat()
if err != nil {
return "", err
}
if info.Mode()&os.ModeNamedPipe != 0 {
reader := bufio.NewReader(os.Stdin)
var output []rune
for {
input, _, err := reader.ReadRune()
if err != nil && err == io.EOF {
break
}
output = append(output, input)
}
return string(output), nil
}
return "", nil
}