|
- package common
-
- import (
- "net/http"
- "strings"
-
- commonpkg "github.com/QuantumNous/new-api/common"
- "github.com/gin-gonic/gin"
- )
-
- const (
- ContextKeyRelayChatID = "relay_chat_id"
- ContextKeyRelayUpstreamID = "relay_upstream_id"
- )
-
- func SetRelayChatID(c *gin.Context, chatID string) {
- if c == nil {
- return
- }
- chatID = strings.TrimSpace(chatID)
- if chatID == "" {
- return
- }
- c.Set(ContextKeyRelayChatID, chatID)
- }
-
- func GetRelayChatID(c *gin.Context) string {
- if c == nil {
- return ""
- }
- return strings.TrimSpace(c.GetString(ContextKeyRelayChatID))
- }
-
- func SetRelayUpstreamID(c *gin.Context, upstreamID string) {
- if c == nil {
- return
- }
- upstreamID = strings.TrimSpace(upstreamID)
- if upstreamID == "" {
- return
- }
- c.Set(ContextKeyRelayUpstreamID, upstreamID)
- }
-
- func GetRelayUpstreamID(c *gin.Context) string {
- if c == nil {
- return ""
- }
- return strings.TrimSpace(c.GetString(ContextKeyRelayUpstreamID))
- }
-
- func ExtractTopLevelChatID(c *gin.Context) (string, error) {
- if c == nil {
- return "", nil
- }
- bodyStorage, err := commonpkg.GetBodyStorage(c)
- if err != nil {
- return "", err
- }
- bodyBytes, err := bodyStorage.Bytes()
- if err != nil || len(bodyBytes) == 0 {
- return "", err
- }
- var payload map[string]interface{}
- if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil {
- return "", nil
- }
- return strings.TrimSpace(commonpkg.Interface2String(payload["chat_id"])), nil
- }
-
- func CaptureUpstreamIDFromHTTPResponse(info *RelayInfo, resp *http.Response) {
- if info == nil || resp == nil {
- return
- }
- for _, key := range []string{"x-request-id", "request-id"} {
- value := strings.TrimSpace(resp.Header.Get(key))
- if value != "" {
- info.UpstreamID = value
- return
- }
- }
- }
|