fix: fix getPreConsumedQuota (#1312)

* fix zhipu embedding error when input is array but not string

* fix: only use the first one

---------

Co-authored-by: 蔡新疆 <cxj@icc.link>
Co-authored-by: JustSong <songquanpeng@foxmail.com>

controller/channel-test.go

@@ -5,6 +5,15 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
 	"github.com/songquanpeng/one-api/common/config"
 	"github.com/songquanpeng/one-api/common/ctxkey"
 	"github.com/songquanpeng/one-api/common/logger"
@@ -18,14 +27,6 @@ import (
 	"github.com/songquanpeng/one-api/relay/meta"
 	relaymodel "github.com/songquanpeng/one-api/relay/model"
 	"github.com/songquanpeng/one-api/relay/relaymode"
-	"io"
-	"net/http"
-	"net/http/httptest"
-	"net/url"
-	"strconv"
-	"strings"
-	"sync"
-	"time"
 
 	"github.com/gin-gonic/gin"
 )
@@ -69,6 +70,7 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
 	adaptor.Init(meta)
 	var modelName string
 	modelList := adaptor.GetModelList()
+	modelMap := channel.GetModelMapping()
 	if len(modelList) != 0 {
 		modelName = modelList[0]
 	}
@@ -77,6 +79,9 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
 		if len(modelNames) > 0 {
 			modelName = modelNames[0]
 		}
+		if modelMap != nil && modelMap[modelName] != "" {
+			modelName = modelMap[modelName]
+		}
 	}
 	request := buildTestRequest()
 	request.Model = modelName

relay/adaptor/anthropic/main.go

@@ -4,6 +4,10 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/common"
 	"github.com/songquanpeng/one-api/common/helper"
@@ -11,9 +15,6 @@ import (
 	"github.com/songquanpeng/one-api/common/logger"
 	"github.com/songquanpeng/one-api/relay/adaptor/openai"
 	"github.com/songquanpeng/one-api/relay/model"
-	"io"
-	"net/http"
-	"strings"
 )
 
 func stopReasonClaude2OpenAI(reason *string) string {
@@ -176,10 +177,10 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 			if len(data) < 6 {
 				continue
 			}
-			if !strings.HasPrefix(data, "data: ") {
+			if !strings.HasPrefix(data, "data:") {
 				continue
 			}
-			data = strings.TrimPrefix(data, "data: ")
+			data = strings.TrimPrefix(data, "data:")
 			dataChan <- data
 		}
 		stopChan <- true
@@ -192,7 +193,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 		select {
 		case data := <-dataChan:
 			// some implementations may add \r at the end of data
-			data = strings.TrimSuffix(data, "\r")
+			data = strings.TrimSpace(data)
 			var claudeResponse StreamResponse
 			err := json.Unmarshal([]byte(data), &claudeResponse)
 			if err != nil {

relay/adaptor/gemini/adaptor.go

@@ -3,6 +3,9 @@ package gemini
 import (
 	"errors"
 	"fmt"
+	"io"
+	"net/http"
+
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/common/config"
 	"github.com/songquanpeng/one-api/common/helper"
@@ -10,8 +13,6 @@ import (
 	"github.com/songquanpeng/one-api/relay/adaptor/openai"
 	"github.com/songquanpeng/one-api/relay/meta"
 	"github.com/songquanpeng/one-api/relay/model"
-	"io"
-	"net/http"
 )
 
 type Adaptor struct {
@@ -25,7 +26,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
 	version := helper.AssignOrDefault(meta.Config.APIVersion, config.GeminiVersion)
 	action := "generateContent"
 	if meta.IsStream {
-		action = "streamGenerateContent"
+		action = "streamGenerateContent?alt=sse"
 	}
 	return fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, meta.ActualModelName, action), nil
 }

relay/adaptor/gemini/main.go

@@ -232,8 +232,6 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *ChatResponse) *openai.ChatC
 
 func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) {
 	responseText := ""
-	dataChan := make(chan string)
-	stopChan := make(chan bool)
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
 		if atEOF && len(data) == 0 {
@@ -247,14 +245,16 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 		}
 		return 0, nil, nil
 	})
