hatter/llm-proxy-go
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

📦 Update dependencies and add request tracing with UUID

Browse Source
This commit is contained in:
Hatter Jiang
2026-04-19 12:19:00 +08:00
parent 70b90216d6
commit 0b906016bc
3 changed files with 21 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

@@ -5,3 +5,5 @@ go 1.24.3
require github.com/joho/godotenv v1.5.1
require github.com/pelletier/go-toml/v2 v2.3.0
require github.com/google/uuid v1.6.0

2
go.sum
View File

@@ -1,3 +1,5 @@
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/pelletier/go-toml/v2 v2.3.0 h1:k59bC/lIZREW0/iVaQR8nDHxVq8OVlIzYCOJf421CaM=

28
main.go
View File

@@ -15,6 +15,7 @@ import (
"syscall"
"time"
"github.com/google/uuid"
"github.com/joho/godotenv"
"github.com/pelletier/go-toml/v2"
)
@@ -71,13 +72,15 @@ func main() {
func handleProxy(cfg Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
log.Println("Handle proxy: ", r.URL)
requestId := uuid.New().String()
log.Println("Handle proxy: ", requestId, r.URL)
if r.URL.Path == "/health" {
w.WriteHeader(http.StatusOK)
return
}
proxyReq := cloneRequest(r, cfg.UpstreamURL)
proxyReq := cloneRequest(requestId, r, cfg.UpstreamURL)
if cfg.APIKey != "" {
proxyReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
@@ -90,7 +93,7 @@ func handleProxy(cfg Config) http.HandlerFunc {
resp, err := client.Do(proxyReq)
if err != nil {
log.Println("Upstream error: ", err)
log.Println("Upstream error: ", requestId, err)
http.Error(w, fmt.Sprintf("upstream error: %v", err), http.StatusBadGateway)
return
}
@@ -102,18 +105,18 @@ func handleProxy(cfg Config) http.HandlerFunc {
w.WriteHeader(resp.StatusCode)
isStreamingRequest := isStreaming(resp)
log.Println("Request streaming: ", isStreamingRequest)
log.Println("Request streaming: ", requestId, isStreamingRequest)
if !isStreamingRequest {
io.Copy(w, resp.Body)
return
}
handleStream(w, resp.Body, cfg)
handleStream(requestId, w, resp.Body, cfg)
}
}
func cloneRequest(r *http.Request, upstreamURL string) *http.Request {
func cloneRequest(requestId string, r *http.Request, upstreamURL string) *http.Request {
upstream, _ := url.Parse(upstreamURL)
proxyReq := r.Clone(context.Background())
@@ -125,7 +128,7 @@ func cloneRequest(r *http.Request, upstreamURL string) *http.Request {
}
proxyReq.Host = upstream.Host
log.Println("Upstream proxy: ", proxyReq.URL)
log.Println("Upstream proxy: ", requestId, proxyReq.URL)
if val := r.Header.Get("Content-Type"); val != "" {
proxyReq.Header.Set("Content-Type", val)
}
@@ -144,14 +147,17 @@ func isStreaming(resp *http.Response) bool {
strings.Contains(ct, "stream")
}
func handleStream(w io.Writer, body io.Reader, cfg Config) {
func handleStream(requestId string, w io.Writer, body io.Reader, cfg Config) {
reader := bufio.NewReader(body)
for {
line, err := reader.ReadString('\n')
if err != nil {
if err != io.EOF {
log.Println("Stream error: ", requestId, err)
fmt.Fprintf(os.Stderr, "stream error: %v\n", err)
} else {
log.Println("Stream end: ", requestId)
}
break
}
@@ -161,12 +167,12 @@ func handleStream(w io.Writer, body io.Reader, cfg Config) {
continue
}
processChunk(w, line, cfg)
processChunk(requestId, w, line, cfg)
}
}
func processChunk(w io.Writer, line string, cfg Config) {
log.Println("Process chunk: ", line)
func processChunk(requestId string, w io.Writer, line string, cfg Config) {
log.Println("Process chunk: ", requestId, line)
if strings.HasPrefix(line, "data: ") {
data := strings.TrimPrefix(line, "data: ")
if data == "[DONE]" {