fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 优化代理服务

Browse Source
This commit is contained in:
tbphp
2025-07-11 17:11:00 +08:00
parent 6ffbb7e9a1
commit 1368792aa0
16 changed files with 249 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
internal/channel/base_channel.go
View File

@@ -2,13 +2,15 @@ package channel
import (
"bytes"
"encoding/json"
"fmt"
"gpt-load/internal/models"
"gpt-load/internal/types"
"net/http"
"net/url"
"reflect"
"strings"
"sync"
"github.com/sirupsen/logrus"
"gorm.io/datatypes"
)
@@ -28,7 +30,7 @@ type BaseChannel struct {
TestModel string
upstreamLock sync.Mutex
groupUpstreams datatypes.JSON
groupConfig datatypes.JSONMap
effectiveConfig *types.SystemSettings
}
// getUpstreamURL selects an upstream URL using a smooth weighted round-robin algorithm.
@@ -64,28 +66,33 @@ func (b *BaseChannel) getUpstreamURL() *url.URL {
return best.URL
}
// BuildUpstreamURL constructs the target URL for the upstream service.
func (b *BaseChannel) BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error) {
base := b.getUpstreamURL()
if base == nil {
return "", fmt.Errorf("no upstream URL configured for channel %s", b.Name)
}
finalURL := *base
proxyPrefix := "/proxy/" + group.Name
if strings.HasPrefix(originalURL.Path, proxyPrefix) {
finalURL.Path = strings.TrimPrefix(originalURL.Path, proxyPrefix)
} else {
finalURL.Path = originalURL.Path
}
finalURL.RawQuery = originalURL.RawQuery
return finalURL.String(), nil
}
// IsConfigStale checks if the channel's configuration is stale compared to the provided group.
func (b *BaseChannel) IsConfigStale(group *models.Group) bool {
// It's important to compare the raw JSON here to detect any changes.
if !bytes.Equal(b.groupUpstreams, group.Upstreams) {
return true
}
// For JSONMap, we need to marshal it to compare.
currentConfigBytes, err := json.Marshal(b.groupConfig)
if err != nil {
// Log the error and assume it's stale to be safe
logrus.Errorf("failed to marshal current group config: %v", err)
return true
}
newConfigBytes, err := json.Marshal(group.Config)
if err != nil {
// Log the error and assume it's stale
logrus.Errorf("failed to marshal new group config: %v", err)
return true
}
if !bytes.Equal(currentConfigBytes, newConfigBytes) {
if !reflect.DeepEqual(b.effectiveConfig, &group.EffectiveConfig) {
return true
}

3
internal/channel/channel.go
View File

@@ -10,17 +10,14 @@ import (
)
// ChannelProxy defines the interface for different API channel proxies.
// It's responsible for channel-specific logic like URL building and request modification.
type ChannelProxy interface {
// BuildUpstreamURL constructs the target URL for the upstream service.
BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error)
// ModifyRequest allows the channel to add specific headers or modify the request
// before it's sent to the upstream service.
ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group)
// IsStreamRequest checks if the request is for a streaming response,
// now using the cached request body to avoid re-reading the stream.
IsStreamRequest(c *gin.Context, bodyBytes []byte) bool
// ValidateKey checks if the given API key is valid.

26
internal/channel/factory.go
View File

@@ -22,7 +22,6 @@ var (
)
// Register adds a new channel constructor to the registry.
// This function is intended to be called from the init() function of each channel implementation.
func Register(channelType string, constructor channelConstructor) {
if _, exists := channelRegistry[channelType]; exists {
panic(fmt.Sprintf("channel type '%s' is already registered", channelType))
@@ -57,7 +56,6 @@ func NewFactory(settingsManager *config.SystemSettingsManager, clientManager *ht
}
// GetChannel returns a channel proxy based on the group's channel type.
// It uses a cache to ensure that only one instance of a channel is created for each group.
func (f *Factory) GetChannel(group *models.Group) (ChannelProxy, error) {
f.cacheLock.Lock()
defer f.cacheLock.Unlock()
@@ -120,20 +118,28 @@ func (f *Factory) newBaseChannel(name string, group *models.Group) (*BaseChannel
MaxIdleConnsPerHost: group.EffectiveConfig.MaxIdleConnsPerHost,
ResponseHeaderTimeout: time.Duration(group.EffectiveConfig.ResponseHeaderTimeout) * time.Second,
DisableCompression: group.EffectiveConfig.DisableCompression,
WriteBufferSize: 32 * 1024, // Use a reasonable default buffer size for regular requests
WriteBufferSize: 32 * 1024,
ReadBufferSize: 32 * 1024,
ForceAttemptHTTP2: true,
TLSHandshakeTimeout: 15 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
}
// Create a dedicated configuration for streaming requests.
// This configuration is optimized for low-latency, long-running connections.
streamConfig := *clientConfig
streamConfig.RequestTimeout = 0 // No overall timeout for the entire request.
streamConfig.DisableCompression = true // Always disable compression for streaming to reduce latency.
streamConfig.WriteBufferSize = 0 // Disable buffering for real-time data transfer.
streamConfig.RequestTimeout = 0
streamConfig.DisableCompression = true
streamConfig.WriteBufferSize = 0
streamConfig.ReadBufferSize = 0
// For stream-specific connection pool, we can use a simple heuristic like doubling the regular one.
// Use a larger, independent connection pool for streaming clients to avoid exhaustion.
streamConfig.MaxIdleConns = group.EffectiveConfig.MaxIdleConns * 2
if streamConfig.MaxIdleConns < 200 {
streamConfig.MaxIdleConns = 200
}
streamConfig.MaxIdleConnsPerHost = group.EffectiveConfig.MaxIdleConnsPerHost * 2
if streamConfig.MaxIdleConnsPerHost < 40 {
streamConfig.MaxIdleConnsPerHost = 40
}
// Get both clients from the manager using their respective configurations.
httpClient := f.clientManager.GetClient(clientConfig)
@@ -145,7 +151,7 @@ func (f *Factory) newBaseChannel(name string, group *models.Group) (*BaseChannel
HTTPClient: httpClient,
StreamClient: streamClient,
TestModel: group.TestModel,
groupUpstreams: group.Upstreams,
groupConfig: group.Config,
groupUpstreams: group.Upstreams,
effectiveConfig: &group.EffectiveConfig,
}, nil
}

29
internal/channel/gemini_channel.go
View File

@@ -9,7 +9,6 @@ import (
"gpt-load/internal/models"
"io"
"net/http"
"net/url"
"strings"
"github.com/gin-gonic/gin"
@@ -34,31 +33,6 @@ func newGeminiChannel(f *Factory, group *models.Group) (ChannelProxy, error) {
}, nil
}
// BuildUpstreamURL constructs the target URL for the Gemini service.
func (ch *GeminiChannel) BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error) {
base := ch.getUpstreamURL()
if base == nil {
// Fallback to default Gemini URL
base, _ = url.Parse("https://generativelanguage.googleapis.com")
}
finalURL := *base
// The originalURL.Path contains the full path, e.g., "/proxy/gemini/v1beta/models/gemini-pro:generateContent".
// We need to strip the proxy prefix to get the correct upstream path.
proxyPrefix := "/proxy/" + group.Name
if strings.HasPrefix(originalURL.Path, proxyPrefix) {
finalURL.Path = strings.TrimPrefix(originalURL.Path, proxyPrefix)
} else {
// Fallback for safety.
finalURL.Path = originalURL.Path
}
// The API key will be added to RawQuery in ModifyRequest.
finalURL.RawQuery = originalURL.RawQuery
return finalURL.String(), nil
}
// ModifyRequest adds the API key as a query parameter for Gemini requests.
func (ch *GeminiChannel) ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group) {
q := req.URL.Query()
@@ -73,11 +47,8 @@ func (ch *GeminiChannel) ValidateKey(ctx context.Context, key string) (bool, err
return false, fmt.Errorf("no upstream URL configured for channel %s", ch.Name)
}
// Use the test model specified in the group settings.
// The path format for Gemini is /v1beta/models/{model}:generateContent
reqURL := fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", upstreamURL.String(), ch.TestModel, key)
// Use a minimal, low-cost payload for validation
payload := gin.H{
"contents": []gin.H{
{"parts": []gin.H{

29
internal/channel/openai_channel.go
View File

@@ -9,7 +9,6 @@ import (
"gpt-load/internal/models"
"io"
"net/http"
"net/url"
"strings"
"github.com/gin-gonic/gin"
@@ -34,34 +33,6 @@ func newOpenAIChannel(f *Factory, group *models.Group) (ChannelProxy, error) {
}, nil
}
// BuildUpstreamURL constructs the target URL for the OpenAI service.
func (ch *OpenAIChannel) BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error) {
// Use the weighted round-robin selection from the base channel.
// This method already handles parsing the group's Upstreams JSON.
base := ch.getUpstreamURL()
if base == nil {
// If no upstreams are configured in the group, fallback to a default.
// This can be considered an error or a feature depending on requirements.
// For now, we'll use the official OpenAI URL as a last resort.
base, _ = url.Parse("https://api.openai.com")
}
// It's crucial to create a copy to avoid modifying the cached URL object in BaseChannel.
finalURL := *base
// The originalURL.Path contains the full path, e.g., "/proxy/openai/v1/chat/completions".
// We need to strip the proxy prefix to get the correct upstream path.
proxyPrefix := "/proxy/" + group.Name
if strings.HasPrefix(originalURL.Path, proxyPrefix) {
finalURL.Path = strings.TrimPrefix(originalURL.Path, proxyPrefix)
} else {
// Fallback for safety, though this case should ideally not be hit.
finalURL.Path = originalURL.Path
}
finalURL.RawQuery = originalURL.RawQuery
return finalURL.String(), nil
}
// ModifyRequest sets the Authorization header for the OpenAI service.
func (ch *OpenAIChannel) ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group) {