Updated llm

Author: djthorpe

cmd/llm/complete.go

@@ -0,0 +1,102 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	// Packages
+	llm "github.com/mutablelogic/go-llm"
+	agent "github.com/mutablelogic/go-llm/pkg/agent"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type CompleteCmd struct {
+	Model    string   `arg:"" help:"Model name"`
+	Prompt   string   `arg:"" optional:"" help:"Prompt"`
+	File     []string `type:"file" help:"Files to attach"`
+	System   string   `flag:"system" help:"Set the system prompt"`
+	NoStream bool     `flag:"no-stream" help:"Do not stream output"`
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func (cmd *CompleteCmd) Run(globals *Globals) error {
+	return runagent(globals, func(ctx context.Context, client llm.Agent) error {
+		var prompt []byte
+
+		// Load the model
+		model, err := client.(*agent.Agent).GetModel(ctx, cmd.Model)
+		if err != nil {
+			return err
+		}
+
+		// 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{}
+	if cmd.System != "" {
+		opts = append(opts, llm.WithSystemPrompt(cmd.System))
+	}
+	return opts
+}

cmd/llm/main.go

@@ -6,6 +6,7 @@ import (
 	"os/signal"
 	"path/filepath"
 	"syscall"
+	"time"
 
 	// Packages
 	kong "github.com/alecthomas/kong"
@@ -21,8 +22,9 @@ import (
 
 type Globals struct {
 	// Debugging
-	Debug   bool `name:"debug" help:"Enable debug output"`
-	Verbose bool `name:"verbose" help:"Enable verbose output"`
+	Debug   bool          `name:"debug" help:"Enable debug output"`
+	Verbose bool          `name:"verbose" help:"Enable verbose output"`
+	Timeout time.Duration `name:"timeout" help:"Timeout for the command"`
 
 	// Agents
 	Ollama    `embed:"" help:"Ollama configuration"`
@@ -71,6 +73,7 @@ type CLI struct {
 	// Commands
 	Download DownloadModelCmd `cmd:"" help:"Download a model"`
 	Chat     ChatCmd          `cmd:"" help:"Start a chat session"`
+	Complete CompleteCmd      `cmd:"" help:"Complete a prompt"`
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -106,13 +109,16 @@ func main() {
 	if cli.Debug || cli.Verbose {
 		clientopts = append(clientopts, client.OptTrace(os.Stderr, cli.Verbose))
 	}
+	if cli.Timeout > 0 {
+		clientopts = append(clientopts, client.OptTimeout(cli.Timeout))
+	}
 
 	// Create an agent
 	opts := []llm.Opt{}
+	if cli.OllamaEndpoint != "" {
+		opts = append(opts, agent.WithOllama(cli.OllamaEndpoint, clientopts...))
+	}
 	/*
-		if cli.OllamaEndpoint != "" {
-			opts = append(opts, agent.WithOllama(cli.OllamaEndpoint, clientopts...))
-		}
 		if cli.AnthropicKey != "" {
 			opts = append(opts, agent.WithAnthropic(cli.AnthropicKey, clientopts...))
 		}

error.go

@@ -13,6 +13,7 @@ const (
 	ErrBadParameter
 	ErrNotImplemented
 	ErrConflict
+	ErrInternalServerError
 )
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -36,6 +37,8 @@ func (e Err) Error() string {
 		return "not implemented"
 	case ErrConflict:
 		return "conflict"
+	case ErrInternalServerError:
+		return "internal server error"
 	}
 	return fmt.Sprintf("error code %d", int(e))
 }

pkg/agent/opt.go

@@ -4,12 +4,13 @@ import (
 	// Packages
 	client "github.com/mutablelogic/go-client"
 	llm "github.com/mutablelogic/go-llm"
+	"github.com/mutablelogic/go-llm/pkg/ollama"
 	openai "github.com/mutablelogic/go-llm/pkg/openai"
 )
 
 ////////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
-/*
+
 func WithOllama(endpoint string, opts ...client.ClientOpt) llm.Opt {
 	return func(o *llm.Opts) error {
 		client, err := ollama.New(endpoint, opts...)
@@ -21,6 +22,7 @@ func WithOllama(endpoint string, opts ...client.ClientOpt) llm.Opt {
 	}
 }
 
+/*
 func WithAnthropic(key string, opts ...client.ClientOpt) llm.Opt {
 	return func(o *llm.Opts) error {
 		client, err := anthropic.New(key, opts...)

pkg/ollama/chat.go

@@ -1,7 +1,6 @@
 package ollama
 
 import (
-
 	// Packages
 	client "github.com/mutablelogic/go-client"
 	llm "github.com/mutablelogic/go-llm"
@@ -14,7 +13,6 @@ type Client struct {
 	*client.Client
 }
 
-// Ensure it satisfies the agent.Agent interface
 var _ llm.Agent = (*Client)(nil)
 
 ///////////////////////////////////////////////////////////////////////////////