chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-18 04:10:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improve getClient logic with optional content generation flag

Browse Source 
- 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:
Luis Pater
2025-07-27 02:30:08 +08:00
parent 00bd6a3e46
commit 3b4634e2dc
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/api/gemini-handlers.go
View File

@@ -266,7 +266,7 @@ func (h *APIHandlers) geminiCountTokens(c *gin.Context, rawJson []byte) {
for {
var errorResponse *client.ErrorMessage
cliClient, errorResponse = h.getClient(modelName)
cliClient, errorResponse = h.getClient(modelName, false)
if errorResponse != nil {
c.Status(errorResponse.StatusCode)
_, _ = fmt.Fprint(c.Writer, errorResponse.Error)

8
internal/api/handlers.go
View File

@@ -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 {
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
mutex.Lock()
startIndex := lastUsedClientIndex
currentIndex := (startIndex + 1) % len(h.cliClients)
lastUsedClientIndex = currentIndex
if (len(isGenerateContent) > 0 && isGenerateContent[0]) || len(isGenerateContent) == 0 {
currentIndex := (startIndex + 1) % len(h.cliClients)
lastUsedClientIndex = currentIndex
}
mutex.Unlock()
// Reorder the client to start from the last used index