api2gpt-mid/controller/relay.go

package controller

import (
	"api2gpt-mid/common"
	"api2gpt-mid/model"
	"api2gpt-mid/service"
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"
)

func Images(c *gin.Context) {
	var imagesRequest model.ImagesRequest
	var modelStr = "images-generations"
	if err := c.ShouldBindJSON(&imagesRequest); err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}
	auth := c.Request.Header.Get("Authorization")
	key := auth[7:]

	serverInfo, err := service.CheckBlanceForImages(key, modelStr, imagesRequest.N)
	if err != nil {
		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, modelStr, err)
		return
	}

	log.Printf("请求的KEY: %v Model: %v", key, modelStr)

	remote, err := url.Parse(serverInfo.ServerAddress)
	if err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(remote)
	newReqBody, err := json.Marshal(imagesRequest)
	if err != nil {
		log.Printf("http request err: %v", err)
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	proxy.Director = func(req *http.Request) {
		req.Header = c.Request.Header
		req.Host = remote.Host
		req.URL.Scheme = remote.Scheme
		req.URL.Host = remote.Host
		req.URL.Path = c.Request.URL.Path
		req.ContentLength = int64(len(newReqBody))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept-Encoding", "")
		serverKey := serverInfo.AvailableKey
		keyList := strings.Split(serverInfo.AvailableKey, ",")
		if len(keyList) > 1 {
			// 随机数种子
			rand.Seed(time.Now().UnixNano())
			// 从数组中随机选择一个元素
			serverKey = keyList[rand.Intn(len(keyList))]
		}
		req.Header.Set("Authorization", "Bearer "+serverKey)
		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
	}
	sss, err := json.Marshal(imagesRequest)
	log.Printf("开始处理返回逻辑 %d", string(sss))

	proxy.ModifyResponse = func(resp *http.Response) error {
		if resp.StatusCode != http.StatusOK {
			//退回预扣除的余额
			service.CheckBlanceReturnForImages(key, modelStr, imagesRequest.N)
			return nil
		}
		resp.Header.Set("Openai-Organization", "api2gpt")
		var imagesResponse model.ImagesResponse
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Printf("读取返回数据出错: %v", err)
			return err
		}
		json.Unmarshal(body, &imagesResponse)
		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
		log.Printf("image size:  %v", len(imagesResponse.Data))
		timestamp := time.Now().Unix()
		timestampID := "img-" + fmt.Sprintf("%d", timestamp)
		//消费余额
		service.ConsumptionForImages(key, modelStr, imagesRequest.N, len(imagesResponse.Data), timestampID)
		return nil
	}

	proxy.ServeHTTP(c.Writer, c.Request)

}

func Edit(c *gin.Context) {
	var prompt_tokens int
	var complate_tokens int
	var total_tokens int
	var chatRequest model.ChatRequest
	if err := c.ShouldBindJSON(&chatRequest); err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	auth := c.Request.Header.Get("Authorization")
	key := auth[7:]
	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
	serverInfo, err := service.CheckBlance(key, chatRequest.Model)
	if err != nil {
		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
		return
	}

	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)

	remote, err := url.Parse(serverInfo.ServerAddress)
	if err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(remote)
	newReqBody, err := json.Marshal(chatRequest)
	if err != nil {
		log.Printf("http request err: %v", err)
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	proxy.Director = func(req *http.Request) {
		req.Header = c.Request.Header
		req.Host = remote.Host
		req.URL.Scheme = remote.Scheme
		req.URL.Host = remote.Host
		req.URL.Path = c.Request.URL.Path
		req.ContentLength = int64(len(newReqBody))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept-Encoding", "")
		serverKey := serverInfo.AvailableKey
		keyList := strings.Split(serverInfo.AvailableKey, ",")
		if len(keyList) > 1 {
			// 随机数种子
			rand.Seed(time.Now().UnixNano())
			// 从数组中随机选择一个元素
			serverKey = keyList[rand.Intn(len(keyList))]
		}
		req.Header.Set("Authorization", "Bearer "+serverKey)
		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
	}
	sss, err := json.Marshal(chatRequest)
	log.Printf("开始处理返回逻辑 %d", string(sss))

	proxy.ModifyResponse = func(resp *http.Response) error {
		resp.Header.Set("Openai-Organization", "api2gpt")
		var chatResponse model.ChatResponse
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Printf("读取返回数据出错: %v", err)
			return err
		}
		json.Unmarshal(body, &chatResponse)
		prompt_tokens = chatResponse.Usage.PromptTokens
		complate_tokens = chatResponse.Usage.CompletionTokens
		total_tokens = chatResponse.Usage.TotalTokens
		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
		log.Printf("prompt_tokens:  %v complate_tokens: %v total_tokens: %v", prompt_tokens, complate_tokens, total_tokens)
		timestamp := time.Now().Unix()
		timestampID := "edit-" + fmt.Sprintf("%d", timestamp)
		//消费余额
		service.Consumption(key, chatRequest.Model, prompt_tokens, 0, total_tokens, timestampID)
		return nil
	}

	proxy.ServeHTTP(c.Writer, c.Request)

}

