diff --git a/.gitignore b/.gitignore index ab951f8..8d766bf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,22 +1,22 @@ -# ---> Rust -# Generated by Cargo -# will have compiled files and executables -debug/ -target/ +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib -# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries -# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html -Cargo.lock +# Test binary, built with `go test -c` +*.test -# These are backup files generated by rustfmt -**/*.rs.bk +# Output of the go coverage tool, specifically when used with LiteIDE +*.out -# MSVC Windows builds of rustc generate these, which store debugging information -*.pdb +# Dependency directories (remove the comment below to include it) +# vendor/ -# RustRover -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +# Go workspace file +go.work +go.work.sum + +# env file +.env \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..47a3f7a --- /dev/null +++ b/go.mod @@ -0,0 +1,8 @@ +module ade9/llmproxymetrics + +go 1.23.3 + +require ( + github.com/caarlos0/env v3.5.0+incompatible + github.com/caarlos0/env/v11 v11.3.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..6c728a5 --- /dev/null +++ b/go.sum @@ -0,0 +1,4 @@ +github.com/caarlos0/env v3.5.0+incompatible h1:Yy0UN8o9Wtr/jGHZDpCBLpNrzcFLLM2yixi/rBrKyJs= +github.com/caarlos0/env v3.5.0+incompatible/go.mod h1:tdCsowwCzMLdkqRYDlHpZCp2UooDD3MspDBjZ2AD02Y= +github.com/caarlos0/env/v11 v11.3.1 h1:cArPWC15hWmEt+gWk7YBi7lEXTXCvpaSdCiZE2X5mCA= +github.com/caarlos0/env/v11 v11.3.1/go.mod h1:qupehSf/Y0TUTsxKywqRt/vJjN5nz6vauiYEUUr8P4U= diff --git a/llmproxymetrics.go b/llmproxymetrics.go new file mode 100644 index 0000000..df02871 --- /dev/null +++ b/llmproxymetrics.go @@ -0,0 +1,105 @@ +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "log" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "strconv" + "time" + + "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 + + lrw := &LoggingResponseWriter{ResponseWriter: w, body: new(bytes.Buffer)} + proxy := httputil.NewSingleHostReverseProxy(target) + + startTime := time.Now() + recorder := httptest.NewRecorder() + proxy.ServeHTTP(recorder, r) + + responseBody := recorder.Body.Bytes() + + var jsonResponse map[string]interface{} + err := json.Unmarshal(responseBody, &jsonResponse) + if err != nil { + log.Printf("Error unmarshalling JSON response: %v", err) + lrw.Write(responseBody) + return + } + + // Add your metrics metadata here + jsonResponse["metrics"] = map[string]interface{}{ + "requestPath": r.URL.Path, + "statusCode": recorder.Code, + "responseTime": time.Since(startTime).Milliseconds(), + // Add more metrics as needed + } + + modifiedResponseBody, err := json.Marshal(jsonResponse) + if err != nil { + log.Printf("Error marshalling modified JSON response: %v", err) + lrw.Write(responseBody) + return + } + + for name, values := range recorder.Header() { + for _, value := range values { + w.Header().Add(name, value) + } + } + + w.WriteHeader(recorder.Code) + lrw.Write(modifiedResponseBody) + + log.Printf("Response with metrics: %s", lrw.Body()) + } +} + +type LoggingResponseWriter struct { + http.ResponseWriter + body *bytes.Buffer +} + +func (lrw *LoggingResponseWriter) Write(b []byte) (int, error) { + lrw.body.Write(b) + return lrw.ResponseWriter.Write(b) +} + +func (lrw *LoggingResponseWriter) Body() string { + return lrw.body.String() +} + +func main() { + env.Parse(&cfg) + + 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(":%s", strconv.Itoa(cfg.Port)), nil) + if err != nil { + log.Fatal(err) + } +}