From c9233d2c9ad2b9ed819c13a11e0c33b3c007d68e Mon Sep 17 00:00:00 2001
From: Benson Wong <mostlygeek@gmail.com>
Date: Mon, 18 Nov 2024 15:58:28 -0800
Subject: [PATCH] use gin instead of standard http lib in main

---
 llama-swap.go         | 14 +++++++++-----
 proxy/proxymanager.go |  4 ++++
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/llama-swap.go b/llama-swap.go
index 1131043c..b7444d81 100644
--- a/llama-swap.go
+++ b/llama-swap.go
@@ -3,9 +3,9 @@ package main
 import (
 	"flag"
 	"fmt"
-	"net/http"
 	"os"
 
+	"github.com/gin-gonic/gin"
 	"github.com/mostlygeek/llama-swap/proxy"
 )
 
@@ -22,12 +22,16 @@ func main() {
 		os.Exit(1)
 	}
 
-	proxyManager := proxy.New(config)
-	http.HandleFunc("/", proxyManager.HandlerFunc)
+	if mode := os.Getenv("GIN_MODE"); mode != "" {
+		gin.SetMode(mode)
+	} else {
+		gin.SetMode(gin.ReleaseMode)
+	}
 
+	proxyManager := proxy.New(config)
 	fmt.Println("llama-swap listening on " + *listenStr)
-	if err := http.ListenAndServe(*listenStr, nil); err != nil {
-		fmt.Printf("Error starting server: %v\n", err)
+	if err := proxyManager.Run(*listenStr); err != nil {
+		fmt.Printf("Server error: %v\n", err)
 		os.Exit(1)
 	}
 }
diff --git a/proxy/proxymanager.go b/proxy/proxymanager.go
index f219a8ff..e9d810a5 100644
--- a/proxy/proxymanager.go
+++ b/proxy/proxymanager.go
@@ -46,6 +46,10 @@ func New(config *Config) *ProxyManager {
 	return pm
 }
 
+func (pm *ProxyManager) Run(addr ...string) error {
+	return pm.ginEngine.Run(addr...)
+}
+
 func (pm *ProxyManager) HandlerFunc(w http.ResponseWriter, r *http.Request) {
 	pm.ginEngine.ServeHTTP(w, r)
 }