Add genimage action

This commit is contained in:
mudler
2025-03-01 12:18:10 +01:00
parent 76a01994f9
commit a57f990576
2 changed files with 87 additions and 0 deletions

View File

@@ -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"},
}
}

View File

@@ -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: