Przeglądaj źródła

feat: support ChatGLM2 (close #274)

JustSong 2 lat temu
rodzic
commit
0e088f7c3e
3 zmienionych plików z 25 dodań i 5 usunięć
  1. 2 1
      README.md
  2. 18 0
      controller/model.go
  3. 5 4
      controller/relay-text.go

+ 2 - 1
README.md

@@ -91,7 +91,8 @@ _✨ All in one 的 OpenAI 接口,整合各种 API 访问方式,开箱即用
     + 邮箱登录注册以及通过邮箱进行密码重置。
     + [GitHub 开放授权](https://github.com/settings/applications/new)。
     + 微信公众号授权(需要额外部署 [WeChat Server](https://github.com/songquanpeng/wechat-server))。
-20. 未来其他大模型开放 API 后,将第一时间支持,并将其封装成同样的 API 访问方式。
+20. 支持 [ChatGLM](https://github.com/THUDM/ChatGLM2-6B)。
+21. 未来其他大模型开放 API 后,将第一时间支持,并将其封装成同样的 API 访问方式。
 
 ## 部署
 ### 基于 Docker 进行部署

+ 18 - 0
controller/model.go

@@ -252,6 +252,24 @@ func init() {
 			Root:       "code-davinci-edit-001",
 			Parent:     nil,
 		},
+		{
+			Id:         "ChatGLM",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "thudm",
+			Permission: permission,
+			Root:       "ChatGLM",
+			Parent:     nil,
+		},
+		{
+			Id:         "ChatGLM2",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "thudm",
+			Permission: permission,
+			Root:       "ChatGLM2",
+			Parent:     nil,
+		},
 	}
 	openAIModelsMap = make(map[string]OpenAIModels)
 	for _, model := range openAIModels {

+ 5 - 4
controller/relay-text.go

@@ -227,8 +227,8 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 				return 0, nil, nil
 			}
 
-			if i := strings.Index(string(data), "\n\n"); i >= 0 {
-				return i + 2, data[0:i], nil
+			if i := strings.Index(string(data), "\n"); i >= 0 {
+				return i + 1, data[0:i], nil
 			}
 
 			if atEOF {
@@ -242,8 +242,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 		go func() {
 			for scanner.Scan() {
 				data := scanner.Text()
-				if len(data) < 6 { // must be something wrong!
-					common.SysError("invalid stream response: " + data)
+				if len(data) < 6 { // ignore blank line or wrong format
 					continue
 				}
 				dataChan <- data
@@ -286,6 +285,8 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 				if strings.HasPrefix(data, "data: [DONE]") {
 					data = data[:12]
 				}
+				// some implementations may add \r at the end of data
+				data = strings.TrimSuffix(data, "\r")
 				c.Render(-1, common.CustomEvent{Data: data})
 				return true
 			case <-stopChan: