Переглянути джерело

fix: charge local input tokens when Gemini returns empty response

Seefs 1 тиждень тому
батько
коміт
ecebd619a4

+ 1 - 2
relay/channel/gemini/relay-gemini.go

@@ -1258,8 +1258,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
 	}
 
 	if usage.CompletionTokens <= 0 {
-		str := responseText.String()
-		if len(str) > 0 {
+		if info.ReceivedResponseCount > 0 {
 			usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
 		} else {
 			usage = &dto.Usage{}

+ 1 - 0
relay/common/relay_info.go

@@ -113,6 +113,7 @@ type RelayInfo struct {
 	UserQuota              int
 	RelayFormat            types.RelayFormat
 	SendResponseCount      int
+	ReceivedResponseCount  int
 	FinalPreConsumedQuota  int // 最终预消耗的配额
 	// BillingSource indicates whether this request is billed from wallet quota or subscription.
 	// "" or "wallet" => wallet; "subscription" => subscription

+ 7 - 7
relay/helper/stream_scanner.go

@@ -90,10 +90,10 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
 
 		// 等待所有 goroutine 退出,最多等待5秒
 		done := make(chan struct{})
-		go func() {
+		gopool.Go(func() {
 			wg.Wait()
 			close(done)
-		}()
+		})
 
 		select {
 		case <-done:
@@ -138,11 +138,11 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
 				case <-pingTicker.C:
 					// 使用超时机制防止写操作阻塞
 					done := make(chan error, 1)
-					go func() {
+					gopool.Go(func() {
 						writeMutex.Lock()
 						defer writeMutex.Unlock()
 						done <- PingData(c)
-					}()
+					})
 
 					select {
 					case err := <-done:
@@ -219,14 +219,14 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
 			data = strings.TrimSuffix(data, "\r")
 			if !strings.HasPrefix(data, "[DONE]") {
 				info.SetFirstResponseTime()
-
+				info.ReceivedResponseCount++
 				// 使用超时机制防止写操作阻塞
 				done := make(chan bool, 1)
-				go func() {
+				gopool.Go(func() {
 					writeMutex.Lock()
 					defer writeMutex.Unlock()
 					done <- dataHandler(data)
-				}()
+				})
 
 				select {
 				case success := <-done: