reface to openteam
This commit is contained in:
Sakurasan
215
llm/openai/handle_proxy.go
Normal file
215
llm/openai/handle_proxy.go
Normal file
@@ -0,0 +1,215 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"opencatd-open/pkg/tokenizer"
|
||||
"opencatd-open/store"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ChatProxy(c *gin.Context, chatReq *ChatCompletionRequest) {
|
||||
usagelog := store.Tokens{Model: chatReq.Model}
|
||||
|
||||
token, _ := c.Get("localuser")
|
||||
|
||||
lu, err := store.GetUserByToken(token.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": err.Error(),
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
usagelog.UserID = int(lu.ID)
|
||||
|
||||
var prompt string
|
||||
for _, msg := range chatReq.Messages {
|
||||
switch ct := msg.Content.(type) {
|
||||
case string:
|
||||
prompt += "<" + msg.Role + ">: " + msg.Content.(string) + "\n"
|
||||
case []any:
|
||||
for _, item := range ct {
|
||||
if m, ok := item.(map[string]interface{}); ok {
|
||||
if m["type"] == "text" {
|
||||
prompt += "<" + msg.Role + ">: " + m["text"].(string) + "\n"
|
||||
} else if m["type"] == "image_url" {
|
||||
if url, ok := m["image_url"].(map[string]interface{}); ok {
|
||||
fmt.Printf(" URL: %v\n", url["url"])
|
||||
if strings.HasPrefix(url["url"].(string), "http") {
|
||||
fmt.Println("网络图片:", url["url"].(string))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
default:
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": "Invalid content type",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
if len(chatReq.Tools) > 0 {
|
||||
tooljson, _ := json.Marshal(chatReq.Tools)
|
||||
prompt += "<tools>: " + string(tooljson) + "\n"
|
||||
}
|
||||
}
|
||||
switch chatReq.Model {
|
||||
case "gpt-4o", "gpt-4o-mini", "chatgpt-4o-latest":
|
||||
chatReq.MaxTokens = 16384
|
||||
}
|
||||
if chatReq.Stream {
|
||||
chatReq.StreamOptions = &StreamOption{IncludeUsage: true}
|
||||
}
|
||||
|
||||
usagelog.PromptCount = tokenizer.NumTokensFromStr(prompt, chatReq.Model)
|
||||
|
||||
// onekey, err := store.SelectKeyCache("openai")
|
||||
onekey, err := store.SelectKeyCacheByModel(chatReq.Model)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
var req *http.Request
|
||||
|
||||
switch onekey.ApiType {
|
||||
case "github":
|
||||
req, err = http.NewRequest(c.Request.Method, Github_Marketplace, bytes.NewReader(chatReq.ToByteJson()))
|
||||
req.Header = c.Request.Header
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", onekey.Key))
|
||||
case "azure":
|
||||
var buildurl string
|
||||
if onekey.EndPoint != "" {
|
||||
buildurl = fmt.Sprintf("%s/openai/deployments/%s/chat/completions?api-version=%s", onekey.EndPoint, modelmap(chatReq.Model), AzureApiVersion)
|
||||
} else {
|
||||
buildurl = fmt.Sprintf("https://%s.openai.azure.com/openai/deployments/%s/chat/completions?api-version=%s", onekey.ResourceNmae, modelmap(chatReq.Model), AzureApiVersion)
|
||||
}
|
||||
req, err = http.NewRequest(c.Request.Method, buildurl, bytes.NewReader(chatReq.ToByteJson()))
|
||||
req.Header = c.Request.Header
|
||||
req.Header.Set("api-key", onekey.Key)
|
||||
default:
|
||||
req, err = http.NewRequest(c.Request.Method, OpenAI_Endpoint, bytes.NewReader(chatReq.ToByteJson())) // default endpoint
|
||||
|
||||
if AIGateWay_Endpoint != "" { // cloudflare gateway的endpoint
|
||||
req, err = http.NewRequest(c.Request.Method, AIGateWay_Endpoint, bytes.NewReader(chatReq.ToByteJson()))
|
||||
}
|
||||
if Custom_Endpoint != "" { // 自定义endpoint
|
||||
req, err = http.NewRequest(c.Request.Method, Custom_Endpoint, bytes.NewReader(chatReq.ToByteJson()))
|
||||
}
|
||||
if onekey.EndPoint != "" { // 优先key的endpoint
|
||||
req, err = http.NewRequest(c.Request.Method, onekey.EndPoint+c.Request.RequestURI, bytes.NewReader(chatReq.ToByteJson()))
|
||||
}
|
||||
|
||||
req.Header = c.Request.Header
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", onekey.Key))
|
||||
}
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var result string
|
||||
if chatReq.Stream {
|
||||
for key, value := range resp.Header {
|
||||
for _, v := range value {
|
||||
c.Writer.Header().Add(key, v)
|
||||
}
|
||||
}
|
||||
c.Writer.WriteHeader(resp.StatusCode)
|
||||
teeReader := io.TeeReader(resp.Body, c.Writer)
|
||||
// 流式响应
|
||||
scanner := bufio.NewScanner(teeReader)
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Bytes()
|
||||
if len(line) > 0 && bytes.HasPrefix(line, []byte("data: ")) {
|
||||
if bytes.HasPrefix(line, []byte("data: [DONE]")) {
|
||||
break
|
||||
}
|
||||
var opiResp ChatCompletionStreamResponse
|
||||
line = bytes.Replace(line, []byte("data: "), []byte(""), -1)
|
||||
line = bytes.TrimSpace(line)
|
||||
if err := json.Unmarshal(line, &opiResp); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if opiResp.Choices != nil && len(opiResp.Choices) > 0 {
|
||||
if opiResp.Choices[0].Delta.Role != "" {
|
||||
result += "<" + opiResp.Choices[0].Delta.Role + "> "
|
||||
}
|
||||
result += opiResp.Choices[0].Delta.Content // 计算Content Token
|
||||
|
||||
if len(opiResp.Choices[0].Delta.ToolCalls) > 0 { // 计算ToolCalls token
|
||||
if opiResp.Choices[0].Delta.ToolCalls[0].Function.Name != "" {
|
||||
result += "name:" + opiResp.Choices[0].Delta.ToolCalls[0].Function.Name + " arguments:"
|
||||
}
|
||||
result += opiResp.Choices[0].Delta.ToolCalls[0].Function.Arguments
|
||||
}
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
} else {
|
||||
|
||||
// 处理非流式响应
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
fmt.Println("Error reading response body:", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
var opiResp ChatCompletionResponse
|
||||
if err := json.Unmarshal(body, &opiResp); err != nil {
|
||||
log.Println("Error parsing JSON:", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if opiResp.Choices != nil && len(opiResp.Choices) > 0 {
|
||||
if opiResp.Choices[0].Message.Role != "" {
|
||||
result += "<" + opiResp.Choices[0].Message.Role + "> "
|
||||
}
|
||||
result += opiResp.Choices[0].Message.Content
|
||||
|
||||
if len(opiResp.Choices[0].Message.ToolCalls) > 0 {
|
||||
if opiResp.Choices[0].Message.ToolCalls[0].Function.Name != "" {
|
||||
result += "name:" + opiResp.Choices[0].Message.ToolCalls[0].Function.Name + " arguments:"
|
||||
}
|
||||
result += opiResp.Choices[0].Message.ToolCalls[0].Function.Arguments
|
||||
}
|
||||
|
||||
}
|
||||
for k, v := range resp.Header {
|
||||
c.Writer.Header().Set(k, v[0])
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, opiResp)
|
||||
}
|
||||
usagelog.CompletionCount = tokenizer.NumTokensFromStr(result, chatReq.Model)
|
||||
usagelog.Cost = fmt.Sprintf("%.6f", tokenizer.Cost(usagelog.Model, usagelog.PromptCount, usagelog.CompletionCount))
|
||||
if err := store.Record(&usagelog); err != nil {
|
||||
log.Println(err)
|
||||
}
|
||||
if err := store.SumDaily(usagelog.UserID); err != nil {
|
||||
log.Println(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user