From de31741d06ee8e4c165e335b96af1171baa8271a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?c=E8=8F=8C?= <littlecjun@gmail.com>
Date: Fri, 15 Sep 2023 23:09:43 +0800
Subject: [PATCH] update

---
 pkg/claude/claude.go | 143 ++++++++++++++++++++++++++++++-------------
 router/router.go     |  10 +++
 2 files changed, 111 insertions(+), 42 deletions(-)

diff --git a/pkg/claude/claude.go b/pkg/claude/claude.go
index 53cb6b7..5218101 100644
--- a/pkg/claude/claude.go
+++ b/pkg/claude/claude.go
@@ -44,6 +44,7 @@ import (
 	"net/url"
 	"opencatd-open/store"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -255,11 +256,11 @@ func TransReq(chatreq *openai.ChatCompletionRequest) (*bytes.Buffer, error) {
 	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 {
 		var completersp CompleteResponse
 		var chatrsp openai.ChatCompletionResponse
-		json.NewDecoder(responseBody).Decode(&completersp)
+		json.NewDecoder(reader).Decode(&completersp)
 		chatrsp.Model = completersp.Model
 		chatrsp.ID = completersp.LogID
 		chatrsp.Object = "chat.completion"
@@ -285,57 +286,115 @@ func TransRsp(c *gin.Context, isStream bool, responseBody *bufio.Reader) {
 		c.JSON(http.StatusOK, payload)
 		return
 	} else {
-		result := TranslateStream(c, responseBody)
-		for _, content := range <-result {
-			c.Writer.WriteString(string(content))
-			c.Writer.Flush()
-		}
+		var (
+			wg       sync.WaitGroup
+			dataChan = make(chan string)
+			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)
 	// stopChan := make(chan bool)
-	go func ()  {
-		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)
+	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(),
 						}
-					} else {
-						log.Println("finish:", result.StopReason)
-						dataChan <- "data: [DONE]"
-						break
+						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 {
-					continue
+					dataChan <- "data: [DONE]"
+					break
 				}
 			} else {
-				log.Println("err:", err)
-				break
+				continue
 			}
-	}()
+		} else {
+			log.Println("err:", err)
+			break
+		}
+	}
+
 }
 
 func Translate(c *gin.Context, chatreq *openai.ChatCompletionRequest) {
diff --git a/router/router.go b/router/router.go
index d0db20d..f261ece 100644
--- a/router/router.go
+++ b/router/router.go
@@ -511,6 +511,16 @@ func HandleProy(c *gin.Context) {
 		onekey := store.FromKeyCacheRandomItemKey()
 		// 创建 API 请求
 		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":
 			var buildurl string
 			var apiVersion = "2023-05-15"