Proxy AI/LLM requests thorugh the balancer to hide the actual AI token

from the individual instances

Author: J12934

balancer/main.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"log"
 	"net/http"
+	"os"
 
 	"github.com/juice-shop/multi-juicer/balancer/pkg/bundle"
+	"github.com/juice-shop/multi-juicer/balancer/pkg/llmgateway"
 	"github.com/juice-shop/multi-juicer/balancer/pkg/notification"
 	"github.com/juice-shop/multi-juicer/balancer/pkg/scoring"
 	"github.com/juice-shop/multi-juicer/balancer/routes"
@@ -29,6 +31,20 @@ func main() {
 	scoringService.CalculateAndCacheScoreBoard(ctx)
 	go scoringService.StartingScoringWorker(ctx)
 	go notificationService.StartNotificationWatcher(ctx)
+
+	if b.Config.JuiceShopConfig.LLM.Enabled {
+		llmAPIKey := os.Getenv("LLM_API_KEY")
+		llmAPIURL := os.Getenv("LLM_API_URL")
+
+		usage := llmgateway.NewUsageTracker()
+		gateway, err := llmgateway.NewGateway(b.Config.CookieConfig.SigningKey, llmAPIURL, llmAPIKey, usage, b.Log)
+		if err != nil {
+			log.Fatalf("Failed to create LLM gateway: %v", err)
+		}
+		go usage.StartFlusher(ctx, b.ClientSet, b.RuntimeEnvironment.Namespace, b.Log)
+		go StartLLMGatewayServer(gateway, b.Log)
+	}
+
 	StartBalancerServer(b)
 }
 
@@ -47,6 +63,19 @@ func StartBalancerServer(b *bundle.Bundle) {
 	}
 }
 
