支持上下文对话方式追加问题

Dockerfile

@@ -2,11 +2,11 @@ FROM alpine:3.15.3
 
 LABEL maintainer="cookeem"
 LABEL email="cookeem@qq.com"
-LABEL version="v1.0.2"
+LABEL version="v1.0.3"
 
 RUN adduser -h /chatgpt-service -u 1000 -D dory
 COPY chatgpt-service /chatgpt-service/
 WORKDIR /chatgpt-service
 USER dory
 
-# docker build -t doryengine/chatgpt-service:v1.0.2-alpine .
+# docker build -t doryengine/chatgpt-service:v1.0.3-alpine .

README.md

@@ -1,4 +1,4 @@
-# Real-time ChatGPT service, based on the latest gpt-3.5-turbo-0301 model
+# Real-time ChatGPT service, support GPT3/GPT4, support conversation and generate pictures from sentences
 
 - [English README](README.md)
 - [中文 README](README_CN.md)
@@ -17,8 +17,13 @@
 
 ## Demo
 
+- Real-time conversation mode
+
 ![](chatgpt-service.gif)
 
+- Generate picture patterns from sentences
+
+![](chatgpt-image.jpeg)
 
 ## Quick start
 
@@ -36,6 +41,9 @@ vi config.yaml
 #  your openai.com API key
 apiKey: "xxxxxx"
 
+# create pictures directory
+mkdir -p assets
+chown -R 1000:1000 assets
 
 # Start the service with docker-compose
 docker-compose up -d
@@ -52,6 +60,9 @@ chatgpt-stream    /docker-entrypoint.sh ngin ...   Up      0.0.0.0:3000->80/tcp,
 # http://localhost:3000
 ```
 
+- Enter the question directly, it will call the ChatGPT interface to return the answer
+- Enter the picture description after `/image`, it will call the DALL-E2 interface to automatically generate pictures through the picture description
+
 ## How to build
 
 ```bash

README_CN.md

@@ -1,4 +1,4 @@
-# 实时ChatGPT服务，基于最新的gpt-3.5-turbo-0301模型
+# 实时ChatGPT服务，支持GPT3/GPT4，支持对话和通过句子生成图片
 
 - [English README](README.md)
 - [中文 README](README_CN.md)
@@ -17,8 +17,13 @@
 
 ## 效果图
 
+- 实时对话模式
+
 ![](chatgpt-service.gif)
 
+- 通过句子生成图片模式
+
+![](chatgpt-image.jpeg)
 
 ## 快速开始
 
@@ -36,6 +41,9 @@ vi config.yaml
 # openai的apiKey，改为你的apiKey
 apiKey: "xxxxxx"
 
+# 创建生成的图片目录
+mkdir -p assets
+chown -R 1000:1000 assets
 
 # 使用docker-compose启动服务
 docker-compose up -d
@@ -52,6 +60,9 @@ chatgpt-stream    /docker-entrypoint.sh ngin ...   Up      0.0.0.0:3000->80/tcp,
 # http://localhost:3000
 ```
 
+- 直接输入问题，则调用ChatGPT接口返回答案
+- `/image `后边输入想要的图片描述，则调用DALL-E2接口，通过图片描述自动生成图片
+
 ## 如何编译
 
 ```bash

chat/common.go

@@ -1,8 +1,10 @@
 package chat
 
 import (
+	"fmt"
 	"github.com/sashabaranov/go-openai"
 	log "github.com/sirupsen/logrus"
+	"math/rand"
 	"os"
 	"time"
 )
@@ -39,6 +41,22 @@ func (logger Logger) LogPanic(args ...interface{}) {
 	log.Panic(args...)
 }
 
+func RandomString(n int) string {
+	var letter []rune
+	lowerChars := "abcdefghijklmnopqrstuvwxyz"
+	numberChars := "0123456789"
+	chars := fmt.Sprintf("%s%s", lowerChars, numberChars)
+	letter = []rune(chars)
+	var str string
+	b := make([]rune, n)
+	seededRand := rand.New(rand.NewSource(time.Now().UnixNano()))
+	for i := range b {
+		b[i] = letter[seededRand.Intn(len(letter))]
+	}
+	str = string(b)
+	return str
+}
+
 const (
 	StatusFail string = "FAIL"
 

chat/service.go

@@ -2,10 +2,12 @@ package chat
 
 import (
 	"context"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"strings"
 	"sync"
 	"time"
@@ -74,7 +76,7 @@ func (api *Api) wsPingMsg(conn *websocket.Conn, chClose, chIsCloseSet chan int)
 	}
 }
 
