mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 04:50:52 +08:00
Improved the /v1/models endpoint
This commit is contained in:
@@ -23,6 +23,7 @@ import (
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/misc"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -55,6 +56,10 @@ type ClaudeClient struct {
|
||||
// - *ClaudeClient: A new Claude client instance.
|
||||
func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeClient {
|
||||
httpClient := util.SetProxy(cfg, &http.Client{})
|
||||
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("claude-%d", time.Now().UnixNano())
|
||||
|
||||
client := &ClaudeClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -67,6 +72,10 @@ func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeC
|
||||
apiKeyIndex: -1,
|
||||
}
|
||||
|
||||
// Initialize model registry and register Claude models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("claude", registry.GetClaudeModels())
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -82,6 +91,10 @@ func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeC
|
||||
// - *ClaudeClient: A new Claude client instance.
|
||||
func NewClaudeClientWithKey(cfg *config.Config, apiKeyIndex int) *ClaudeClient {
|
||||
httpClient := util.SetProxy(cfg, &http.Client{})
|
||||
|
||||
// Generate unique client ID for API key client
|
||||
clientID := fmt.Sprintf("claude-apikey-%d-%d", apiKeyIndex, time.Now().UnixNano())
|
||||
|
||||
client := &ClaudeClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -94,6 +107,10 @@ func NewClaudeClientWithKey(cfg *config.Config, apiKeyIndex int) *ClaudeClient {
|
||||
apiKeyIndex: apiKeyIndex,
|
||||
}
|
||||
|
||||
// Initialize model registry and register Claude models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("claude", registry.GetClaudeModels())
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -174,10 +191,14 @@ func (c *ClaudeClient) SendRawMessage(ctx context.Context, modelName string, raw
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -234,11 +255,15 @@ func (c *ClaudeClient) SendRawMessageStream(ctx context.Context, modelName strin
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/luispater/CLIProxyAPI/internal/auth"
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
)
|
||||
|
||||
// ClientBase provides a common base structure for all AI API clients.
|
||||
@@ -34,6 +35,12 @@ type ClientBase struct {
|
||||
// modelQuotaExceeded tracks when models have exceeded their quota.
|
||||
// The map key is the model name, and the value is the time when the quota was exceeded.
|
||||
modelQuotaExceeded map[string]*time.Time
|
||||
|
||||
// clientID is the unique identifier for this client instance.
|
||||
clientID string
|
||||
|
||||
// modelRegistry is the global model registry for tracking model availability.
|
||||
modelRegistry *registry.ModelRegistry
|
||||
}
|
||||
|
||||
// GetRequestMutex returns the mutex used to synchronize requests for this client.
|
||||
@@ -71,3 +78,50 @@ func (c *ClientBase) AddAPIResponseData(ctx context.Context, line []byte) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// InitializeModelRegistry initializes the model registry for this client
|
||||
// This should be called by all client implementations during construction
|
||||
func (c *ClientBase) InitializeModelRegistry(clientID string) {
|
||||
c.clientID = clientID
|
||||
c.modelRegistry = registry.GetGlobalRegistry()
|
||||
}
|
||||
|
||||
// RegisterModels registers the models that this client can provide
|
||||
// Parameters:
|
||||
// - provider: The provider name (e.g., "gemini", "claude", "openai")
|
||||
// - models: The list of models this client supports
|
||||
func (c *ClientBase) RegisterModels(provider string, models []*registry.ModelInfo) {
|
||||
if c.modelRegistry != nil && c.clientID != "" {
|
||||
c.modelRegistry.RegisterClient(c.clientID, provider, models)
|
||||
}
|
||||
}
|
||||
|
||||
// UnregisterClient removes this client from the model registry
|
||||
func (c *ClientBase) UnregisterClient() {
|
||||
if c.modelRegistry != nil && c.clientID != "" {
|
||||
c.modelRegistry.UnregisterClient(c.clientID)
|
||||
}
|
||||
}
|
||||
|
||||
// SetModelQuotaExceeded marks a model as quota exceeded in the registry
|
||||
// Parameters:
|
||||
// - modelID: The model that exceeded quota
|
||||
func (c *ClientBase) SetModelQuotaExceeded(modelID string) {
|
||||
if c.modelRegistry != nil && c.clientID != "" {
|
||||
c.modelRegistry.SetModelQuotaExceeded(c.clientID, modelID)
|
||||
}
|
||||
}
|
||||
|
||||
// ClearModelQuotaExceeded clears quota exceeded status for a model
|
||||
// Parameters:
|
||||
// - modelID: The model to clear quota status for
|
||||
func (c *ClientBase) ClearModelQuotaExceeded(modelID string) {
|
||||
if c.modelRegistry != nil && c.clientID != "" {
|
||||
c.modelRegistry.ClearModelQuotaExceeded(c.clientID, modelID)
|
||||
}
|
||||
}
|
||||
|
||||
// GetClientID returns the unique identifier for this client
|
||||
func (c *ClientBase) GetClientID() string {
|
||||
return c.clientID
|
||||
}
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -50,6 +51,10 @@ type CodexClient struct {
|
||||
// - error: An error if the client creation fails.
|
||||
func NewCodexClient(cfg *config.Config, ts *codex.CodexTokenStorage) (*CodexClient, error) {
|
||||
httpClient := util.SetProxy(cfg, &http.Client{})
|
||||
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("codex-%d", time.Now().UnixNano())
|
||||
|
||||
client := &CodexClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -61,6 +66,10 @@ func NewCodexClient(cfg *config.Config, ts *codex.CodexTokenStorage) (*CodexClie
|
||||
codexAuth: codex.NewCodexAuth(cfg),
|
||||
}
|
||||
|
||||
// Initialize model registry and register OpenAI models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("codex", registry.GetOpenAIModels())
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
@@ -123,10 +132,14 @@ func (c *CodexClient) SendRawMessage(ctx context.Context, modelName string, rawJ
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -184,11 +197,15 @@ func (c *CodexClient) SendRawMessageStream(ctx context.Context, modelName string
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -57,6 +58,9 @@ type GeminiCLIClient struct {
|
||||
// Returns:
|
||||
// - *GeminiCLIClient: A new Gemini CLI client instance.
|
||||
func NewGeminiCLIClient(httpClient *http.Client, ts *geminiAuth.GeminiTokenStorage, cfg *config.Config) *GeminiCLIClient {
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("gemini-cli-%d", time.Now().UnixNano())
|
||||
|
||||
client := &GeminiCLIClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -66,6 +70,11 @@ func NewGeminiCLIClient(httpClient *http.Client, ts *geminiAuth.GeminiTokenStora
|
||||
modelQuotaExceeded: make(map[string]*time.Time),
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize model registry and register Gemini models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("gemini-cli", registry.GetGeminiModels())
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -426,6 +435,8 @@ func (c *GeminiCLIClient) SendRawTokenCount(ctx context.Context, modelName strin
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
if c.cfg.QuotaExceeded.SwitchPreviewModel {
|
||||
continue
|
||||
}
|
||||
@@ -433,6 +444,8 @@ func (c *GeminiCLIClient) SendRawTokenCount(ctx context.Context, modelName strin
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -485,6 +498,8 @@ func (c *GeminiCLIClient) SendRawMessage(ctx context.Context, modelName string,
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
if c.cfg.QuotaExceeded.SwitchPreviewModel {
|
||||
continue
|
||||
}
|
||||
@@ -492,6 +507,8 @@ func (c *GeminiCLIClient) SendRawMessage(ctx context.Context, modelName string,
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -562,6 +579,8 @@ func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName st
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
if c.cfg.QuotaExceeded.SwitchPreviewModel {
|
||||
continue
|
||||
}
|
||||
@@ -570,6 +589,8 @@ func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName st
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
break
|
||||
}
|
||||
defer func() {
|
||||
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -44,6 +45,9 @@ type GeminiClient struct {
|
||||
// Returns:
|
||||
// - *GeminiClient: A new Gemini client instance.
|
||||
func NewGeminiClient(httpClient *http.Client, cfg *config.Config, glAPIKey string) *GeminiClient {
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("gemini-apikey-%s-%d", glAPIKey[:8], time.Now().UnixNano()) // Use first 8 chars of API key
|
||||
|
||||
client := &GeminiClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -53,6 +57,11 @@ func NewGeminiClient(httpClient *http.Client, cfg *config.Config, glAPIKey strin
|
||||
},
|
||||
glAPIKey: glAPIKey,
|
||||
}
|
||||
|
||||
// Initialize model registry and register Gemini models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("gemini", registry.GetGeminiModels())
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -195,10 +204,14 @@ func (c *GeminiClient) SendRawTokenCount(ctx context.Context, modelName string,
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -240,10 +253,14 @@ func (c *GeminiClient) SendRawMessage(ctx context.Context, modelName string, raw
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -297,11 +314,15 @@ func (c *GeminiClient) SendRawMessageStream(ctx context.Context, modelName strin
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -53,6 +54,10 @@ func NewOpenAICompatibilityClient(cfg *config.Config, compatConfig *config.OpenA
|
||||
}
|
||||
|
||||
httpClient := util.SetProxy(cfg, &http.Client{})
|
||||
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("openai-compatibility-%s-%d", compatConfig.Name, time.Now().UnixNano())
|
||||
|
||||
client := &OpenAICompatibilityClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -64,6 +69,25 @@ func NewOpenAICompatibilityClient(cfg *config.Config, compatConfig *config.OpenA
|
||||
currentAPIKeyIndex: 0,
|
||||
}
|
||||
|
||||
// Initialize model registry
|
||||
client.InitializeModelRegistry(clientID)
|
||||
|
||||
// Convert compatibility models to registry models and register them
|
||||
registryModels := make([]*registry.ModelInfo, 0, len(compatConfig.Models))
|
||||
for _, model := range compatConfig.Models {
|
||||
registryModel := ®istry.ModelInfo{
|
||||
ID: model.Alias,
|
||||
Object: "model",
|
||||
Created: time.Now().Unix(),
|
||||
OwnedBy: compatConfig.Name,
|
||||
Type: "openai-compatibility",
|
||||
DisplayName: model.Name,
|
||||
}
|
||||
registryModels = append(registryModels, registryModel)
|
||||
}
|
||||
|
||||
client.RegisterModels(compatConfig.Name, registryModels)
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
@@ -216,10 +240,14 @@ func (c *OpenAICompatibilityClient) SendRawMessage(ctx context.Context, modelNam
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -270,11 +298,15 @@ func (c *OpenAICompatibilityClient) SendRawMessageStream(ctx context.Context, mo
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"github.com/luispater/CLIProxyAPI/internal/config"
|
||||
. "github.com/luispater/CLIProxyAPI/internal/constant"
|
||||
"github.com/luispater/CLIProxyAPI/internal/interfaces"
|
||||
"github.com/luispater/CLIProxyAPI/internal/registry"
|
||||
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
|
||||
"github.com/luispater/CLIProxyAPI/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -49,6 +50,10 @@ type QwenClient struct {
|
||||
// - *QwenClient: A new Qwen client instance.
|
||||
func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient {
|
||||
httpClient := util.SetProxy(cfg, &http.Client{})
|
||||
|
||||
// Generate unique client ID
|
||||
clientID := fmt.Sprintf("qwen-%d", time.Now().UnixNano())
|
||||
|
||||
client := &QwenClient{
|
||||
ClientBase: ClientBase{
|
||||
RequestMutex: &sync.Mutex{},
|
||||
@@ -60,6 +65,10 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient {
|
||||
qwenAuth: qwen.NewQwenAuth(cfg),
|
||||
}
|
||||
|
||||
// Initialize model registry and register Qwen models
|
||||
client.InitializeModelRegistry(clientID)
|
||||
client.RegisterModels("qwen", registry.GetQwenModels())
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -119,10 +128,14 @@ func (c *QwenClient) SendRawMessage(ctx context.Context, modelName string, rawJS
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
bodyBytes, errReadAll := io.ReadAll(respBody)
|
||||
if errReadAll != nil {
|
||||
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
|
||||
@@ -182,11 +195,15 @@ func (c *QwenClient) SendRawMessageStream(ctx context.Context, modelName string,
|
||||
if err.StatusCode == 429 {
|
||||
now := time.Now()
|
||||
c.modelQuotaExceeded[modelName] = &now
|
||||
// Update model registry quota status
|
||||
c.SetModelQuotaExceeded(modelName)
|
||||
}
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
delete(c.modelQuotaExceeded, modelName)
|
||||
// Clear quota status in model registry
|
||||
c.ClearModelQuotaExceeded(modelName)
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
Reference in New Issue
Block a user