Merge pull request #326 from ibuildthecloud/main

chore: cache github and http lookups in loader

Author: ibuildthecloud

Makefile

@@ -9,6 +9,9 @@ build-ui:
 	touch static/ui/placeholder static/ui/_nuxt/_placeholder
 	cp -rp ui/.output/public/* static/ui/
 
+build-exe:
+	GOOS=windows go build -o bin/gptscript.exe -tags "${GO_TAGS}" .
+
 build:
 	CGO_ENABLED=0 go build -o bin/gptscript -tags "${GO_TAGS}" -ldflags "-s -w" .
 

pkg/cache/cache.go

@@ -2,12 +2,17 @@ package cache
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/gob"
+	"encoding/hex"
 	"errors"
 	"io/fs"
 	"os"
 	"path/filepath"
 
 	"github.com/adrg/xdg"
+	"github.com/getkin/kin-openapi/openapi3"
+	openai "github.com/gptscript-ai/chat-completion-client"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/gptscript-ai/gptscript/pkg/version"
 )
@@ -22,6 +27,11 @@ type Options struct {
 	CacheDir     string `usage:"Directory to store cache (default: $XDG_CACHE_HOME/gptscript)"`
 }
 
+func init() {
+	gob.Register(openai.ChatCompletionRequest{})
+	gob.Register(openapi3.Schema{})
+}
+
 func Complete(opts ...Options) (result Options) {
 	for _, opt := range opts {
 		result.CacheDir = types.FirstSet(opt.CacheDir, result.CacheDir)
@@ -59,22 +69,62 @@ func (c *Client) CacheDir() string {
 	return c.dir
 }
 
-func (c *Client) Store(key string, content []byte) error {
-	if c == nil || c.noop {
+func (c *Client) cacheKey(key any) (string, error) {
+	hash := sha256.New()
+	if err := gob.NewEncoder(hash).Encode(key); err != nil {
+		return "", err
+	}
+	digest := hash.Sum(nil)
+	return hex.EncodeToString(digest), nil
+}
+
+func (c *Client) Store(ctx context.Context, key, value any) error {
+	if c == nil {
 		return nil
 	}
-	return os.WriteFile(filepath.Join(c.dir, key), content, 0644)
+
+	if c.noop || IsNoCache(ctx) {
+		keyValue, err := c.cacheKey(key)
+		if err == nil {
+			p := filepath.Join(c.dir, keyValue)
+			if _, err := os.Stat(p); err == nil {
+				_ = os.Remove(p)
+			}
+		}
+		return nil
+	}
+
+	keyValue, err := c.cacheKey(key)
+	if err != nil {
+		return err
+	}
+
+	f, err := os.Create(filepath.Join(c.dir, keyValue))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	return gob.NewEncoder(f).Encode(value)
 }
 
-func (c *Client) Get(key string) ([]byte, bool, error) {
-	if c == nil || c.noop {
-		return nil, false, nil
+func (c *Client) Get(ctx context.Context, key, out any) (bool, error) {
+	if c == nil || c.noop || IsNoCache(ctx) {
+		return false, nil
 	}
-	data, err := os.ReadFile(filepath.Join(c.dir, key))
+
+	keyValue, err := c.cacheKey(key)
+	if err != nil {
+		return false, err
+	}
+
+	f, err := os.Open(filepath.Join(c.dir, keyValue))
 	if errors.Is(err, fs.ErrNotExist) {
-		return nil, false, nil
+		return false, nil
 	} else if err != nil {
-		return nil, false, err
+		return false, err
 	}
-	return data, true, nil
+	defer f.Close()
+
+	return gob.NewDecoder(f).Decode(out) == nil, nil
 }

pkg/cli/eval.go

@@ -49,17 +49,19 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
 		tool.Temperature = &temp32
 	}
 
-	prg, err := loader.ProgramFromSource(cmd.Context(), tool.String(), "")
+	opts, err := e.gptscript.NewGPTScriptOpts()
 	if err != nil {
 		return err
 	}
 
-	opts, err := e.gptscript.NewGPTScriptOpts()
+	runner, err := gptscript.New(&opts)
 	if err != nil {
 		return err
 	}
 
-	runner, err := gptscript.New(&opts)
+	prg, err := loader.ProgramFromSource(cmd.Context(), tool.String(), "", loader.Options{
+		Cache: runner.Cache,
+	})
 	if err != nil {
 		return err
 	}

pkg/cli/gptscript.go

@@ -259,7 +259,7 @@ func (r *GPTScript) listModels(ctx context.Context, gptScript *gptscript.GPTScri
 	return nil
 }
 
-func (r *GPTScript) readProgram(ctx context.Context, args []string) (prg types.Program, err error) {
+func (r *GPTScript) readProgram(ctx context.Context, runner *gptscript.GPTScript, args []string) (prg types.Program, err error) {
 	if len(args) == 0 {
 		return
 	}
@@ -278,10 +278,14 @@ func (r *GPTScript) readProgram(ctx context.Context, args []string) (prg types.P
 			}
 			r.readData = data
 		}
-		return loader.ProgramFromSource(ctx, string(data), r.SubTool)
+		return loader.ProgramFromSource(ctx, string(data), r.SubTool, loader.Options{
+			Cache: runner.Cache,
+		})
 	}
 
-	return loader.Program(ctx, args[0], r.SubTool)
+	return loader.Program(ctx, args[0], r.SubTool, loader.Options{
+		Cache: runner.Cache,
+	})
 }
 
 func (r *GPTScript) PrintOutput(toolInput, toolOutput string) (err error) {
@@ -337,7 +341,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 		return r.listModels(ctx, gptScript, args)
 	}
 
-	prg, err := r.readProgram(ctx, args)
+	prg, err := r.readProgram(ctx, gptScript, args)
 	if err != nil {
 		return err
 	}
@@ -392,7 +396,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 
 	if prg.IsChat() || r.ForceChat {
 		return chat.Start(r.NewRunContext(cmd), nil, gptScript, func() (types.Program, error) {
-			return r.readProgram(ctx, args)
+			return r.readProgram(ctx, gptScript, args)
 		}, os.Environ(), toolInput)
 	}
 

pkg/debugcmd/debug.go

@@ -16,6 +16,10 @@ type WrappedCmd struct {
 	Dir string
 }
 
+func (w *WrappedCmd) Stdout() string {
+	return w.r.Stdout()
+}
+
 func (w *WrappedCmd) Run() error {
 	if len(w.Env) > 0 {
 		w.c.Env = w.Env
@@ -51,6 +55,16 @@ type recorder struct {
 	entries []entry
 }
 
+func (r *recorder) Stdout() string {
+	buf := strings.Builder{}
+	for _, e := range r.entries {
+		if !e.err {
+			buf.Write(e.data)
+		}
+	}
+	return buf.String()
+}
+
 func (r *recorder) dump() string {
 	var errMessage strings.Builder
 	for _, entry := range r.entries {

pkg/gptscript/gptscript.go

@@ -24,6 +24,7 @@ var log = mvl.Package()
 type GPTScript struct {
 	Registry               *llm.Registry
 	Runner                 *runner.Runner
+	Cache                  *cache.Client
 	WorkspacePath          string
 	DeleteWorkspaceOnClose bool
 }
@@ -99,6 +100,7 @@ func New(opts *Options) (*GPTScript, error) {
 	return &GPTScript{
 		Registry:               registry,
 		Runner:                 runner,
+		Cache:                  cacheClient,
 		WorkspacePath:          opts.Workspace,
 		DeleteWorkspaceOnClose: opts.Workspace == "",
 	}, nil

pkg/hash/seed.go

@@ -1,10 +1,14 @@
 package hash
 
-import "hash/fnv"
+import (
+	"encoding/gob"
+	"hash/fnv"
+)
 
 func Seed(input any) int {
-	s := Encode(input)
 	h := fnv.New32a()
-	_, _ = h.Write([]byte(s))
+	if err := gob.NewEncoder(h).Encode(input); err != nil {
+		panic(err)
+	}
 	return int(h.Sum32())
 }

pkg/hash/sha256.go

@@ -2,8 +2,8 @@ package hash
 
 import (
 	"crypto/sha256"
+	"encoding/gob"
 	"encoding/hex"
-	"encoding/json"
 )
 
 func ID(parts ...string) string {
@@ -19,31 +19,9 @@ func ID(parts ...string) string {
 }
 
 func Digest(obj any) string {
-	data, err := json.Marshal(obj)
-	if err != nil {
+	hash := sha256.New()
+	if err := gob.NewEncoder(hash).Encode(obj); err != nil {
 		panic(err)
 	}
-
-	hash := sha256.Sum256(data)
-	return hex.EncodeToString(hash[:])
-}
-
-func Encode(obj any) string {
-	data, err := json.Marshal(obj)
-	if err != nil {
-		panic(err)
-	}
-
-	asMap := map[string]any{}
-	if err := json.Unmarshal(data, &asMap); err != nil {
-		panic(err)
-	}
-
-	data, err = json.Marshal(asMap)
-	if err != nil {
-		panic(err)
-	}
-
-	hash := sha256.Sum256(data)
-	return hex.EncodeToString(hash[:])
+	return hex.EncodeToString(hash.Sum(nil))
 }

pkg/loader/github/github.go

@@ -1,6 +1,7 @@
 package github
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -9,7 +10,9 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/loader"
+	"github.com/gptscript-ai/gptscript/pkg/repos/git"
 	"github.com/gptscript-ai/gptscript/pkg/system"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 )
@@ -29,11 +32,14 @@ func init() {
 	loader.AddVSC(Load)
 }
 
-func getCommit(account, repo, ref string) (string, error) {
-	url := fmt.Sprintf(githubCommitURL, account, repo, ref)
-	client := &http.Client{}
+func getCommitLsRemote(ctx context.Context, account, repo, ref string) (string, error) {
+	url := fmt.Sprintf(githubRepoURL, account, repo)
+	return git.LsRemote(ctx, url, ref)
+}
 
-	req, err := http.NewRequest(http.MethodGet, url, nil)
+func getCommit(ctx context.Context, account, repo, ref string) (string, error) {
+	url := fmt.Sprintf(githubCommitURL, account, repo, ref)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
 		return "", fmt.Errorf("failed to create request of %s/%s at %s: %w", account, repo, url, err)
 	}
@@ -42,13 +48,15 @@ func getCommit(account, repo, ref string) (string, error) {
 		req.Header.Add("Authorization", "Bearer "+githubAuthToken)
 	}
 
-	resp, err := client.Do(req)
-
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return "", err
 	} else if resp.StatusCode != http.StatusOK {
 		c, _ := io.ReadAll(resp.Body)
 		resp.Body.Close()
+		if commit, err := getCommitLsRemote(ctx, account, repo, ref); err == nil {
+			return commit, nil
+		}
 		return "", fmt.Errorf("failed to get GitHub commit of %s/%s at %s: %s %s",
 			account, repo, ref, resp.Status, c)
 	}
@@ -68,7 +76,7 @@ func getCommit(account, repo, ref string) (string, error) {
 	return commit.SHA, nil
 }
 
-func Load(urlName string) (string, *types.Repo, bool, error) {
+func Load(ctx context.Context, _ *cache.Client, urlName string) (string, *types.Repo, bool, error) {
 	if !strings.HasPrefix(urlName, GithubPrefix) {
 		return "", nil, false, nil
 	}
@@ -93,7 +101,7 @@ func Load(urlName string) (string, *types.Repo, bool, error) {
 		path += "/tool.gpt"
 	}
 
-	ref, err := getCommit(account, repo, ref)
+	ref, err := getCommit(ctx, account, repo, ref)
 	if err != nil {
 		return "", nil, false, err
 	}