feat: allow bypassing of built-in proxy

If specific proxy environment variables are set, the built-in proxy will be bypassed. Signed-off-by: Donnie Adams <[email protected]>
gptscript-ai · Nov 6, 2024 · 266bc58 · 266bc58
1 parent 53a6ae0
commit 266bc58
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 8 deletions.
diff --git a/pkg/builtin/builtin.go b/pkg/builtin/builtin.go
@@ -695,9 +695,9 @@ func invalidArgument(input string, err error) string {
 	return fmt.Sprintf("Failed to parse arguments %s: %v", input, err)
 }
 
-func SysModelProviderCredential(ctx context.Context, _ []string, _ string, _ chan<- string) (string, error) {
+func SysModelProviderCredential(ctx context.Context, env []string, _ string, _ chan<- string) (string, error) {
 	engineContext, _ := engine.FromContext(ctx)
-	auth, url, err := engineContext.Engine.Model.ProxyInfo()
+	auth, url, err := engineContext.Engine.Model.ProxyInfo(env)
 	if err != nil {
 		return "", err
 	}

diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
@@ -14,7 +14,7 @@ import (
 
 type Model interface {
 	Call(ctx context.Context, messageRequest types.CompletionRequest, env []string, status chan<- types.CompletionStatus) (*types.CompletionMessage, error)
-	ProxyInfo() (string, string, error)
+	ProxyInfo([]string) (string, string, error)
 }
 
 type RuntimeManager interface {

diff --git a/pkg/gptscript/gptscript.go b/pkg/gptscript/gptscript.go
@@ -317,7 +317,7 @@ func (n noopModel) Call(_ context.Context, _ types.CompletionRequest, _ []string
 	return nil, errors.New("unsupported")
 }
 
-func (n noopModel) ProxyInfo() (string, string, error) {
+func (n noopModel) ProxyInfo([]string) (string, string, error) {
 	return "", "", errors.New("unsupported")
 }
 

diff --git a/pkg/llm/proxy.go b/pkg/llm/proxy.go
@@ -15,7 +15,20 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/openai"
 )
 
-func (r *Registry) ProxyInfo() (string, string, error) {
+func (r *Registry) ProxyInfo(env []string) (string, string, error) {
+	var proxyURL, proxyToken string
+	for _, e := range env {
+		if url, ok := strings.CutPrefix(e, "GPTSCRIPT_MODEL_PROVIDER_PROXY_URL="); ok {
+			proxyURL = url
+		} else if token, ok := strings.CutPrefix(e, "GPTSCRIPT_MODEL_PROVIDER_PROXY_TOKEN="); ok {
+			proxyToken = token
+		}
+	}
+
+	if proxyToken != "" && proxyURL != "" {
+		return proxyToken, proxyURL, nil
+	}
+
 	r.proxyLock.Lock()
 	defer r.proxyLock.Unlock()
 
@@ -77,7 +90,7 @@ func (r *Registry) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	auth, targetURL := oai.ProxyInfo()
+	auth, targetURL := oai.ProxyInfo(nil)
 	if targetURL == "" {
 		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
 		return

diff --git a/pkg/openai/client.go b/pkg/openai/client.go
@@ -132,7 +132,7 @@ func NewClient(ctx context.Context, credStore credentials.CredentialStore, opts
 	}, nil
 }
 
-func (c *Client) ProxyInfo() (token, urlBase string) {
+func (c *Client) ProxyInfo([]string) (token, urlBase string) {
 	if c.invalidAuth {
 		return "", ""
 	}

diff --git a/pkg/tests/tester/runner.go b/pkg/tests/tester/runner.go
@@ -31,7 +31,7 @@ type Result struct {
 	Err     error
 }
 
-func (c *Client) ProxyInfo() (string, string, error) {
+func (c *Client) ProxyInfo([]string) (string, string, error) {
 	return "test-auth", "test-url", nil
 }