diff --git a/commands/image.go b/commands/image.go index 360245b..9e4a56a 100644 --- a/commands/image.go +++ b/commands/image.go @@ -5,6 +5,7 @@ import ( "github.com/LordPax/aicli/lang" "github.com/LordPax/aicli/sdk" + "github.com/LordPax/aicli/service" cli "github.com/urfave/cli/v2" ) @@ -19,7 +20,7 @@ func ImageCommand() (*cli.Command, error) { Name: "image", Usage: l.Get("image-usage"), ArgsUsage: "[image|-]", - Aliases: []string{"i"}, + Aliases: []string{"im"}, Action: imageAction, Flags: imageFlags(), }, nil @@ -94,5 +95,6 @@ func imageFlags() []cli.Flag { } func imageAction(c *cli.Context) error { - return nil + prompt := c.Args().First() + return service.SendImageRequest(prompt) } diff --git a/commands/text.go b/commands/text.go index 32327af..2fd7c33 100644 --- a/commands/text.go +++ b/commands/text.go @@ -25,7 +25,7 @@ func TextCommand() (*cli.Command, error) { Name: "text", Usage: l.Get("text-usage"), ArgsUsage: "[prompt|-]", - Aliases: []string{"t"}, + Aliases: []string{"te"}, Action: textAction, Flags: textFlags(), }, nil diff --git a/lang/en.go b/lang/en.go index db44be9..aa5f540 100644 --- a/lang/en.go +++ b/lang/en.go @@ -37,4 +37,5 @@ var EN_STRINGS = LangString{ "image-size-usage": "Set the size of the image", "image-nb-usage": "Set the number of images", "image-output-usage": "Set the name of the output file", + "image-save": "Image saved under \"%s\"\n", } diff --git a/lang/fr.go b/lang/fr.go index 050f355..34997e0 100644 --- a/lang/fr.go +++ b/lang/fr.go @@ -37,4 +37,5 @@ var FR_STRINGS = LangString{ "image-size-usage": "Définir la taille de l'image", "image-nb-usage": "Définir le nombre d'images", "image-output-usage": "Définir le nom du fichier de sortie", + "image-save": "L'image a été enregistrée sous \"%s\"\n", } diff --git a/sdk/image.go b/sdk/image.go index cd470af..366fc61 100644 --- a/sdk/image.go +++ b/sdk/image.go @@ -16,7 +16,7 @@ type IImageService interface { } type ISdkImage interface { - SendRequest(prompt string) (OpenaiImageResponse, error) + SendRequest(prompt string) error SetModel(model string) GetModel() string SetSize(size string) diff --git a/sdk/openai-image.go b/sdk/openai-image.go index cc96a34..bdf0d02 100644 --- a/sdk/openai-image.go +++ b/sdk/openai-image.go @@ -3,9 +3,11 @@ package sdk import ( "encoding/json" "errors" + "fmt" "io" "net/http" + "github.com/LordPax/aicli/lang" "github.com/LordPax/aicli/utils" ) @@ -18,10 +20,10 @@ type OpenaiImageBody struct { } type OpenaiImageResponse struct { - Images []struct { + Data []struct { Url string `json:"url"` B64Json string `json:"b64_json"` - } `json:"images"` + } `json:"data"` } type OpenaiImage struct { @@ -45,12 +47,12 @@ func NewOpenaiImage(apiKey string) (*OpenaiImage, error) { }, nil } -func (o *OpenaiImage) SendRequest(prompt string) (OpenaiImageResponse, error) { +func (o *OpenaiImage) SendRequest(prompt string) error { var openaiResponse OpenaiImageResponse format := "url" if o.GetOutput() != "" { - format = "b4_json" + format = "b64_json" } jsonBody, err := json.Marshal(OpenaiImageBody{ @@ -61,7 +63,7 @@ func (o *OpenaiImage) SendRequest(prompt string) (OpenaiImageResponse, error) { ResponseFormat: format, }) if err != nil { - return OpenaiImageResponse{}, err + return err } resp, err := utils.PostRequest(o.ApiUrl, jsonBody, map[string]string{ @@ -69,26 +71,63 @@ func (o *OpenaiImage) SendRequest(prompt string) (OpenaiImageResponse, error) { "Authorization": "Bearer " + o.ApiKey, }) if err != nil { - return OpenaiImageResponse{}, err + return err } defer resp.Body.Close() respBody, err := io.ReadAll(resp.Body) if err != nil { - return OpenaiImageResponse{}, err + return err } if resp.StatusCode != http.StatusOK { var errorMsg ErrorMsg if err := json.Unmarshal(respBody, &errorMsg); err != nil { - return OpenaiImageResponse{}, err + return err } - return OpenaiImageResponse{}, errors.New(errorMsg.Error.Message) + return errors.New(errorMsg.Error.Message) } if err := json.Unmarshal(respBody, &openaiResponse); err != nil { - return OpenaiImageResponse{}, err + return err } - return openaiResponse, nil + if o.GetOutput() != "" { + if err := convertFile(openaiResponse, o.GetOutput()); err != nil { + return err + } + return nil + } + + for _, image := range openaiResponse.Data { + println(image.Url) + } + + return nil +} + +func convertFile(openaiResponse OpenaiImageResponse, output string) error { + l := lang.GetLocalize() + log, err := utils.GetLog() + if err != nil { + return err + } + + if (len(openaiResponse.Data)) == 1 { + if err := utils.ConvertB64ToImage(output, openaiResponse.Data[0].B64Json); err != nil { + return err + } + log.Printf(l.Get("image-save"), output) + return nil + } + + for i, image := range openaiResponse.Data { + name := fmt.Sprintf("%s-%d", output, i) + if err := utils.ConvertB64ToImage(name, image.B64Json); err != nil { + return err + } + log.Printf(l.Get("image-save"), name) + } + + return nil } diff --git a/service/image.go b/service/image.go new file mode 100644 index 0000000..70f28bf --- /dev/null +++ b/service/image.go @@ -0,0 +1,27 @@ +package service + +import ( + "io" + "os" + + "github.com/LordPax/aicli/sdk" +) + +func SendImageRequest(prompt string) error { + imageSdk := sdk.GetSdkImage() + + if prompt == "-" { + stdin, err := io.ReadAll(os.Stdin) + if err != nil { + return err + } + + prompt = string(stdin) + } + + if err := imageSdk.SendRequest(prompt); err != nil { + return err + } + + return nil +} diff --git a/utils/utils.go b/utils/utils.go index 394d7b0..576af7c 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -2,6 +2,7 @@ package utils import ( "bytes" + "encoding/base64" "io" "net/http" "os" @@ -91,3 +92,12 @@ func GetFileFromUrl(url string) ([]byte, error) { return io.ReadAll(resp.Body) } + +func ConvertB64ToImage(name, b64 string) error { + dec, err := base64.StdEncoding.DecodeString(b64) + if err != nil { + return err + } + + return os.WriteFile(name, dec, 0644) +}