tbphp
14a85a7b05
* feat: Model redirect * fix: Hint error * fix: Prefix models error * feat: Override models API * feat: Code style
55 lines
1.7 KiB
Go
55 lines
1.7 KiB
Go
package proxy
|
|
|
|
import (
|
|
"gpt-load/internal/channel"
|
|
"gpt-load/internal/models"
|
|
"gpt-load/internal/utils"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
// shouldInterceptModelList checks if this is a model list request that should be intercepted
|
|
func shouldInterceptModelList(path string, method string) bool {
|
|
if method != "GET" {
|
|
return false
|
|
}
|
|
|
|
// Check various model list endpoints
|
|
return strings.HasSuffix(path, "/v1/models") ||
|
|
strings.HasSuffix(path, "/v1beta/models") ||
|
|
strings.Contains(path, "/v1beta/openai/v1/models")
|
|
}
|
|
|
|
// handleModelListResponse processes the model list response and applies filtering based on redirect rules
|
|
func (ps *ProxyServer) handleModelListResponse(c *gin.Context, resp *http.Response, group *models.Group, channelHandler channel.ChannelProxy) {
|
|
// Read the upstream response body
|
|
bodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
logrus.WithError(err).Error("Failed to read model list response body")
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to read response"})
|
|
return
|
|
}
|
|
|
|
// Decompress response data based on Content-Encoding
|
|
contentEncoding := resp.Header.Get("Content-Encoding")
|
|
decompressed, err := utils.DecompressResponse(contentEncoding, bodyBytes)
|
|
if err != nil {
|
|
logrus.WithError(err).Warn("Decompression failed, using original data")
|
|
decompressed = bodyBytes
|
|
}
|
|
|
|
// Transform model list (returns map[string]any directly, no marshaling)
|
|
response, err := channelHandler.TransformModelList(c.Request, decompressed, group)
|
|
if err != nil {
|
|
logrus.WithError(err).Error("Failed to transform model list")
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to process response"})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, response)
|
|
}
|