LLM-Proxy-Metrics/llmproxymetrics.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"

	"github.com/caarlos0/env/v11"
)

var cfg config

type config struct {
	BaseURL string `env:"BASE_URL"`
	Port    int    `env:"PORT"`
}

func createProxy(target *url.URL) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		r.Host = target.Host
		r.URL.Scheme = target.Scheme
		r.URL.Host = target.Host

		// startTime := time.Now()
		proxy := httputil.NewSingleHostReverseProxy(target)
		proxy.Director = func(req *http.Request) {
			req.Header.Set("X-Forwarded-For", r.RemoteAddr)
			req.Host = target.Host
		}

		recorder := httptest.NewRecorder()

		proxy.ServeHTTP(recorder, r)

		responseBody := recorder.Body.Bytes()
		contentType := recorder.Result().Header.Get("Content-Type")

		log.Printf("Response Content-Type: %s", contentType)
		log.Printf("Response Body: %s", string(responseBody))

		if strings.Contains(contentType, "application/json") {
			var jsonResponse map[string]interface{}
			err := json.Unmarshal(responseBody, &jsonResponse)
			if err != nil {
				log.Printf("Error unmarshalling JSON response: %v", err)
				w.WriteHeader(recorder.Code)
				w.Write(responseBody)
				return
			}

			// jsonResponse["metrics"] = map[string]interface{}{
			// 	"requestPath":  r.URL.Path,
			// 	"statusCode":   recorder.Code,
			// 	"responseTime": time.Since(startTime).Milliseconds(),
			// }

			modifiedResponseBody, err := json.Marshal(jsonResponse)
			if err != nil {
				log.Printf("Error marshalling modified JSON response: %v", err)
				w.WriteHeader(recorder.Code)
				w.Write(responseBody)
				return
			}

			for name, values := range recorder.Header() {
				for _, value := range values {
					w.Header().Add(name, value)
				}
			}

			w.WriteHeader(recorder.Code)
			w.Write(modifiedResponseBody)

			// log.Printf("Modified Response Body: %s", string(modifiedResponseBody))
		} else if strings.Contains(contentType, "application/x-ndjson") {
			var modifiedResponseBody []string
			scanner := bufio.NewScanner(strings.NewReader(string(responseBody)))
			for scanner.Scan() {
				line := scanner.Text()
				if line != "" {
					var jsonResponse map[string]interface{}
					err := json.Unmarshal([]byte(line), &jsonResponse)
					if err != nil {
						log.Printf("Error unmarshalling NDJSON line: %v", err)
						modifiedResponseBody = append(modifiedResponseBody, line)
						continue
					}

					// jsonResponse["metrics"] = map[string]interface{}{
					// 	"requestPath":  r.URL.Path,
					// 	"statusCode":   recorder.Code,
					// 	"responseTime": time.Since(startTime).Milliseconds(),
					// }

					modifiedLine, err := json.Marshal(jsonResponse)
					if err != nil {
						log.Printf("Error marshalling modified NDJSON line: %v", err)
						modifiedResponseBody = append(modifiedResponseBody, string([]byte(line)))
						continue
					}
					modifiedResponseBody = append(modifiedResponseBody, string(modifiedLine))
				}
			}

			if err := scanner.Err(); err != nil {
				log.Printf("Error scanning NDJSON: %v", err)
				w.WriteHeader(recorder.Code)
				w.Write(responseBody)
				return
			}

			for name, values := range recorder.Header() {
				for _, value := range values {
					w.Header().Add(name, value)
				}
			}

			w.WriteHeader(recorder.Code)
			w.Write([]byte(strings.Join(modifiedResponseBody, "\n")))

			// log.Printf("Modified Response Body: %s", strings.Join(modifiedResponseBody, "\n"))
		} else {
			for name, values := range recorder.Header() {
				for _, value := range values {
					w.Header().Add(name, value)
				}
			}

			w.WriteHeader(recorder.Code)
			w.Write(responseBody)

			// log.Printf("Response without metrics: %s", string(responseBody))
		}
	}
}

func main() {
	err := env.Parse(&cfg)
	if err != nil {
		log.Fatalf("Error parsing environment variables: %v", err)
	}

	targetURL, err := url.Parse(cfg.BaseURL)
	if err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("/", createProxy(targetURL))

	log.Printf("Starting proxy server on :%s", strconv.Itoa(cfg.Port))
	err = http.ListenAndServe(fmt.Sprintf(":%d", cfg.Port), nil)
	if err != nil {
		log.Fatal(err)
	}
}