Files
LLM-Proxy-Metrics/llmproxymetrics.go
2025-04-23 14:46:49 +03:00

106 lines
2.0 KiB
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)
}
}