Просмотр исходного кода

fix: fix Azure channel not working in stream mode (#57)

JustSong 2 лет назад
Родитель
Сommit
2edd52e851
1 измененных файлов с 9 добавлено и 4 удалено
  1. 9 4
      controller/relay.go

+ 9 - 4
controller/relay.go

@@ -94,10 +94,12 @@ func relayHelper(c *gin.Context) error {
 	if channelType == common.ChannelTypeAzure {
 	if channelType == common.ChannelTypeAzure {
 		// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
 		// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
 		query := c.Request.URL.Query()
 		query := c.Request.URL.Query()
-		if query.Get("api-version") == "" {
-			apiVersion := c.GetString("api_version")
-			requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion)
+		apiVersion := query.Get("api-version")
+		if apiVersion == "" {
+			apiVersion = c.GetString("api_version")
 		}
 		}
+		requestURL := strings.Split(requestURL, "?")[0]
+		requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion)
 		baseURL = c.GetString("base_url")
 		baseURL = c.GetString("base_url")
 		task := strings.TrimPrefix(requestURL, "/v1/")
 		task := strings.TrimPrefix(requestURL, "/v1/")
 		model_ := textRequest.Model
 		model_ := textRequest.Model
@@ -186,7 +188,7 @@ func relayHelper(c *gin.Context) error {
 				data := scanner.Text()
 				data := scanner.Text()
 				dataChan <- data
 				dataChan <- data
 				data = data[6:]
 				data = data[6:]
-				if data != "[DONE]" {
+				if !strings.HasPrefix(data, "[DONE]") {
 					var streamResponse StreamResponse
 					var streamResponse StreamResponse
 					err = json.Unmarshal([]byte(data), &streamResponse)
 					err = json.Unmarshal([]byte(data), &streamResponse)
 					if err != nil {
 					if err != nil {
@@ -207,6 +209,9 @@ func relayHelper(c *gin.Context) error {
 		c.Stream(func(w io.Writer) bool {
 		c.Stream(func(w io.Writer) bool {
 			select {
 			select {
 			case data := <-dataChan:
 			case data := <-dataChan:
+				if strings.HasPrefix(data, "data: [DONE]") {
+					data = data[:len(data)-1]
+				}
 				c.Render(-1, common.CustomEvent{Data: data})
 				c.Render(-1, common.CustomEvent{Data: data})
 				return true
 				return true
 			case <-stopChan:
 			case <-stopChan: