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"
|
ActionWikipedia = "wikipedia"
|
||||||
ActionBrowse = "browse"
|
ActionBrowse = "browse"
|
||||||
ActionSendMail = "send_mail"
|
ActionSendMail = "send_mail"
|
||||||
|
ActionGenerateImage = "generate_image"
|
||||||
)
|
)
|
||||||
|
|
||||||
var AvailableActions = []string{
|
var AvailableActions = []string{
|
||||||
@@ -37,6 +38,7 @@ var AvailableActions = []string{
|
|||||||
ActionBrowse,
|
ActionBrowse,
|
||||||
ActionWikipedia,
|
ActionWikipedia,
|
||||||
ActionSendMail,
|
ActionSendMail,
|
||||||
|
ActionGenerateImage,
|
||||||
}
|
}
|
||||||
|
|
||||||
func Actions(a *state.AgentConfig) func(ctx context.Context) []agent.Action {
|
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
|
continue
|
||||||
}
|
}
|
||||||
allActions = append(allActions, customAction)
|
allActions = append(allActions, customAction)
|
||||||
|
case ActionGenerateImage:
|
||||||
|
allActions = append(allActions, actions.NewGenImage(config))
|
||||||
case ActionSearch:
|
case ActionSearch:
|
||||||
allActions = append(allActions, actions.NewSearch(config))
|
allActions = append(allActions, actions.NewSearch(config))
|
||||||
case ActionGithubIssueLabeler:
|
case ActionGithubIssueLabeler:
|
||||||
|
|||||||
Reference in New Issue
Block a user