-func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *sync.Mutex, requestMsg string) {
+func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *sync.Mutex, reqMsgs []openai.ChatCompletionMessage) {
 	var err error
 	var strResp string
 
@@ -82,16 +84,12 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 
 	switch api.Config.Model {
 	case openai.GPT3Dot5Turbo0301, openai.GPT3Dot5Turbo, openai.GPT4, openai.GPT40314, openai.GPT432K0314, openai.GPT432K:
+		prompt := reqMsgs[len(reqMsgs)-1].Content
 		req := openai.ChatCompletionRequest{
-			Model:       api.Config.Model,
-			MaxTokens:   api.Config.MaxLength,
-			Temperature: 1.0,
-			Messages: []openai.ChatCompletionMessage{
-				{
-					Role:    openai.ChatMessageRoleUser,
-					Content: requestMsg,
-				},
-			},
+			Model:            api.Config.Model,
+			MaxTokens:        api.Config.MaxLength,
+			Temperature:      1.0,
+			Messages:         reqMsgs,
 			Stream:           true,
 			TopP:             1,
 			FrequencyPenalty: 0.1,
@@ -100,7 +98,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 
 		stream, err := cli.CreateChatCompletionStream(ctx, req)
 		if err != nil {
-			err = fmt.Errorf("[ERROR] create chatGPT stream model=%s error: %s", api.Config.Model, err.Error())
+			err = fmt.Errorf("[ERROR] create ChatGPT stream model=%s error: %s", api.Config.Model, err.Error())
 			chatMsg := Message{
 				Kind:       "error",
 				Msg:        err.Error(),
@@ -149,7 +147,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 			if len(response.Choices) > 0 {
 				var s string
 				if i == 0 {
-					s = fmt.Sprintf(`%s# %s`, s, requestMsg)
+					s = fmt.Sprintf("%s# %s\n\n", s, prompt)
 				}
 				for _, choice := range response.Choices {
 					s = s + choice.Delta.Content
@@ -171,11 +169,12 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 			api.Logger.LogInfo(fmt.Sprintf("[RESPONSE] %s\n", strResp))
 		}
 	case openai.GPT3TextDavinci003, openai.GPT3TextDavinci002, openai.GPT3TextCurie001, openai.GPT3TextBabbage001, openai.GPT3TextAda001, openai.GPT3TextDavinci001, openai.GPT3DavinciInstructBeta, openai.GPT3Davinci, openai.GPT3CurieInstructBeta, openai.GPT3Curie, openai.GPT3Ada, openai.GPT3Babbage:
+		prompt := reqMsgs[len(reqMsgs)-1].Content
 		req := openai.CompletionRequest{
 			Model:       api.Config.Model,
 			MaxTokens:   api.Config.MaxLength,
 			Temperature: 0.6,
-			Prompt:      requestMsg,
+			Prompt:      prompt,
 			Stream:      true,
 			//Stop:             []string{"\n\n\n"},
 			TopP:             1,
@@ -185,7 +184,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 
 		stream, err := cli.CreateCompletionStream(ctx, req)
 		if err != nil {
-			err = fmt.Errorf("[ERROR] create chatGPT stream model=%s error: %s", api.Config.Model, err.Error())
+			err = fmt.Errorf("[ERROR] create ChatGPT stream model=%s error: %s", api.Config.Model, err.Error())
 			chatMsg := Message{
 				Kind:       "error",
 				Msg:        err.Error(),
@@ -234,7 +233,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 			if len(response.Choices) > 0 {
 				var s string
 				if i == 0 {
-					s = fmt.Sprintf(`%s# %s`, s, requestMsg)
+					s = fmt.Sprintf("%s# %s\n\n", s, prompt)
 				}
 				for _, choice := range response.Choices {
 					s = s + choice.Text
@@ -262,6 +261,83 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *openai.Client, mutex *
 	}
 }
 
+func (api *Api) GetImageMessage(conn *websocket.Conn, cli *openai.Client, mutex *sync.Mutex, requestMsg string) {
+	var err error
+
+	ctx := context.Background()
+
+	prompt := strings.TrimPrefix(requestMsg, "/image ")
+	req := openai.ImageRequest{
+		Prompt:         prompt,
+		Size:           openai.CreateImageSize256x256,
+		ResponseFormat: openai.CreateImageResponseFormatB64JSON,
+		N:              1,
+	}
+
+	sendError := func(err error) {
+		err = fmt.Errorf("[ERROR] generate image error: %s", err.Error())
+		chatMsg := Message{
+			Kind:       "error",
+			Msg:        err.Error(),
+			MsgId:      uuid.New().String(),
+			CreateTime: time.Now().Format("2006-01-02 15:04:05"),
+		}
+		mutex.Lock()
+		_ = conn.WriteJSON(chatMsg)
+		mutex.Unlock()
+		api.Logger.LogError(err.Error())
+	}
+
+	resp, err := cli.CreateImage(ctx, req)
+	if err != nil {
+		err = fmt.Errorf("[ERROR] generate image error: %s", err.Error())
+		sendError(err)
+		return
+	}
+	if len(resp.Data) == 0 {
+		err = fmt.Errorf("[ERROR] generate image error: result is empty")
+		sendError(err)
+		return
+	}
+
+	imgBytes, err := base64.StdEncoding.DecodeString(resp.Data[0].B64JSON)
+	if err != nil {
+		err = fmt.Errorf("[ERROR] image base64 decode error: %s", err.Error())
+		sendError(err)
+		return
+	}
+
+	date := time.Now().Format("2006-01-02")
+	imageDir := fmt.Sprintf("assets/images/%s", date)
+	err = os.MkdirAll(imageDir, 0700)
+	if err != nil {
+		err = fmt.Errorf("[ERROR] create image directory error: %s", err.Error())
+		sendError(err)
+		return
+	}
+
+	imageFileName := fmt.Sprintf("%s.png", RandomString(16))
+	err = os.WriteFile(fmt.Sprintf("%s/%s", imageDir, imageFileName), imgBytes, 0600)
+	if err != nil {
+		err = fmt.Errorf("[ERROR] write png image error: %s", err.Error())
+		sendError(err)
+		return
+	}
+
+	msg := fmt.Sprintf("api/%s/%s", imageDir, imageFileName)
+	chatMsg := Message{
+		Kind:       "image",
+		Msg:        msg,
+		MsgId:      uuid.New().String(),
+		CreateTime: time.Now().Format("2006-01-02 15:04:05"),
+	}
+	mutex.Lock()
+	_ = conn.WriteJSON(chatMsg)
+	mutex.Unlock()
+	api.Logger.LogInfo(fmt.Sprintf("[IMAGE] # %s\n%s", requestMsg, msg))
+	return
+}
+
 func (api *Api) WsChat(c *gin.Context) {
 	startTime := time.Now()
 	status := StatusFail
@@ -314,6 +390,8 @@ func (api *Api) WsChat(c *gin.Context) {
 	api.Logger.LogInfo(fmt.Sprintf("websocket connection open"))
 	cli := openai.NewClient(api.Config.ApiKey)
 
+	reqMsgs := make([]openai.ChatCompletionMessage, 0)
+
 	var latestRequestTime time.Time
 	for {
 		if isClosed {
@@ -366,16 +444,33 @@ func (api *Api) WsChat(c *gin.Context) {
 					mutex.Unlock()
 					api.Logger.LogError(err.Error())
 				} else {
-					chatMsg := Message{
-						Kind:       "receive",
-						Msg:        requestMsg,
-						MsgId:      uuid.New().String(),
-						CreateTime: time.Now().Format("2006-01-02 15:04:05"),
+					if strings.HasPrefix(requestMsg, "/image ") {
+						chatMsg := Message{
+							Kind:       "receive",
+							Msg:        requestMsg,
+							MsgId:      uuid.New().String(),
+							CreateTime: time.Now().Format("2006-01-02 15:04:05"),
+						}
+						mutex.Lock()
+						_ = conn.WriteJSON(chatMsg)
+						mutex.Unlock()
+						go api.GetImageMessage(conn, cli, mutex, requestMsg)
+					} else {
+						chatMsg := Message{
+							Kind:       "receive",
+							Msg:        requestMsg,
+							MsgId:      uuid.New().String(),
+							CreateTime: time.Now().Format("2006-01-02 15:04:05"),
+						}
+						mutex.Lock()
+						_ = conn.WriteJSON(chatMsg)
+						mutex.Unlock()
+						reqMsgs = append(reqMsgs, openai.ChatCompletionMessage{
+							Role:    openai.ChatMessageRoleUser,
+							Content: requestMsg,
+						})
+						go api.GetChatMessage(conn, cli, mutex, reqMsgs)
 					}
-					mutex.Lock()
-					_ = conn.WriteJSON(chatMsg)
-					mutex.Unlock()
-					go api.GetChatMessage(conn, cli, mutex, requestMsg)
 				}
 			}
 		case websocket.CloseMessage:

docker-compose.yaml

@@ -1,7 +1,7 @@
 version: "3"
 services:
   chatgpt-stream:
-    image: "doryengine/chatgpt-stream:v1.0.2"
+    image: "doryengine/chatgpt-stream:v1.0.3"
     hostname: chatgpt-stream
     container_name: chatgpt-stream
     ports:
@@ -11,12 +11,13 @@ services:
       - chatgpt-service
     restart: always
   chatgpt-service:
-    image: "doryengine/chatgpt-service:v1.0.2-alpine"
+    image: "doryengine/chatgpt-service:v1.0.3-alpine"
     hostname: chatgpt-service
     container_name: chatgpt-service
     ports:
       - "9000:9000"
     volumes:
       - ./config.yaml:/chatgpt-service/config.yaml
+      - ./assets:/chatgpt-service/assets
     command: /chatgpt-service/chatgpt-service
     restart: always

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/gin-gonic/gin v1.8.2
 	github.com/google/uuid v1.3.0
 	github.com/gorilla/websocket v1.5.0
-	github.com/sashabaranov/go-openai v1.5.4
+	github.com/sashabaranov/go-openai v1.5.7
 	github.com/sirupsen/logrus v1.9.0
 	gopkg.in/yaml.v3 v3.0.1
 )

main.go

@@ -56,6 +56,7 @@ func main() {
 	}
 
 	groupApi := r.Group("/api")
+	groupApi.Static("/assets", "assets")
 	groupWs := groupApi.Group("/ws")
 	groupWs.GET("chat", api.WsChat)