+	dataChan := make(chan string)
+	stopChan := make(chan bool)
 	go func() {
 		for scanner.Scan() {
 			data := scanner.Text()
 			data = strings.TrimSpace(data)
-			if !strings.HasPrefix(data, "\"text\": \"") {
+			if !strings.HasPrefix(data, "data: ") {
 				continue
 			}
-			data = strings.TrimPrefix(data, "\"text\": \"")
+			data = strings.TrimPrefix(data, "data: ")
 			data = strings.TrimSuffix(data, "\"")
 			dataChan <- data
 		}
@@ -264,23 +264,17 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:
-			// this is used to prevent annoying \ related format bug
-			data = fmt.Sprintf("{\"content\": \"%s\"}", data)
-			type dummyStruct struct {
-				Content string `json:"content"`
+			var geminiResponse ChatResponse
+			err := json.Unmarshal([]byte(data), &geminiResponse)
+			if err != nil {
+				logger.SysError("error unmarshalling stream response: " + err.Error())
+				return true
 			}
-			var dummy dummyStruct
-			err := json.Unmarshal([]byte(data), &dummy)
-			responseText += dummy.Content
-			var choice openai.ChatCompletionsStreamResponseChoice
-			choice.Delta.Content = dummy.Content
-			response := openai.ChatCompletionsStreamResponse{
-				Id:      fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
-				Object:  "chat.completion.chunk",
-				Created: helper.GetTimestamp(),
-				Model:   "gemini-pro",
-				Choices: []openai.ChatCompletionsStreamResponseChoice{choice},
+			response := streamResponseGeminiChat2OpenAI(&geminiResponse)
+			if response == nil {
+				return true
 			}
+			responseText += response.Choices[0].Delta.StringContent()
 			jsonResponse, err := json.Marshal(response)
 			if err != nil {
 				logger.SysError("error marshalling stream response: " + err.Error())

relay/adaptor/ollama/main.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/image"
 	"github.com/songquanpeng/one-api/common/logger"
 	"github.com/songquanpeng/one-api/relay/adaptor/openai"
 	"github.com/songquanpeng/one-api/relay/constant"
@@ -32,9 +33,22 @@ func ConvertRequest(request model.GeneralOpenAIRequest) *ChatRequest {
 		Stream: request.Stream,
 	}
 	for _, message := range request.Messages {
+		openaiContent := message.ParseContent()
+		var imageUrls []string
+		var contentText string
+		for _, part := range openaiContent {
+			switch part.Type {
+			case model.ContentTypeText:
+				contentText = part.Text
+			case model.ContentTypeImageURL:
+				_, data, _ := image.GetImageFromUrl(part.ImageURL.Url)
+				imageUrls = append(imageUrls, data)
+			}
+		}
 		ollamaRequest.Messages = append(ollamaRequest.Messages, Message{
 			Role:    message.Role,
-			Content: message.StringContent(),
+			Content: contentText,
+			Images:  imageUrls,
 		})
 	}
 	return &ollamaRequest

relay/adaptor/zhipu/adaptor.go

@@ -62,8 +62,8 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
 	}
 	switch relayMode {
 	case relaymode.Embeddings:
-		baiduEmbeddingRequest := ConvertEmbeddingRequest(*request)
-		return baiduEmbeddingRequest, nil
+		baiduEmbeddingRequest, err := ConvertEmbeddingRequest(*request)
+		return baiduEmbeddingRequest, err
 	default:
 		// TopP (0.0, 1.0)
 		request.TopP = math.Min(0.99, request.TopP)
@@ -129,11 +129,15 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Met
 	return
 }
 
-func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) *EmbeddingRequest {
-	return &EmbeddingRequest{
-		Model: "embedding-2",
-		Input: request.Input.(string),
+func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) (*EmbeddingRequest, error) {
+	inputs := request.ParseInput()
+	if len(inputs) != 1 {
+		return nil, errors.New("invalid input length, zhipu only support one input")
 	}
+	return &EmbeddingRequest{
+		Model: request.Model,
+		Input: inputs[0],
+	}, nil
 }
 
 func (a *Adaptor) GetModelList() []string {

relay/controller/helper.go

@@ -125,9 +125,9 @@ func getPromptTokens(textRequest *relaymodel.GeneralOpenAIRequest, relayMode int
 }
 
 func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64) int64 {
-	preConsumedTokens := config.PreConsumedQuota
+	preConsumedTokens := config.PreConsumedQuota + int64(promptTokens)
 	if textRequest.MaxTokens != 0 {
-		preConsumedTokens = int64(promptTokens) + int64(textRequest.MaxTokens)
+		preConsumedTokens += int64(textRequest.MaxTokens)
 	}
 	return int64(float64(preConsumedTokens) * ratio)
 }