func Embeddings(c *gin.Context) {
	var prompt_tokens int
	var total_tokens int
	var chatRequest model.ChatRequest
	if err := c.ShouldBindJSON(&chatRequest); err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	auth := c.Request.Header.Get("Authorization")
	//key := strings.Trim(auth, "Bearer ")
	key := auth[7:]
	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
	serverInfo, err := service.CheckBlance(key, chatRequest.Model)
	if err != nil {
		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
		return
	}

	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)

	remote, err := url.Parse(serverInfo.ServerAddress)
	if err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(remote)
	newReqBody, err := json.Marshal(chatRequest)
	if err != nil {
		log.Printf("http request err: %v", err)
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	proxy.Director = func(req *http.Request) {
		req.Header = c.Request.Header
		req.Host = remote.Host
		req.URL.Scheme = remote.Scheme
		req.URL.Host = remote.Host
		req.URL.Path = c.Request.URL.Path
		req.ContentLength = int64(len(newReqBody))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept-Encoding", "")
		serverKey := serverInfo.AvailableKey
		keyList := strings.Split(serverInfo.AvailableKey, ",")
		if len(keyList) > 1 {
			// 随机数种子
			rand.Seed(time.Now().UnixNano())
			// 从数组中随机选择一个元素
			serverKey = keyList[rand.Intn(len(keyList))]
		}
		req.Header.Set("Authorization", "Bearer "+serverKey)
		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
	}
	sss, err := json.Marshal(chatRequest)
	log.Printf("开始处理返回逻辑 %d", string(sss))

	proxy.ModifyResponse = func(resp *http.Response) error {
		resp.Header.Set("Openai-Organization", "api2gpt")
		var chatResponse model.ChatResponse
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Printf("读取返回数据出错: %v", err)
			return err
		}
		json.Unmarshal(body, &chatResponse)
		prompt_tokens = chatResponse.Usage.PromptTokens
		total_tokens = chatResponse.Usage.TotalTokens
		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
		log.Printf("prompt_tokens:  %v total_tokens: %v", prompt_tokens, total_tokens)
		timestamp := time.Now().Unix()
		timestampID := "emb-" + fmt.Sprintf("%d", timestamp)
		//消费余额
		service.Consumption(key, chatRequest.Model, prompt_tokens, 0, total_tokens, timestampID)
		return nil
	}

	proxy.ServeHTTP(c.Writer, c.Request)

}

func Completions(c *gin.Context) {
	var prompt_tokens int
	var completion_tokens int
	var total_tokens int
	var chatRequest model.ChatRequest
	if err := c.ShouldBindJSON(&chatRequest); err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	auth := c.Request.Header.Get("Authorization")
	key := auth[7:]
	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
	serverInfo, err := service.CheckBlance(key, chatRequest.Model)
	if err != nil {
		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
		return
	}

	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)

	remote, err := url.Parse(serverInfo.ServerAddress)
	if err != nil {
		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(remote)
	newReqBody, err := json.Marshal(chatRequest)
	if err != nil {
		log.Printf("http request err: %v", err)
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	proxy.Director = func(req *http.Request) {
		req.Header = c.Request.Header
		req.Host = remote.Host
		req.URL.Scheme = remote.Scheme
		req.URL.Host = remote.Host
		req.URL.Path = c.Request.URL.Path
		req.ContentLength = int64(len(newReqBody))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept-Encoding", "")
		serverKey := serverInfo.AvailableKey
		keyList := strings.Split(serverInfo.AvailableKey, ",")
		if len(keyList) > 1 {
			// 随机数种子
			rand.Seed(time.Now().UnixNano())
			// 从数组中随机选择一个元素
			serverKey = keyList[rand.Intn(len(keyList))]
		}
		req.Header.Set("Authorization", "Bearer "+serverKey)
		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
	}
	sss, err := json.Marshal(chatRequest)
	if err != nil {
		log.Printf("chatRequest 转化出错 %v", err)
	}
	log.Printf("开始处理返回逻辑: %v", string(sss))
	if chatRequest.Stream {
		// 流式回应，处理
		proxy.ModifyResponse = func(resp *http.Response) error {
			log.Printf("流式回应 http status code: %v", resp.StatusCode)
			if resp.StatusCode != http.StatusOK {
				//退回预扣除的余额
				service.CheckBlanceReturn(key, chatRequest.Model)
				return nil
			}
			chatRequestId := ""
			reqContent := ""
			reader := bufio.NewReader(resp.Body)
			headers := resp.Header
			for k, v := range headers {
				c.Writer.Header().Set(k, v[0])
			}
			c.Writer.Header().Set("Openai-Organization", "api2gpt")
			for {
				chunk, err := reader.ReadBytes('\n')
				if err != nil {
					if err == io.EOF {
						break
					}
					log.Printf("流式回应，处理 err %v:", err.Error())
					break
					//return err
				}
				var chatResponse model.ChatResponse
				//去除回应中的data:前缀
				var trimStr = strings.Trim(string(chunk), "data: ")
				if trimStr != "\n" {
					json.Unmarshal([]byte(trimStr), &chatResponse)
					if chatResponse.Choices != nil {
						if chatResponse.Choices[0].Text != "" {
							reqContent += chatResponse.Choices[0].Text
						} else {
							reqContent += chatResponse.Choices[0].Delta.Content
						}

						chatRequestId = chatResponse.Id
					}

					// 写回数据
					_, err = c.Writer.Write([]byte(string(chunk) + "\n"))
					if err != nil {
						log.Printf("写回数据 err: %v", err.Error())
						return err
					}
					c.Writer.(http.Flusher).Flush()
				}
			}
			if chatRequest.Model == "text-davinci-003" {
				prompt_tokens = common.NumTokensFromString(chatRequest.Prompt, chatRequest.Model)
			} else {
				prompt_tokens = common.NumTokensFromMessages(chatRequest.Messages, chatRequest.Model)
			}
			completion_tokens = common.NumTokensFromString(reqContent, chatRequest.Model)
			log.Printf("返回内容：%v", reqContent)
			total_tokens = prompt_tokens + completion_tokens
			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", prompt_tokens, completion_tokens, total_tokens)
			//消费余额
			service.Consumption(key, chatRequest.Model, prompt_tokens, completion_tokens, total_tokens, chatRequestId)
			return nil
		}
	} else {
		// 非流式回应，处理
		proxy.ModifyResponse = func(resp *http.Response) error {
			resp.Header.Set("Openai-Organization", "api2gpt")
			var chatResponse model.ChatResponse
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Printf("非流式回应，处理 err: %v", err)
				return err
			}
			json.Unmarshal(body, &chatResponse)
			prompt_tokens = chatResponse.Usage.PromptTokens
			completion_tokens = chatResponse.Usage.CompletionTokens
			total_tokens = chatResponse.Usage.TotalTokens
			resp.Body = ioutil.NopCloser(bytes.NewReader(body))
			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", prompt_tokens, completion_tokens, total_tokens)
			//消费余额
			service.Consumption(key, chatRequest.Model, prompt_tokens, completion_tokens, total_tokens, chatResponse.Id)
			return nil
		}
	}

	proxy.ServeHTTP(c.Writer, c.Request)
}

// 针对接口预检OPTIONS的处理
func HandleOptions(c *gin.Context) {
	// BUGFIX: fix options request, see https://github.com/diemus/azure-openai-proxy/issues/1
	c.Header("Access-Control-Allow-Origin", "*")
	c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
	c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
	c.Status(200)
	return
}