155 lines
3.8 KiB
Go
155 lines
3.8 KiB
Go
// Package scriptorium adapts the external scriptorium CLI.
|
|
package scriptorium
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"time"
|
|
)
|
|
|
|
type CommandRunner interface {
|
|
Run(ctx context.Context, name string, args []string, timeout time.Duration) (CommandResult, error)
|
|
}
|
|
|
|
type CommandResult struct {
|
|
Stdout []byte
|
|
Stderr []byte
|
|
ExitCode int
|
|
}
|
|
|
|
type ExecRunner struct{}
|
|
|
|
func (ExecRunner) Run(ctx context.Context, name string, args []string, timeout time.Duration) (CommandResult, error) {
|
|
runCtx := ctx
|
|
cancel := func() {}
|
|
if timeout > 0 {
|
|
runCtx, cancel = context.WithTimeout(ctx, timeout)
|
|
}
|
|
defer cancel()
|
|
|
|
cmd := exec.CommandContext(runCtx, name, args...)
|
|
var stdout bytes.Buffer
|
|
var stderr bytes.Buffer
|
|
cmd.Stdout = &stdout
|
|
cmd.Stderr = &stderr
|
|
err := cmd.Run()
|
|
result := CommandResult{Stdout: stdout.Bytes(), Stderr: stderr.Bytes(), ExitCode: 0}
|
|
if err == nil {
|
|
return result, nil
|
|
}
|
|
if exitErr, ok := err.(*exec.ExitError); ok {
|
|
result.ExitCode = exitErr.ExitCode()
|
|
return result, nil
|
|
}
|
|
if runCtx.Err() != nil {
|
|
return result, runCtx.Err()
|
|
}
|
|
return result, err
|
|
}
|
|
|
|
type Runner struct {
|
|
Binary string
|
|
ConfigPath string
|
|
Profile string
|
|
Timeout time.Duration
|
|
ExtraArgs []string
|
|
Commands CommandRunner
|
|
}
|
|
|
|
type RenderRequest struct {
|
|
PromptID string
|
|
DataPackagePath string
|
|
}
|
|
|
|
type RenderResult struct {
|
|
Command []string `json:"command"`
|
|
Stdout string `json:"stdout"`
|
|
Stderr string `json:"stderr"`
|
|
ExitCode int `json:"exitCode"`
|
|
}
|
|
|
|
func (r Runner) Render(ctx context.Context, req RenderRequest) (*RenderResult, error) {
|
|
if req.PromptID == "" {
|
|
return nil, fmt.Errorf("prompt id is required")
|
|
}
|
|
if req.DataPackagePath == "" {
|
|
return nil, fmt.Errorf("data package path is required")
|
|
}
|
|
binary := r.Binary
|
|
if binary == "" {
|
|
binary = "scriptorium"
|
|
}
|
|
commands := r.Commands
|
|
if commands == nil {
|
|
commands = ExecRunner{}
|
|
}
|
|
args := r.renderArgs(req)
|
|
commandResult, err := commands.Run(ctx, binary, args, r.Timeout)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("run scriptorium render: %w", err)
|
|
}
|
|
result := &RenderResult{
|
|
Command: append([]string{binary}, args...),
|
|
Stdout: string(commandResult.Stdout),
|
|
Stderr: string(commandResult.Stderr),
|
|
ExitCode: commandResult.ExitCode,
|
|
}
|
|
if commandResult.ExitCode != 0 {
|
|
return result, fmt.Errorf("scriptorium render exited with code %d: %s", commandResult.ExitCode, result.Stderr)
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (r Runner) renderArgs(req RenderRequest) []string {
|
|
args := []string{"render"}
|
|
if r.ConfigPath != "" {
|
|
args = append(args, "--config", r.ConfigPath)
|
|
}
|
|
if r.Profile != "" {
|
|
args = append(args, "--profile", r.Profile)
|
|
}
|
|
args = append(args,
|
|
"--prompt", req.PromptID,
|
|
"--input", "data_package="+req.DataPackagePath,
|
|
"--format", "json",
|
|
)
|
|
args = append(args, r.ExtraArgs...)
|
|
return args
|
|
}
|
|
|
|
func SaveRenderResult(path string, result *RenderResult) error {
|
|
if result == nil {
|
|
return fmt.Errorf("render result is required")
|
|
}
|
|
data, err := json.MarshalIndent(result, "", " ")
|
|
if err != nil {
|
|
return fmt.Errorf("marshal render result: %w", err)
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
|
return fmt.Errorf("create preflight directory %q: %w", filepath.Dir(path), err)
|
|
}
|
|
tmp, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*.tmp")
|
|
if err != nil {
|
|
return fmt.Errorf("create temporary preflight file: %w", err)
|
|
}
|
|
tmpName := tmp.Name()
|
|
defer os.Remove(tmpName)
|
|
|
|
if _, err := tmp.Write(data); err != nil {
|
|
tmp.Close()
|
|
return fmt.Errorf("write temporary preflight file: %w", err)
|
|
}
|
|
if err := tmp.Close(); err != nil {
|
|
return fmt.Errorf("close temporary preflight file: %w", err)
|
|
}
|
|
if err := os.Rename(tmpName, path); err != nil {
|
|
return fmt.Errorf("save preflight %q: %w", path, err)
|
|
}
|
|
return nil
|
|
}
|