mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 21:00:52 +08:00
Improve getClient logic with optional content generation flag
- Added `isGenerateContent` optional parameter to `getClient` for conditional client selection. - Updated `gemini-handlers` to utilize the new parameter for enhanced control.
This commit is contained in:
@@ -266,7 +266,7 @@ func (h *APIHandlers) geminiCountTokens(c *gin.Context, rawJson []byte) {
|
|||||||
|
|
||||||
for {
|
for {
|
||||||
var errorResponse *client.ErrorMessage
|
var errorResponse *client.ErrorMessage
|
||||||
cliClient, errorResponse = h.getClient(modelName)
|
cliClient, errorResponse = h.getClient(modelName, false)
|
||||||
if errorResponse != nil {
|
if errorResponse != nil {
|
||||||
c.Status(errorResponse.StatusCode)
|
c.Status(errorResponse.StatusCode)
|
||||||
_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
|
_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
|
||||||
|
|||||||
@@ -85,7 +85,7 @@ func (h *APIHandlers) Models(c *gin.Context) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *APIHandlers) getClient(modelName string) (*client.Client, *client.ErrorMessage) {
|
func (h *APIHandlers) getClient(modelName string, isGenerateContent ...bool) (*client.Client, *client.ErrorMessage) {
|
||||||
if len(h.cliClients) == 0 {
|
if len(h.cliClients) == 0 {
|
||||||
return nil, &client.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("no clients available")}
|
return nil, &client.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("no clients available")}
|
||||||
}
|
}
|
||||||
@@ -95,8 +95,10 @@ func (h *APIHandlers) getClient(modelName string) (*client.Client, *client.Error
|
|||||||
// Lock the mutex to update the last used client index
|
// Lock the mutex to update the last used client index
|
||||||
mutex.Lock()
|
mutex.Lock()
|
||||||
startIndex := lastUsedClientIndex
|
startIndex := lastUsedClientIndex
|
||||||
currentIndex := (startIndex + 1) % len(h.cliClients)
|
if (len(isGenerateContent) > 0 && isGenerateContent[0]) || len(isGenerateContent) == 0 {
|
||||||
lastUsedClientIndex = currentIndex
|
currentIndex := (startIndex + 1) % len(h.cliClients)
|
||||||
|
lastUsedClientIndex = currentIndex
|
||||||
|
}
|
||||||
mutex.Unlock()
|
mutex.Unlock()
|
||||||
|
|
||||||
// Reorder the client to start from the last used index
|
// Reorder the client to start from the last used index
|
||||||
|
|||||||
Reference in New Issue
Block a user