126 lines
3.5 KiB
Go
126 lines
3.5 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"g.fc.run/ai/byre/internal/codex"
|
|
"g.fc.run/ai/byre/internal/mcp"
|
|
)
|
|
|
|
const imageGenGuidance = "You make images for another agent. Call your image_gen tool exactly once with the prompt you are given, unchanged, then reply with only the word done."
|
|
|
|
func (m *Manager) generateImageTool(l *live) mcp.Tool {
|
|
return mcp.Tool{
|
|
Name: "generate_image",
|
|
Description: "Generate an image (a photo, illustration, texture, sprite, mockup) with OpenAI's image model through byre's Codex sign-in, saved as a PNG in the task's media folder. " +
|
|
"The reply gives its path: show it inline by embedding it in your final message as a Markdown image, , or copy it into the worktree when the project uses it. " +
|
|
"Use it whenever the human asks for a picture, rather than drawing one in SVG or HTML. It takes a minute or two.",
|
|
InputSchema: json.RawMessage(`{
|
|
"type": "object",
|
|
"properties": {
|
|
"prompt": {"type": "string", "description": "What to draw: subject, style, composition, lighting, and any exact text."}
|
|
},
|
|
"required": ["prompt"]
|
|
}`),
|
|
Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
|
|
return m.generateImage(ctx, l, raw)
|
|
},
|
|
}
|
|
}
|
|
|
|
func (m *Manager) generateImage(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
|
|
args := struct {
|
|
Prompt string `json:"prompt"`
|
|
}{}
|
|
err := json.Unmarshal(raw, &args)
|
|
if err != nil {
|
|
return "", fmt.Errorf("malformed generate_image call: %w", err)
|
|
}
|
|
if strings.TrimSpace(args.Prompt) == "" {
|
|
return "", fmt.Errorf("generate_image needs a prompt")
|
|
}
|
|
|
|
media := MediaDir(l.dir)
|
|
err = os.MkdirAll(m.codexHome, 0o700)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
s, err := codex.Start(ctx, codex.Options{
|
|
Cwd: media,
|
|
Env: []string{"CODEX_HOME=" + m.codexHome},
|
|
Guidance: imageGenGuidance,
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("start codex: %w", err)
|
|
}
|
|
defer s.Close()
|
|
|
|
_, err = s.Send(ctx, args.Prompt)
|
|
if err != nil {
|
|
return "", fmt.Errorf("start codex turn: %w", err)
|
|
}
|
|
|
|
saved := ""
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
s.Kill()
|
|
return "", fmt.Errorf("no image after %s", mcp.Timeout(ctx))
|
|
case ev, ok := <-s.Events():
|
|
if !ok {
|
|
return "", fmt.Errorf("codex exited before the image was made")
|
|
}
|
|
switch ev.Method {
|
|
case codex.MethodItemCompleted:
|
|
n := codex.ItemCompletedNotification{}
|
|
err := json.Unmarshal(ev.Params, &n)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if n.Item.Type != codex.ThreadItemTypeImageGeneration {
|
|
continue
|
|
}
|
|
if n.Item.Failure != nil {
|
|
return "", fmt.Errorf("image generation failed: %s (limit %s)", n.Item.Failure.Type, n.Item.Failure.LimitID)
|
|
}
|
|
saved = string(n.Item.SavedPath)
|
|
case codex.MethodTurnCompleted:
|
|
n := codex.TurnCompletedNotification{}
|
|
err := json.Unmarshal(ev.Params, &n)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if n.Turn.Error != nil {
|
|
return "", fmt.Errorf("codex: %s", n.Turn.Error.Message)
|
|
}
|
|
if saved == "" {
|
|
return "", fmt.Errorf("codex finished without making an image")
|
|
}
|
|
path, err := copyToMedia(media, saved)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return fmt.Sprintf("Saved %s", path), nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func copyToMedia(media, src string) (string, error) {
|
|
data, err := os.ReadFile(src)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
err = os.MkdirAll(media, 0o755)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
dest := filepath.Join(media, filepath.Base(src))
|
|
return dest, os.WriteFile(dest, data, 0o644)
|
|
}
|