+func StartLLMGatewayServer(gateway *llmgateway.Gateway, logger *log.Logger) {
+	router := http.NewServeMux()
+	router.Handle("/", gateway)
+	server := &http.Server{
+		Addr:    ":8082",
+		Handler: router,
+	}
+	logger.Println("Starting LLM gateway on :8082")
+	if err := server.ListenAndServe(); err != nil {
+		log.Fatalf("Failed to start LLM gateway server: %v", err)
+	}
+}
+
 func StartMetricsServer() {
 	metricsRouter := http.NewServeMux()
 	metricsRouter.Handle("GET /balancer/metrics", promhttp.Handler())

balancer/pkg/bundle/bundle.go

@@ -68,6 +68,12 @@ type CookieConfig struct {
 	Secure bool `json:"secure"`
 }
 
+type LLMConfig struct {
+	Enabled bool   `json:"enabled"`
+	Model   string `json:"model"`
+	ApiUrl  string `json:"apiUrl"`
+}
+
 type JuiceShopConfig struct {
 	Image            string                        `json:"image"`
 	Tag              string                        `json:"tag"`
@@ -87,6 +93,8 @@ type JuiceShopConfig struct {
 	VolumeMounts             []corev1.VolumeMount        `json:"volumeMounts"`
 	RuntimeClassName         *string                     `json:"runtimeClassName"`
 
+	LLM LLMConfig `json:"llm"`
+
 	JuiceShopPodConfig JuiceShopPodConfig `json:"pod"`
 }
 

balancer/pkg/llmgateway/gateway.go

@@ -0,0 +1,157 @@
+package llmgateway
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"log"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"strings"
+
+	"github.com/juice-shop/multi-juicer/balancer/pkg/signutil"
+)
+
+// openAIResponse is a minimal representation of an OpenAI chat completion response for usage extraction.
+type openAIResponse struct {
+	Usage *openAIUsage `json:"usage,omitempty"`
+}
+
+type openAIUsage struct {
+	InputTokens  int64 `json:"prompt_tokens"`
+	OutputTokens int64 `json:"completion_tokens"`
+}
+
+// Gateway proxies LLM requests from JuiceShop instances to an upstream LLM API.
+type Gateway struct {
+	signingKey  string
+	upstreamURL *url.URL
+	apiKey      string
+	usage       *UsageTracker
+	logger      *log.Logger
+}
+
+// NewGateway creates a new LLM gateway.
+func NewGateway(signingKey string, upstreamURL string, apiKey string, usage *UsageTracker, logger *log.Logger) (*Gateway, error) {
+	u, err := url.Parse(upstreamURL)
+	if err != nil {
+		return nil, err
+	}
+	return &Gateway{
+		signingKey:  signingKey,
+		upstreamURL: u,
+		apiKey:      apiKey,
+		usage:       usage,
+		logger:      logger,
+	}, nil
+}
+
+func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Extract bearer token
+	authHeader := r.Header.Get("Authorization")
+	if !strings.HasPrefix(authHeader, "Bearer ") {
+		http.Error(w, `{"error":"missing or invalid Authorization header"}`, http.StatusUnauthorized)
+		return
+	}
+	teamToken := strings.TrimPrefix(authHeader, "Bearer ")
+
+	// Validate token by verifying the HMAC signature and extracting the team name
+	team, err := signutil.Unsign(teamToken, g.signingKey)
+	if err != nil {
+		http.Error(w, `{"error":"invalid token"}`, http.StatusUnauthorized)
+		return
+	}
+
+	// Check if this is a chat completions request (for usage tracking)
+	isChatCompletion := strings.Contains(r.URL.Path, "/chat/completions")
+	g.logger.Printf("LLM gateway: request from team '%s': %s %s (isChatCompletion=%v)", team, r.Method, r.URL.Path, isChatCompletion)
+
+	// Create reverse proxy
+	proxy := &httputil.ReverseProxy{
+		Rewrite: func(pr *httputil.ProxyRequest) {
+			pr.SetURL(g.upstreamURL)
+			pr.Out.Host = g.upstreamURL.Host
+			// Replace the authorization header with the real API key
+			pr.Out.Header.Set("Authorization", "Bearer "+g.apiKey)
+		},
+	}
+
+	if isChatCompletion {
+		proxy.ModifyResponse = func(resp *http.Response) error {
+			return g.extractUsage(resp, team)
+		}
+	}
+
+	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+		g.logger.Printf("LLM gateway proxy error for team '%s': %v", team, err)
+		http.Error(w, `{"error":"upstream LLM API error"}`, http.StatusBadGateway)
+	}
+
+	proxy.ServeHTTP(w, r)
+}
+
+func (g *Gateway) extractUsage(resp *http.Response, team string) error {
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil
+	}
+
+	contentType := resp.Header.Get("Content-Type")
+	isSSE := strings.Contains(contentType, "text/event-stream")
+
+	body, err := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if err != nil {
+		g.logger.Printf("LLM gateway: failed to read response body for team '%s': %v", team, err)
+		resp.Body = io.NopCloser(bytes.NewReader(body))
+		return nil
+	}
+
+	// Restore the body for the client
+	resp.Body = io.NopCloser(bytes.NewReader(body))
+
+	if isSSE {
+		g.extractUsageFromSSE(body, team)
+	} else {
+		g.extractUsageFromJSON(body, team)
+	}
+	return nil
+}
+
+func (g *Gateway) extractUsageFromJSON(body []byte, team string) {
+	var result openAIResponse
+	if err := json.Unmarshal(body, &result); err != nil {
+		return
+	}
+	if result.Usage != nil {
+		g.logger.Printf("LLM gateway: usage for team '%s': input_tokens=%d, output_tokens=%d", team, result.Usage.InputTokens, result.Usage.OutputTokens)
+		g.usage.Add(team, result.Usage.InputTokens, result.Usage.OutputTokens)
+	}
+}
+
+// extractUsageFromSSE scans SSE events for usage data, which typically appears in the last chunk.
+func (g *Gateway) extractUsageFromSSE(body []byte, team string) {
+	// SSE format: lines starting with "data: " contain JSON payloads
+	// Scan all data lines for usage (it's usually in the last real chunk before "data: [DONE]")
+	lines := strings.Split(string(body), "\n")
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if !strings.HasPrefix(line, "data: ") {
+			continue
+		}
+		payload := strings.TrimPrefix(line, "data: ")
+		if payload == "[DONE]" {
+			continue
+		}
+		var chunk openAIResponse
+		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
+			continue
+		}
+		if chunk.Usage != nil {
+			g.logger.Printf("LLM gateway: SSE usage for team '%s': input_tokens=%d, output_tokens=%d", team, chunk.Usage.InputTokens, chunk.Usage.OutputTokens)
+			g.usage.Add(team, chunk.Usage.InputTokens, chunk.Usage.OutputTokens)
+			return
+		}
+	}
+	g.logger.Printf("LLM gateway: no usage data found in SSE stream for team '%s'", team)
+}