Added some image support

Author: djthorpe

attachment.go

@@ -8,11 +8,13 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
+// General attachment metadata
 type AttachmentMeta struct {
 	Id        string `json:"id,omitempty"`
 	Filename  string `json:"filename,omitempty"`
@@ -21,9 +23,17 @@ type AttachmentMeta struct {
 	Data      []byte `json:"data"`
 }
 
+// OpenAI image metadata
+type ImageMeta struct {
+	Url    string `json:"url,omitempty"`
+	Data   []byte `json:"b64_json,omitempty"`
+	Prompt string `json:"revised_prompt,omitempty"`
+}
+
 // Attachment for messages
 type Attachment struct {
-	meta AttachmentMeta
+	meta  *AttachmentMeta
+	image *ImageMeta
 }
 
 const (
@@ -38,6 +48,11 @@ func NewAttachment() *Attachment {
 	return new(Attachment)
 }
 
+// NewAttachment with OpenAI image
+func NewAttachmentWithImage(image *ImageMeta) *Attachment {
+	return &Attachment{image: image}
+}
+
 // ReadAttachment returns an attachment from a reader object.
 // It is the responsibility of the caller to close the reader.
 func ReadAttachment(r io.Reader) (*Attachment, error) {
@@ -50,7 +65,7 @@ func ReadAttachment(r io.Reader) (*Attachment, error) {
 		filename = f.Name()
 	}
 	return &Attachment{
-		meta: AttachmentMeta{
+		meta: &AttachmentMeta{
 			Filename: filename,
 			Data:     data,
 		},
@@ -73,19 +88,25 @@ func (a *Attachment) MarshalJSON() ([]byte, error) {
 		Filename string `json:"filename,omitempty"`
 		Type     string `json:"type"`
 		Bytes    uint64 `json:"bytes"`
-		Caption  string `json:"transcript,omitempty"`
+		Caption  string `json:"caption,omitempty"`
 	}
-	j.Id = a.meta.Id
-	j.Filename = a.meta.Filename
+
 	j.Type = a.Type()
-	j.Bytes = uint64(len(a.meta.Data))
-	j.Caption = a.meta.Caption
+	j.Caption = a.Caption()
+	if a.meta != nil {
+		j.Id = a.meta.Id
+		j.Filename = a.meta.Filename
+		j.Bytes = uint64(len(a.meta.Data))
+	} else if a.image != nil {
+		j.Bytes = uint64(len(a.image.Data))
+	}
+
 	return json.Marshal(j)
 }
 
 // Stringify an attachment
 func (a *Attachment) String() string {
-	data, err := json.MarshalIndent(a.meta, "", "  ")
+	data, err := json.MarshalIndent(a, "", "  ")
 	if err != nil {
 		return err.Error()
 	}
@@ -95,68 +116,98 @@ func (a *Attachment) String() string {
 ////////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
+// Write out attachment
+func (a *Attachment) Write(w io.Writer) (int, error) {
+	if a.meta != nil {
+		return w.Write(a.meta.Data)
+	}
+	if a.image != nil {
+		return w.Write(a.image.Data)
+	}
+	return 0, io.EOF
+}
+
 // Return the filename of an attachment
 func (a *Attachment) Filename() string {
-	return a.meta.Filename
+	if a.meta != nil {
+		return a.meta.Filename
+	} else {
+		return ""
+	}
 }
 
 // Return the raw attachment data
 func (a *Attachment) Data() []byte {
-	return a.meta.Data
+	if a.meta != nil {
+		return a.meta.Data
+	}
+	if a.image != nil {
+		return a.image.Data
+	}
+	return nil
 }
 
 // Return the caption for the attachment
 func (a *Attachment) Caption() string {
-	return a.meta.Caption
+	if a.meta != nil {
+		return strings.TrimSpace(a.meta.Caption)
+	}
+	if a.image != nil {
+		return strings.TrimSpace(a.image.Prompt)
+	}
+	return ""
 }
 
 // Return the mime media type for the attachment, based
 // on the data and/or filename extension. Returns an empty string if
 // there is no data or filename
 func (a *Attachment) Type() string {
 	// If there's no data or filename, return empty
-	if len(a.meta.Data) == 0 && a.meta.Filename == "" {
+	if len(a.Data()) == 0 && a.Filename() == "" {
 		return ""
 	}
 
 	// Mimetype based on content
 	mimetype := defaultMimetype
-	if len(a.meta.Data) > 0 {
-		mimetype = http.DetectContentType(a.meta.Data)
+	if len(a.Data()) > 0 {
+		mimetype = http.DetectContentType(a.Data())
 		if mimetype != defaultMimetype {
 			return mimetype
 		}
 	}
 
 	// Mimetype based on filename
-	if a.meta.Filename != "" {
+	if a.Filename() != "" {
 		// Detect mimetype from extension
-		mimetype = mime.TypeByExtension(filepath.Ext(a.meta.Filename))
+		mimetype = mime.TypeByExtension(filepath.Ext(a.Filename()))
 	}
 
 	// Return the default mimetype
 	return mimetype
 }
 
 func (a *Attachment) Url() string {
-	return "data:" + a.Type() + ";base64," + base64.StdEncoding.EncodeToString(a.meta.Data)
+	return "data:" + a.Type() + ";base64," + base64.StdEncoding.EncodeToString(a.Data())
 }
 
 // Streaming includes the ability to append data
 func (a *Attachment) Append(other *Attachment) {
-	if other.meta.Id != "" {
-		a.meta.Id = other.meta.Id
-	}
-	if other.meta.Filename != "" {
-		a.meta.Filename = other.meta.Filename
-	}
-	if other.meta.ExpiresAt != 0 {
-		a.meta.ExpiresAt = other.meta.ExpiresAt
-	}
-	if other.meta.Caption != "" {
-		a.meta.Caption += other.meta.Caption
-	}
-	if len(other.meta.Data) > 0 {
-		a.meta.Data = append(a.meta.Data, other.meta.Data...)
+	if a.meta != nil {
+		if other.meta.Id != "" {
+			a.meta.Id = other.meta.Id
+		}
+		if other.meta.Filename != "" {
+			a.meta.Filename = other.meta.Filename
+		}
+		if other.meta.ExpiresAt != 0 {
+			a.meta.ExpiresAt = other.meta.ExpiresAt
+		}
+		if other.meta.Caption != "" {
+			a.meta.Caption += other.meta.Caption
+		}
+		if len(other.meta.Data) > 0 {
+			a.meta.Data = append(a.meta.Data, other.meta.Data...)
+		}
 	}
+	// TODO: Append for image
 }

cmd/llm/chat.go

@@ -26,7 +26,7 @@ type ChatCmd struct {
 // PUBLIC METHODS
 
 func (cmd *ChatCmd) Run(globals *Globals) error {
-	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+	return run(globals, AudioType, cmd.Model, func(ctx context.Context, model llm.Model) error {
 		// Current buffer
 		var buf string
 

cmd/llm/chat2.go

@@ -57,7 +57,7 @@ func NewTelegramServer(token string, model llm.Model, toolkit llm.ToolKit, opts
 // PUBLIC METHODS
 
 func (cmd *Chat2Cmd) Run(globals *Globals) error {
-	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+	return run(globals, ChatType, cmd.Model, func(ctx context.Context, model llm.Model) error {
 		server, err := NewTelegramServer(cmd.Token, model, globals.toolkit, telegram.WithDebug(globals.Debug))
 		if err != nil {
 			return err

cmd/llm/complete.go

@@ -9,26 +9,41 @@ import (
 
 	// Packages
 	llm "github.com/mutablelogic/go-llm"
+	"github.com/mutablelogic/go-llm/pkg/openai"
 )
 
 ////////////////////////////////////////////////////////////////////////////////
 // TYPES
 
 type CompleteCmd struct {
-	Model       string   `arg:"" help:"Model name"`
 	Prompt      string   `arg:"" optional:"" help:"Prompt"`
+	Model       string   `flag:"model" help:"Model name"`
 	File        []string `type:"file" short:"f" help:"Files to attach"`
 	System      string   `flag:"system" help:"Set the system prompt"`
 	NoStream    bool     `flag:"no-stream" help:"Do not stream output"`
-	Format      string   `flag:"format" enum:"text,markdown,json" default:"text" help:"Output format"`
+	Format      string   `flag:"format" enum:"text,markdown,json,image,audio" default:"text" help:"Output format"`
+	Size        string   `flag:"size" enum:"256x256,512x512,1024x1024,1792x1024,1024x1792" default:"1024x1024" help:"Image size"`
+	Style       string   `flag:"style" enum:"vivid,natural" default:"vivid" help:"Image style"`
+	Quality     string   `flag:"quality" enum:"standard,hd" default:"standard" help:"Image quality"`
 	Temperature *float64 `flag:"temperature" short:"t"  help:"Temperature for sampling"`
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
+func typeFromFormat(format string) Type {
+	switch format {
+	case "image":
+		return ImageType
+	case "audio":
+		return AudioType
+	default:
+		return TextType
+	}
+}
+
 func (cmd *CompleteCmd) Run(globals *Globals) error {
-	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+	return run(globals, typeFromFormat(cmd.Format), cmd.Model, func(ctx context.Context, model llm.Model) error {
 		var prompt []byte
 
 		// If we are pipeline content in via stdin
@@ -76,13 +91,20 @@ func (cmd *CompleteCmd) Run(globals *Globals) error {
 			return err
 		}
 
-		// Print the completion
+		// Print the completion - text
 		if cmd.NoStream {
 			fmt.Println(completion.Text(0))
 		} else {
 			fmt.Println("")
 		}
 
+		// Print the completion - attachments
+		for i := 0; i < completion.Num(); i++ {
+			if attachment := completion.Attachment(i); attachment != nil {
+				fmt.Println(attachment)
+			}
+		}
+
 		// Return success
 		return nil
 	})
@@ -106,9 +128,12 @@ func (cmd *CompleteCmd) opts() []llm.Opt {
 	}
 
 	// Set format
-	if cmd.Format == "json" {
-		opts = append(opts, llm.WithFormat("json"))
-	}
+	opts = append(opts, llm.WithFormat(cmd.Format))
+
+	// Set image parameters
+	opts = append(opts, openai.WithSize(cmd.Size))
+	opts = append(opts, openai.WithStyle(cmd.Style))
+	opts = append(opts, openai.WithQuality(cmd.Quality))
 
 	// Set temperature
 	if cmd.Temperature != nil {