LLM-Proxy-Metrics/llmproxymetrics.go

package main

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

	"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

		director := func(req *http.Request) {
			req.Header.Set("X-Forwarded-For", r.RemoteAddr)
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
		}

		modifyResponse := func(response *http.Response) error {
			pr, pw := io.Pipe()
			body := response.Body
			response.Body = pr

			go func() {
				defer pw.Close()

				reader := bufio.NewReader(body)
				for {
					line, err := reader.ReadBytes('\n')
					if err != nil {
						if err == io.EOF {
							handleJsonLine([]byte(string(line)))
							pw.Write(line)
							break
						}
						return
					}
					handleJsonLine(line)
					pw.Write(line)
				}

			}()

			return nil
		}

		proxy := httputil.NewSingleHostReverseProxy(target)
		proxy.Director = director
		proxy.ModifyResponse = modifyResponse

		proxy.ServeHTTP(w, r)
	}
}

func handleJsonLine(line []byte) {
	var jsonData map[string]interface{}
	err := json.Unmarshal([]byte(line), &jsonData)
	if err != nil {
		fmt.Println("Error parsing JSON:", err)
		return
	}

	if jsonData["done"].(bool) {
		duration := jsonData["eval_duration"].(float64)
		fmt.Printf("Duration: %.2f seconds\n", duration/1000000000.0)
	}
}

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)
	}
}