From 81142ac54e7169338de5cc82a18303d016ab78c6 Mon Sep 17 00:00:00 2001
From: Donnie Adams <donnie@acorn.io>
Date: Wed, 6 Nov 2024 13:08:31 -0500
Subject: [PATCH] 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 <donnie@acorn.io>
---
 pkg/builtin/builtin.go     |  4 ++--
 pkg/engine/engine.go       |  2 +-
 pkg/gptscript/gptscript.go |  2 +-
 pkg/llm/proxy.go           | 17 +++++++++++++++--
 pkg/openai/client.go       |  2 +-
 pkg/tests/tester/runner.go |  2 +-
 6 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/pkg/builtin/builtin.go b/pkg/builtin/builtin.go
index c17f2c80..90f36919 100644
--- 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
index 7232157e..ae5a1368 100644
--- 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
index 65aa2209..dfb1771a 100644
--- 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
index aa8802be..82b33d02 100644
--- 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
index 3dbbca44..1894bdda 100644
--- 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
index 1f59ea03..44ec4e3c 100644
--- 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
 }