feat: record request_id

common/helper/helper.go

@@ -1,9 +1,8 @@
 package helper
 
 import (
+	"context"
 	"fmt"
-	"github.com/gin-gonic/gin"
-	"github.com/songquanpeng/one-api/common/random"
 	"html/template"
 	"log"
 	"net"
@@ -11,6 +10,10 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/songquanpeng/one-api/common/random"
 )
 
 func OpenBrowser(url string) {
@@ -106,6 +109,18 @@ func GenRequestID() string {
 	return GetTimeString() + random.GetRandomNumberString(8)
 }
 
+func SetRequestID(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, RequestIdKey, id)
+}
+
+func GetRequestID(ctx context.Context) string {
+	rawRequestId := ctx.Value(RequestIdKey)
+	if rawRequestId == nil {
+		return ""
+	}
+	return rawRequestId.(string)
+}
+
 func GetResponseID(c *gin.Context) string {
 	logID := c.GetString(RequestIdKey)
 	return fmt.Sprintf("chatcmpl-%s", logID)