Merge pull request #7 from mutablelogic/dev

Added new agents

Author: djthorpe

README.md

@@ -13,15 +13,31 @@ import (
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
+type AttachmentMeta struct {
+	Id        string `json:"id,omitempty"`
+	Filename  string `json:"filename,omitempty"`
+	ExpiresAt uint64 `json:"expires_at,omitempty"`
+	Caption   string `json:"transcript,omitempty"`
+	Data      []byte `json:"data"`
+}
+
 // Attachment for messages
 type Attachment struct {
-	filename string
-	data     []byte
+	meta AttachmentMeta
 }
 
+const (
+	defaultMimetype = "application/octet-stream"
+)
+
 ////////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
+// NewAttachment creates a new, empty attachment
+func NewAttachment() *Attachment {
+	return new(Attachment)
+}
+
 // 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) {
@@ -33,22 +49,43 @@ func ReadAttachment(r io.Reader) (*Attachment, error) {
 	if f, ok := r.(*os.File); ok {
 		filename = f.Name()
 	}
-	return &Attachment{filename: filename, data: data}, nil
+	return &Attachment{
+		meta: AttachmentMeta{
+			Filename: filename,
+			Data:     data,
+		},
+	}, nil
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // STRINGIFY
 
-func (a *Attachment) String() string {
+// Convert JSON into an attachment
+func (a *Attachment) UnmarshalJSON(data []byte) error {
+	return json.Unmarshal(data, &a.meta)
+}
+
+// Convert an attachment into JSON
+func (a *Attachment) MarshalJSON() ([]byte, error) {
+	// Create a JSON representation
 	var j struct {
-		Filename string `json:"filename"`
+		Id       string `json:"id,omitempty"`
+		Filename string `json:"filename,omitempty"`
 		Type     string `json:"type"`
 		Bytes    uint64 `json:"bytes"`
+		Caption  string `json:"transcript,omitempty"`
 	}
-	j.Filename = a.filename
+	j.Id = a.meta.Id
+	j.Filename = a.meta.Filename
 	j.Type = a.Type()
-	j.Bytes = uint64(len(a.data))
-	data, err := json.MarshalIndent(j, "", "  ")
+	j.Bytes = uint64(len(a.meta.Data))
+	j.Caption = a.meta.Caption
+	return json.Marshal(j)
+}
+
+// Stringify an attachment
+func (a *Attachment) String() string {
+	data, err := json.MarshalIndent(a.meta, "", "  ")
 	if err != nil {
 		return err.Error()
 	}
@@ -58,24 +95,68 @@ func (a *Attachment) String() string {
 ////////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
+// Return the filename of an attachment
 func (a *Attachment) Filename() string {
-	return a.filename
+	return a.meta.Filename
 }
 
+// Return the raw attachment data
 func (a *Attachment) Data() []byte {
-	return a.data
+	return a.meta.Data
+}
+
+// Return the caption for the attachment
+func (a *Attachment) Caption() string {
+	return a.meta.Caption
 }
 
+// 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 == "" {
+		return ""
+	}
+
 	// Mimetype based on content
-	mimetype := http.DetectContentType(a.data)
-	if mimetype == "application/octet-stream" && a.filename != "" {
+	mimetype := defaultMimetype
+	if len(a.meta.Data) > 0 {
+		mimetype = http.DetectContentType(a.meta.Data)
+		if mimetype != defaultMimetype {
+			return mimetype
+		}
+	}
+
+	// Mimetype based on filename
+	if a.meta.Filename != "" {
 		// Detect mimetype from extension
-		mimetype = mime.TypeByExtension(filepath.Ext(a.filename))
+		mimetype = mime.TypeByExtension(filepath.Ext(a.meta.Filename))
 	}
+
+	// Return the default mimetype
 	return mimetype
 }
 
 func (a *Attachment) Url() string {
-	return "data:" + a.Type() + ";base64," + base64.StdEncoding.EncodeToString(a.data)
+	return "data:" + a.Type() + ";base64," + base64.StdEncoding.EncodeToString(a.meta.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...)
+	}
 }

attachment.go

@@ -9,7 +9,6 @@ import (
 
 	// Packages
 	llm "github.com/mutablelogic/go-llm"
-	agent "github.com/mutablelogic/go-llm/pkg/agent"
 )
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -27,26 +26,17 @@ type ChatCmd struct {
 // PUBLIC METHODS
 
 func (cmd *ChatCmd) Run(globals *Globals) error {
-	return runagent(globals, func(ctx context.Context, client llm.Agent) error {
-		// Get the model
-		a, ok := client.(*agent.Agent)
-		if !ok {
-			return fmt.Errorf("No agents found")
-		}
-		model, err := a.GetModel(ctx, cmd.Model)
-		if err != nil {
-			return err
-		}
+	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+		// Current buffer
+		var buf string
 
 		// Set the options
 		opts := []llm.Opt{}
 		if !cmd.NoStream {
 			opts = append(opts, llm.WithStream(func(cc llm.Completion) {
-				if text := cc.Text(0); text != "" {
-					count := strings.Count(text, "\n")
-					fmt.Print(strings.Repeat("\033[F", count) + strings.Repeat(" ", count) + "\r")
-					fmt.Print(text)
-				}
+				text := cc.Text(0)
+				fmt.Print(strings.TrimPrefix(text, buf))
+				buf = text
 			}))
 		}
 		if cmd.System != "" {
@@ -66,6 +56,7 @@ func (cmd *ChatCmd) Run(globals *Globals) error {
 				input = cmd.Prompt
 				cmd.Prompt = ""
 			} else {
+				var err error
 				input, err = globals.term.ReadLine(model.Name() + "> ")
 				if errors.Is(err, io.EOF) {
 					return nil
@@ -91,6 +82,7 @@ func (cmd *ChatCmd) Run(globals *Globals) error {
 				if len(calls) == 0 {
 					break
 				}
+
 				if session.Text(0) != "" {
 					globals.term.Println(session.Text(0))
 				} else {
@@ -100,15 +92,20 @@ func (cmd *ChatCmd) Run(globals *Globals) error {
 					}
 					globals.term.Println("Calling ", strings.Join(names, ", "))
 				}
+
 				if results, err := globals.toolkit.Run(ctx, calls...); err != nil {
 					return err
 				} else if err := session.FromTool(ctx, results...); err != nil {
 					return err
 				}
 			}
 
-			// Print the response
-			globals.term.Println("\n" + session.Text(0) + "\n")
+			// Print the response, if not streaming
+			if cmd.NoStream {
+				globals.term.Println("\n" + session.Text(0) + "\n")
+			} else {
+				globals.term.Println()
+			}
 		}
 	})
 }

cmd/llm/chat.go

@@ -0,0 +1,119 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	// Packages
+	llm "github.com/mutablelogic/go-llm"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type CompleteCmd struct {
+	Model       string   `arg:"" help:"Model name"`
+	Prompt      string   `arg:"" optional:"" help:"Prompt"`
+	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"`
+	Temperature *float64 `flag:"temperature" short:"t"  help:"Temperature for sampling"`
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func (cmd *CompleteCmd) Run(globals *Globals) error {
+	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+		var prompt []byte
+
+		// If we are pipeline content in via stdin
+		fileInfo, err := os.Stdin.Stat()
+		if err != nil {
+			return llm.ErrInternalServerError.Withf("Failed to get stdin stat: %v", err)
+		}
+		if (fileInfo.Mode() & os.ModeCharDevice) == 0 {
+			if data, err := io.ReadAll(os.Stdin); err != nil {
+				return err
+			} else if len(data) > 0 {
+				prompt = data
+			}
+		}
+
+		// Append any further prompt
+		if len(cmd.Prompt) > 0 {
+			prompt = append(prompt, []byte("\n\n")...)
+			prompt = append(prompt, []byte(cmd.Prompt)...)
+		}
+
+		opts := cmd.opts()
+		if !cmd.NoStream {
+			// Add streaming callback
+			var buf string
+			opts = append(opts, llm.WithStream(func(c llm.Completion) {
+				fmt.Print(strings.TrimPrefix(c.Text(0), buf))
+				buf = c.Text(0)
+			}))
+		}
+
+		// Add attachments
+		for _, file := range cmd.File {
+			f, err := os.Open(file)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			opts = append(opts, llm.WithAttachment(f))
+		}
+
+		// Make the completion
+		completion, err := model.Completion(ctx, string(prompt), opts...)
+		if err != nil {
+			return err
+		}
+
+		// Print the completion
+		if cmd.NoStream {
+			fmt.Println(completion.Text(0))
+		} else {
+			fmt.Println("")
+		}
+
+		// Return success
+		return nil
+	})
+}
+
+func (cmd *CompleteCmd) opts() []llm.Opt {
+	opts := []llm.Opt{}
+
+	// Set system prompt
+	var system []string
+	if cmd.Format == "markdown" {
+		system = append(system, "Structure your output in markdown format.")
+	} else if cmd.Format == "json" {
+		system = append(system, "Structure your output in JSON format.")
+	}
+	if cmd.System != "" {
+		system = append(system, cmd.System)
+	}
+	if len(system) > 0 {
+		opts = append(opts, llm.WithSystemPrompt(strings.Join(system, "\n")))
+	}
+
+	// Set format
+	if cmd.Format == "json" {
+		opts = append(opts, llm.WithFormat("json"))
+	}
+
+	// Set temperature
+	if cmd.Temperature != nil {
+		opts = append(opts, llm.WithTemperature(*cmd.Temperature))
+	}
+
+	return opts
+}

cmd/llm/complete.go

@@ -0,0 +1,36 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	// Packages
+	llm "github.com/mutablelogic/go-llm"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type EmbeddingCmd struct {
+	Model  string `arg:"" help:"Model name"`
+	Prompt string `arg:"" help:"Prompt"`
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func (cmd *EmbeddingCmd) Run(globals *Globals) error {
+	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+		vector, err := model.Embedding(ctx, cmd.Prompt)
+		if err != nil {
+			return err
+		}
+		data, err := json.Marshal(vector)
+		if err != nil {
+			return err
+		}
+		fmt.Println(string(data))
+		return nil
+	})
+}