Add genimage action
This commit is contained in:
83
services/actions/genimage.go
Normal file
83
services/actions/genimage.go
Normal 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"},
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user