From a57f9905767a2f876211a314c99816d5b6d3f7d3 Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Sat, 1 Mar 2025 12:18:10 +0100
Subject: [PATCH] Add genimage action

---
 services/actions/genimage.go | 83 ++++++++++++++++++++++++++++++++++++
 webui/actions.go             |  4 ++
 2 files changed, 87 insertions(+)
 create mode 100644 services/actions/genimage.go

diff --git a/services/actions/genimage.go b/services/actions/genimage.go
new file mode 100644
index 0000000..0417632
--- /dev/null
+++ b/services/actions/genimage.go
@@ -0,0 +1,83 @@
+package actions
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/mudler/LocalAgent/core/action"
+	"github.com/sashabaranov/go-openai"
+	"github.com/sashabaranov/go-openai/jsonschema"
+)
+
+func NewGenImage(config map[string]string) *GenImageAction {
+	defaultConfig := openai.DefaultConfig(config["apiKey"])
+	defaultConfig.BaseURL = config["apiURL"]
+
+	return &GenImageAction{
+		client:     openai.NewClientWithConfig(defaultConfig),
+		imageModel: config["model"],
+	}
+}
+
+type GenImageAction struct {
+	client     *openai.Client
+	imageModel string
+}
+
+func (a *GenImageAction) Run(ctx context.Context, params action.ActionParams) (string, error) {
+	result := struct {
+		Prompt string `json:"prompt"`
+		Size   string `json:"size"`
+	}{}
+	err := params.Unmarshal(&result)
+	if err != nil {
+		fmt.Printf("error: %v", err)
+
+		return "", err
+	}
+
+	req := openai.ImageRequest{
+		Prompt: result.Prompt,
+	}
+
+	switch result.Size {
+	case "256x256":
+		req.Size = openai.CreateImageSize256x256
+	case "512x512":
+		req.Size = openai.CreateImageSize512x512
+	case "1024x1024":
+		req.Size = openai.CreateImageSize1024x1024
+	default:
+		req.Size = openai.CreateImageSize256x256
+	}
+
+	resp, err := a.client.CreateImage(ctx, req)
+	if err != nil {
+		return "Failed to generate image " + err.Error(), err
+	}
+
+	if len(resp.Data) == 0 {
+		return "Failed to generate image", nil
+	}
+
+	return resp.Data[0].URL, nil
+}
+
+func (a *GenImageAction) Definition() action.ActionDefinition {
+	return action.ActionDefinition{
+		Name:        "generate_image",
+		Description: "Generate image with.",
+		Properties: map[string]jsonschema.Definition{
+			"prompt": {
+				Type:        jsonschema.String,
+				Description: "The image prompt to generate the image.",
+			},
+			"size": {
+				Type:        jsonschema.String,
+				Description: "The image prompt to generate the image.",
+				Enum:        []string{"256x256", "512x512", "1024x1024"},
+			},
+		},
+		Required: []string{"prompt"},
+	}
+}
diff --git a/webui/actions.go b/webui/actions.go
index dc7677b..c96b23d 100644
--- a/webui/actions.go
+++ b/webui/actions.go
@@ -24,6 +24,7 @@ const (
 	ActionWikipedia           = "wikipedia"
 	ActionBrowse              = "browse"
 	ActionSendMail            = "send_mail"
+	ActionGenerateImage       = "generate_image"
 )
 
 var AvailableActions = []string{
@@ -37,6 +38,7 @@ var AvailableActions = []string{
 	ActionBrowse,
 	ActionWikipedia,
 	ActionSendMail,
+	ActionGenerateImage,
 }
 
 func Actions(a *state.AgentConfig) func(ctx context.Context) []agent.Action {
@@ -58,6 +60,8 @@ func Actions(a *state.AgentConfig) func(ctx context.Context) []agent.Action {
 					continue
 				}
 				allActions = append(allActions, customAction)
+			case ActionGenerateImage:
+				allActions = append(allActions, actions.NewGenImage(config))
 			case ActionSearch:
 				allActions = append(allActions, actions.NewSearch(config))
 			case ActionGithubIssueLabeler: