Merge pull request #174 from ibuildthecloud/models

feat: add support for 3rd party model shims

Author: ibuildthecloud

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.8.0
 	github.com/stretchr/testify v1.8.4
+	golang.org/x/exp v0.0.0-20240103183307-be819d1f06fc
 	golang.org/x/sync v0.6.0
 	golang.org/x/term v0.16.0
 )
@@ -66,7 +67,6 @@ require (
 	github.com/therootcompany/xz v1.0.1 // indirect
 	github.com/ulikunitz/xz v0.5.10 // indirect
 	go4.org v0.0.0-20200411211856-f5505b9728dd // indirect
-	golang.org/x/exp v0.0.0-20240103183307-be819d1f06fc // indirect
 	golang.org/x/mod v0.15.0 // indirect
 	golang.org/x/net v0.20.0 // indirect
 	golang.org/x/sys v0.16.0 // indirect

pkg/cache/cache.go

@@ -58,6 +58,10 @@ func New(opts ...Options) (*Client, error) {
 	}, nil
 }
 
+func (c *Client) CacheDir() string {
+	return c.dir
+}
+
 func (c *Client) Store(key string, content []byte) error {
 	if c == nil || c.noop {
 		return nil

pkg/cli/gptscript.go

@@ -1,7 +1,6 @@
 package cli
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"os"
@@ -12,15 +11,12 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/builtin"
 	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/confirm"
-	"github.com/gptscript-ai/gptscript/pkg/engine"
+	"github.com/gptscript-ai/gptscript/pkg/gptscript"
 	"github.com/gptscript-ai/gptscript/pkg/input"
-	"github.com/gptscript-ai/gptscript/pkg/llm"
 	"github.com/gptscript-ai/gptscript/pkg/loader"
 	"github.com/gptscript-ai/gptscript/pkg/monitor"
 	"github.com/gptscript-ai/gptscript/pkg/mvl"
 	"github.com/gptscript-ai/gptscript/pkg/openai"
-	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
-	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/gptscript-ai/gptscript/pkg/server"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/gptscript-ai/gptscript/pkg/version"
@@ -50,8 +46,6 @@ type GPTScript struct {
 	Server        bool   `usage:"Start server"`
 	ListenAddress string `usage:"Server listen address" default:"127.0.0.1:9090"`
 	Chdir         string `usage:"Change current working directory" short:"C"`
-
-	_client llm.Client `usage:"-"`
 }
 
 func New() *cobra.Command {
@@ -78,33 +72,6 @@ func (r *GPTScript) Customize(cmd *cobra.Command) {
 	}
 }
 
-func (r *GPTScript) getClient(ctx context.Context) (llm.Client, error) {
-	if r._client != nil {
-		return r._client, nil
-	}
-
-	cacheClient, err := cache.New(cache.Options(r.CacheOptions))
-	if err != nil {
-		return nil, err
-	}
-
-	oaClient, err := openai.NewClient(openai.Options(r.OpenAIOptions), openai.Options{
-		Cache: cacheClient,
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	registry := llm.NewRegistry()
-
-	if err := registry.AddClient(ctx, oaClient); err != nil {
-		return nil, err
-	}
-
-	r._client = registry
-	return r._client, nil
-}
-
 func (r *GPTScript) listTools() error {
 	var lines []string
 	for _, tool := range builtin.ListTools() {
@@ -114,24 +81,6 @@ func (r *GPTScript) listTools() error {
 	return nil
 }
 
-func (r *GPTScript) listModels(ctx context.Context) error {
-	c, err := r.getClient(ctx)
-	if err != nil {
-		return err
-	}
-
-	models, err := c.ListModels(ctx)
-	if err != nil {
-		return err
-	}
-
-	for _, model := range models {
-		fmt.Println(model)
-	}
-
-	return nil
-}
-
 func (r *GPTScript) Pre(*cobra.Command, []string) error {
 	// chdir as soon as possible
 	if r.Chdir != "" {
@@ -164,37 +113,50 @@ func (r *GPTScript) Pre(*cobra.Command, []string) error {
 }
 
 func (r *GPTScript) Run(cmd *cobra.Command, args []string) error {
-	defer engine.CloseDaemons()
-
-	if r.ListModels {
-		return r.listModels(cmd.Context())
-	}
-
-	if r.ListTools {
-		return r.listTools()
+	gptOpt := gptscript.Options{
+		Cache:   cache.Options(r.CacheOptions),
+		OpenAI:  openai.Options(r.OpenAIOptions),
+		Monitor: monitor.Options(r.DisplayOptions),
+		Quiet:   r.Quiet,
+		Env:     os.Environ(),
 	}
 
 	if r.Server {
-		c, err := r.getClient(cmd.Context())
-		if err != nil {
-			return err
-		}
-		s, err := server.New(c, server.Options{
+		s, err := server.New(&server.Options{
 			ListenAddress: r.ListenAddress,
+			GPTScript:     gptOpt,
 		})
 		if err != nil {
 			return err
 		}
+		defer s.Close()
 		return s.Start(cmd.Context())
 	}
 
+	gptScript, err := gptscript.New(&gptOpt)
+	if err != nil {
+		return err
+	}
+	defer gptScript.Close()
+
+	if r.ListModels {
+		models, err := gptScript.ListModels(cmd.Context())
+		if err != nil {
+			return err
+		}
+		fmt.Println(strings.Join(models, "\n"))
+	}
+
+	if r.ListTools {
+		return r.listTools()
+	}
+
 	if len(args) == 0 {
 		return cmd.Help()
 	}
 
 	var (
 		prg types.Program
-		err error
 	)
 
 	if args[0] == "-" {
@@ -227,21 +189,6 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) error {
 		return assemble.Assemble(prg, out)
 	}
 
-	client, err := r.getClient(cmd.Context())
-	if err != nil {
-		return err
-	}
-
-	runner, err := runner.New(client, runner.Options{
-		MonitorFactory: monitor.NewConsole(monitor.Options(r.DisplayOptions), monitor.Options{
-			DisplayProgress: !*r.Quiet,
-		}),
-		RuntimeManager: runtimes.Default(cache.Complete(cache.Options(r.CacheOptions)).CacheDir),
-	})
-	if err != nil {
-		return err
-	}
-
 	toolInput, err := input.FromCLI(r.Input, args)
 	if err != nil {
 		return err
@@ -251,7 +198,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) error {
 	if r.Confirm {
 		ctx = confirm.WithConfirm(ctx, confirm.TextPrompt{})
 	}
-	s, err := runner.Run(ctx, prg, os.Environ(), toolInput)
+	s, err := gptScript.Run(ctx, prg, os.Environ(), toolInput)
 	if err != nil {
 		return err
 	}

pkg/engine/cmd.go

@@ -109,13 +109,13 @@ var ignoreENV = map[string]struct{}{
 	"GPTSCRIPT_TOOL_DIR": {},
 }
 
-func appendEnv(env []string, k, v string) []string {
-	for _, k := range []string{k, strings.ToUpper(strings.ReplaceAll(k, "-", "_"))} {
+func appendEnv(envs []string, k, v string) []string {
+	for _, k := range []string{k, env.ToEnvLike(k)} {
 		if _, ignore := ignoreENV[k]; !ignore {
-			env = append(env, k+"="+v)
+			envs = append(envs, k+"="+v)
 		}
 	}
-	return env
+	return envs
 }
 
 func appendInputAsEnv(env []string, input string) []string {

pkg/engine/http.go

@@ -55,6 +55,12 @@ func (e *Engine) runHTTP(ctx context.Context, prg *types.Program, tool types.Too
 		toolURL = parsed.String()
 	}
 
+	if tool.Blocking {
+		return &Return{
+			Result: &toolURL,
+		}, nil
+	}
+
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, toolURL, strings.NewReader(input))
 	if err != nil {
 		return nil, err

pkg/env/env.go

@@ -12,6 +12,10 @@ func execEquals(bin, check string) bool {
 		bin == check+".exe"
 }
 
+func ToEnvLike(v string) string {
+	return strings.ToUpper(strings.ReplaceAll(v, "-", "_"))
+}
+
 func Matches(cmd []string, bin string) bool {
 	switch len(cmd) {
 	case 0:

pkg/gptscript/gptscript.go

@@ -0,0 +1,108 @@
+package gptscript
+
+import (
+	"context"
+	"os"
+
+	"github.com/gptscript-ai/gptscript/pkg/cache"
+	"github.com/gptscript-ai/gptscript/pkg/engine"
+	"github.com/gptscript-ai/gptscript/pkg/llm"
+	"github.com/gptscript-ai/gptscript/pkg/monitor"
+	"github.com/gptscript-ai/gptscript/pkg/openai"
+	"github.com/gptscript-ai/gptscript/pkg/remote"
+	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
+	"github.com/gptscript-ai/gptscript/pkg/runner"
+	"github.com/gptscript-ai/gptscript/pkg/types"
+)
+
+type GPTScript struct {
+	Registry *llm.Registry
+	Runner   *runner.Runner
+}
+
+type Options struct {
+	Cache   cache.Options
+	OpenAI  openai.Options
+	Monitor monitor.Options
+	Runner  runner.Options
+	Quiet   *bool    `usage:"No output logging (set --quiet=false to force on even when there is no TTY)" short:"q"`
+	Env     []string `usage:"-"`
+}
+
+func complete(opts *Options) (result *Options) {
+	result = opts
+	if result == nil {
+		result = &Options{}
+	}
+	if result.Quiet == nil {
+		result.Quiet = new(bool)
+	}
+	if len(result.Env) == 0 {
+		result.Env = os.Environ()
+	}
+	return
+}
+
+func New(opts *Options) (*GPTScript, error) {
+	opts = complete(opts)
+
+	registry := llm.NewRegistry()
+
+	cacheClient, err := cache.New(opts.Cache)
+	if err != nil {
+		return nil, err
+	}
+
+	oAIClient, err := openai.NewClient(append([]openai.Options{opts.OpenAI}, openai.Options{
+		Cache: cacheClient,
+	})...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := registry.AddClient(oAIClient); err != nil {
+		return nil, err
+	}
+
+	if opts.Runner.MonitorFactory == nil {
+		opts.Runner.MonitorFactory = monitor.NewConsole(append([]monitor.Options{opts.Monitor}, monitor.Options{
+			DisplayProgress: !*opts.Quiet,
+		})...)
+	}
+
+	if opts.Runner.RuntimeManager == nil {
+		opts.Runner.RuntimeManager = runtimes.Default(cacheClient.CacheDir())
+	}
+
+	runner, err := runner.New(registry, opts.Runner)
+	if err != nil {
+		return nil, err
+	}
+
+	remoteClient := remote.New(runner, opts.Env, cacheClient)
+
+	if err := registry.AddClient(remoteClient); err != nil {
+		return nil, err
+	}
+
+	return &GPTScript{
+		Registry: registry,
+		Runner:   runner,
+	}, nil
+}
+
+func (g *GPTScript) Run(ctx context.Context, prg types.Program, envs []string, input string) (string, error) {
+	return g.Runner.Run(ctx, prg, envs, input)
+}
+
+func (g *GPTScript) Close() {
+	engine.CloseDaemons()
+}
+
+func (g *GPTScript) GetModel() engine.Model {
+	return g.Registry
+}
+
+func (g *GPTScript) ListModels(ctx context.Context) ([]string, error) {
+	return g.Registry.ListModels(ctx)
+}