From 554c2f63a309683e8daae81383d3fa12a6c5128b Mon Sep 17 00:00:00 2001
From: Erik Seliger <erikseliger@me.com>
Date: Wed, 10 May 2023 02:56:12 +0200
Subject: [PATCH] llmproxy: Make dotcom API URL configurable (#51629)

---
 .../cmd/llm-proxy/internal/dotcom/dotcom.go   |  4 +--
 enterprise/cmd/llm-proxy/shared/config.go     | 12 +++++++
 enterprise/cmd/llm-proxy/shared/main.go       |  2 +-
 sg.config.yaml                                | 36 ++++++++++++++++---
 4 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/enterprise/cmd/llm-proxy/internal/dotcom/dotcom.go b/enterprise/cmd/llm-proxy/internal/dotcom/dotcom.go
index 05129c000a3..0beb2c051d2 100644
--- a/enterprise/cmd/llm-proxy/internal/dotcom/dotcom.go
+++ b/enterprise/cmd/llm-proxy/internal/dotcom/dotcom.go
@@ -11,8 +11,6 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-const endpoint = "https://sourcegraph.com/.api/graphql"
-
 // NewClient returns a new GraphQL client for the Sourcegraph.com API authenticated
 // with the given Sourcegraph access token.
 //
@@ -27,7 +25,7 @@ const endpoint = "https://sourcegraph.com/.api/graphql"
 //	println(resp.GetDotcom().ProductSubscriptionByAccessToken.LlmProxyAccess.Enabled)
 //
 // The client generator automatically ensures we're up-to-date with the GraphQL schema.
-func NewClient(token string) graphql.Client {
+func NewClient(endpoint, token string) graphql.Client {
 	return &tracedClient{graphql.NewClient(endpoint, &http.Client{
 		Transport: &tokenAuthTransport{
 			token:   token,
diff --git a/enterprise/cmd/llm-proxy/shared/config.go b/enterprise/cmd/llm-proxy/shared/config.go
index c08101c6564..a3ab3a05c72 100644
--- a/enterprise/cmd/llm-proxy/shared/config.go
+++ b/enterprise/cmd/llm-proxy/shared/config.go
@@ -1,6 +1,7 @@
 package shared
 
 import (
+	"net/url"
 	"time"
 
 	"github.com/sourcegraph/sourcegraph/internal/env"
@@ -12,6 +13,7 @@ type Config struct {
 	Address string
 
 	Dotcom struct {
+		URL         string
 		AccessToken string
 	}
 
@@ -27,7 +29,17 @@ type Config struct {
 func (c *Config) Load() {
 	c.Address = c.Get("LLM_PROXY_ADDR", ":9992", "Address to serve LLM proxy on.")
 	c.Dotcom.AccessToken = c.Get("LLM_PROXY_DOTCOM_ACCESS_TOKEN", "", "The Sourcegraph.com access token to be used.")
+	c.Dotcom.URL = c.Get("LLM_PROXY_DOTCOM_API_URL", "https://sourcegraph.com/.api/graphql", "Custom override for the dotcom API endpoint")
 	c.Anthropic.AccessToken = c.Get("LLM_PROXY_ANTHROPIC_ACCESS_TOKEN", "", "The Anthropic access token to be used.")
 	c.AllowAnonymous = c.GetBool("LLM_PROXY_ALLOW_ANONYMOUS", "false", "Allow anonymous access to LLM proxy.")
 	c.SourcesSyncInterval = c.GetInterval("LLM_PROXY_SOURCES_SYNC_INTERVAL", "2m", "The interval at which to sync actor sources.")
 }
+
+func (c *Config) Validate() error {
+	_, err := url.Parse(c.Dotcom.URL)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/enterprise/cmd/llm-proxy/shared/main.go b/enterprise/cmd/llm-proxy/shared/main.go
index 4fd03d383c3..9aa25c8306c 100644
--- a/enterprise/cmd/llm-proxy/shared/main.go
+++ b/enterprise/cmd/llm-proxy/shared/main.go
@@ -47,7 +47,7 @@ func Main(ctx context.Context, obctx *observation.Context, ready service.ReadyFu
 		productsubscription.NewSource(
 			obctx.Logger,
 			rcache.New("product-subscriptions"),
-			dotcom.NewClient(config.Dotcom.AccessToken)),
+			dotcom.NewClient(config.Dotcom.URL, config.Dotcom.AccessToken)),
 	}
 
 	// Set up our handler chain, which is run from the bottom up
diff --git a/sg.config.yaml b/sg.config.yaml
index ae5d57e53a3..4f14aafac27 100644
--- a/sg.config.yaml
+++ b/sg.config.yaml
@@ -396,6 +396,11 @@ commands:
 
       go build -gcflags="$GCFLAGS" -o .bin/llm-proxy github.com/sourcegraph/sourcegraph/enterprise/cmd/llm-proxy
     checkBinary: .bin/llm-proxy
+    env:
+      LLM_PROXY_ANTHROPIC_ACCESS_TOKEN: foobar
+      LLM_PROXY_DOTCOM_ACCESS_TOKEN: foobar
+      LLM_PROXY_DOTCOM_API_URL: https://sourcegraph.test:3443/.api/graphql
+      SRC_LOG_LEVEL: dbug
     watch:
       - lib
       - internal
@@ -1118,6 +1123,7 @@ commandsets:
       - zoekt-web-0
       - zoekt-web-1
       - caddy
+  # If you modify this command set, please consider also updating the dotcom runset.
   enterprise: &enterprise_set
     requiresDevPrivate: true
     checks:
@@ -1152,7 +1158,32 @@ commandsets:
       EXTSVC_CONFIG_FILE: ''
 
   dotcom:
-    <<: *enterprise_set
+    # This is 95% the enterprise runset, with the addition of LLM-proxy.
+    requiresDevPrivate: true
+    checks:
+      - docker
+      - redis
+      - postgres
+      - git
+    commands:
+      - frontend
+      - worker
+      - repo-updater
+      - web
+      - gitserver-0
+      - gitserver-1
+      - searcher
+      - symbols
+      - caddy
+      - docsite
+      - syntax-highlighter
+      - github-proxy
+      - zoekt-index-0
+      - zoekt-index-1
+      - zoekt-web-0
+      - zoekt-web-1
+      - blobstore
+      - llm-proxy
     env:
       SOURCEGRAPHDOTCOM_MODE: true
 
@@ -1426,9 +1457,6 @@ commandsets:
       - redis
     commands:
       - llm-proxy
-    env:
-      LLM_PROXY_ANTHROPIC_ACCESS_TOKEN: foobar
-      LLM_PROXY_DOTCOM_ACCESS_TOKEN: foobar
 
 tests:
   # These can be run with `sg test [name]`