This commit is contained in:
c菌
2023-09-15 23:09:43 +08:00
parent 7df0b2817c
commit de31741d06
2 changed files with 111 additions and 42 deletions

View File

@@ -44,6 +44,7 @@ import (
"net/url" "net/url"
"opencatd-open/store" "opencatd-open/store"
"strings" "strings"
"sync"
"time" "time"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
@@ -255,11 +256,11 @@ func TransReq(chatreq *openai.ChatCompletionRequest) (*bytes.Buffer, error) {
return payload, nil return payload, nil
} }
func TransRsp(c *gin.Context, isStream bool, responseBody *bufio.Reader) { func TransRsp(c *gin.Context, isStream bool, reader *bufio.Reader) {
if !isStream { if !isStream {
var completersp CompleteResponse var completersp CompleteResponse
var chatrsp openai.ChatCompletionResponse var chatrsp openai.ChatCompletionResponse
json.NewDecoder(responseBody).Decode(&completersp) json.NewDecoder(reader).Decode(&completersp)
chatrsp.Model = completersp.Model chatrsp.Model = completersp.Model
chatrsp.ID = completersp.LogID chatrsp.ID = completersp.LogID
chatrsp.Object = "chat.completion" chatrsp.Object = "chat.completion"
@@ -285,57 +286,115 @@ func TransRsp(c *gin.Context, isStream bool, responseBody *bufio.Reader) {
c.JSON(http.StatusOK, payload) c.JSON(http.StatusOK, payload)
return return
} else { } else {
result := TranslateStream(c, responseBody) var (
for _, content := range <-result { wg sync.WaitGroup
c.Writer.WriteString(string(content)) dataChan = make(chan string)
c.Writer.Flush() stopChan = make(chan bool)
} )
go func() {
defer wg.Done()
for {
line, err := reader.ReadString('\n')
if err == nil {
if strings.HasPrefix(line, "data: ") {
var result CompleteResponse
json.NewDecoder(strings.NewReader(line[6:])).Decode(&result)
if result.StopReason == "" {
if result.Completion != "" {
chatrsp := openai.ChatCompletionStreamResponse{
ID: result.LogID,
Model: result.Model,
Object: "chat.completion",
Created: time.Now().Unix(),
}
choice := openai.ChatCompletionStreamChoice{
Delta: openai.ChatCompletionStreamChoiceDelta{
Role: "assistant",
Content: result.Completion,
},
FinishReason: "",
}
chatrsp.Choices = append(chatrsp.Choices, choice)
bytedate, _ := json.Marshal(chatrsp)
dataChan <- string(bytedate)
}
} else {
dataChan <- "[DONE]"
break
}
} else {
continue
}
} else {
break
}
}
close(dataChan)
stopChan <- true
close(stopChan)
}()
go func() {
defer wg.Done()
Loop:
for {
select {
case data := <-dataChan:
if data != "" {
c.Writer.WriteString("data: " + data)
c.Writer.Flush()
}
case <-stopChan:
break Loop
}
}
}()
wg.Wait()
} }
} }
func TranslateStream(ctx *gin.Context, reader *bufio.Reader,dataChan chan string,stopChan chan bool) { func TranslateStream(ctx *gin.Context, reader *bufio.Reader, dataChan chan string, stopChan chan bool) {
// dataChan := make(chan string) // dataChan := make(chan string)
// stopChan := make(chan bool) // stopChan := make(chan bool)
go func () { for {
for { line, err := reader.ReadString('\n')
line, err := reader.ReadString('\n') if err == nil {
if err == nil { if strings.HasPrefix(line, "data: ") {
if strings.HasPrefix(line, "data: ") { var result CompleteResponse
json.NewDecoder(strings.NewReader(line[6:])).Decode(&result)
var result CompleteResponse if result.StopReason == "" {
json.NewDecoder(strings.NewReader(line[6:])).Decode(&result) if result.Completion != "" {
if result.StopReason == "" { chatrsp := openai.ChatCompletionStreamResponse{
if result.Completion != "" { ID: result.LogID,
chatrsp := openai.ChatCompletionStreamResponse{ Model: result.Model,
ID: result.LogID, Object: "chat.completion",
Model: result.Model, Created: time.Now().Unix(),
Object: "chat.completion",
Created: time.Now().Unix(),
}
choice := openai.ChatCompletionStreamChoice{
Delta: openai.ChatCompletionStreamChoiceDelta{
Role: "assistant",
Content: result.Completion,
},
FinishReason: "",
}
chatrsp.Choices = append(chatrsp.Choices, choice)
bytedate, _ := json.Marshal(chatrsp)
dataChan <- string(bytedate)
} }
} else { choice := openai.ChatCompletionStreamChoice{
log.Println("finish:", result.StopReason) Delta: openai.ChatCompletionStreamChoiceDelta{
dataChan <- "data: [DONE]" Role: "assistant",
break Content: result.Completion,
},
FinishReason: "",
}
chatrsp.Choices = append(chatrsp.Choices, choice)
bytedate, _ := json.Marshal(chatrsp)
dataChan <- string(bytedate)
} }
} else { } else {
continue dataChan <- "data: [DONE]"
break
} }
} else { } else {
log.Println("err:", err) continue
break
} }
}() } else {
log.Println("err:", err)
break
}
}
} }
func Translate(c *gin.Context, chatreq *openai.ChatCompletionRequest) { func Translate(c *gin.Context, chatreq *openai.ChatCompletionRequest) {

View File

@@ -511,6 +511,16 @@ func HandleProy(c *gin.Context) {
onekey := store.FromKeyCacheRandomItemKey() onekey := store.FromKeyCacheRandomItemKey()
// 创建 API 请求 // 创建 API 请求
switch onekey.ApiType { switch onekey.ApiType {
case "claude":
payload, _ := claude.TransReq(&chatreq)
buildurl := "https://api.anthropic.com/v1/complete"
req, err = http.NewRequest("POST", buildurl, payload)
req.Header.Add("accept", "application/json")
req.Header.Add("anthropic-version", "2023-06-01")
req.Header.Add("x-api-key", onekey.Key)
req.Header.Add("content-type", "application/json")
case "azure":
fallthrough
case "azure_openai": case "azure_openai":
var buildurl string var buildurl string
var apiVersion = "2023-05-15" var apiVersion = "2023-05-15"