diff --git a/addition/web/call.go b/addition/web/call.go index 3f6dc08..fdd7d60 100644 --- a/addition/web/call.go +++ b/addition/web/call.go @@ -26,8 +26,8 @@ func ChatWithWeb(hook Hook, message []globals.Message, long bool) (string, []glo return keyword, utils.Insert(message, 0, globals.Message{ Role: "system", Content: fmt.Sprintf("你将扮演AI问答助手,你的知识库不是截止于2021年9月,而是可以实时联网的,你可以提供实时联网的知识库。"+ - "网站使用markdown链接包裹,多使用emoji和颜文字,给予用户精确的答复。"+ - "公式使用 $$ [公式] $$ 包裹,使用latex语法,例如:$$ \\frac{1}{2} $$。"+ + //"网站使用markdown链接包裹,多使用emoji和颜文字,给予用户精确的答复。"+ + //"公式使用 $$ [公式] $$ 包裹,使用latex语法,例如:$$ \\frac{1}{2} $$。"+ "当前时间: %s, 实时联网搜索结果:%s", time.Now().Format("2006-01-02 15:04:05"), data, ), diff --git a/go.mod b/go.mod index 25f3b78..d201c0e 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module chat go 1.20 require ( + github.com/alexandrevicenzi/go-sse v1.6.0 github.com/bincooo/claude-api v1.0.2 github.com/dgrijalva/jwt-go v3.2.0+incompatible github.com/gin-gonic/gin v1.9.1 @@ -10,8 +11,10 @@ require ( github.com/go-sql-driver/mysql v1.7.1 github.com/google/uuid v1.3.1 github.com/gorilla/websocket v1.5.0 + github.com/natefinch/lumberjack v2.0.0+incompatible github.com/pkoukk/tiktoken-go v0.1.6 github.com/russross/blackfriday/v2 v2.1.0 + github.com/sirupsen/logrus v1.9.3 github.com/spf13/viper v1.16.0 golang.org/x/net v0.15.0 ) @@ -45,10 +48,8 @@ require ( github.com/mitchellh/mapstructure v1.5.0 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect - github.com/natefinch/lumberjack v2.0.0+incompatible // indirect github.com/pelletier/go-toml/v2 v2.1.0 // indirect github.com/refraction-networking/utls v1.3.2 // indirect - github.com/sirupsen/logrus v1.9.3 // indirect github.com/spf13/afero v1.10.0 // indirect github.com/spf13/cast v1.5.1 // indirect github.com/spf13/jwalterweatherman v1.1.0 // indirect @@ -62,6 +63,8 @@ require ( golang.org/x/sys v0.12.0 // indirect golang.org/x/text v0.13.0 // indirect google.golang.org/protobuf v1.31.0 // indirect + gopkg.in/antage/eventsource.v1 v1.0.0-20150318155416-803f4c5af225 // indirect gopkg.in/ini.v1 v1.67.0 // indirect + gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 21d94b7..94180e9 100644 --- a/go.sum +++ b/go.sum @@ -36,8 +36,11 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= +github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= +github.com/alexandrevicenzi/go-sse v1.6.0 h1:3KvOzpuY7UrbqZgAtOEmub9/V5ykr7Myudw+PA+H1Ik= +github.com/alexandrevicenzi/go-sse v1.6.0/go.mod h1:jdrNAhMgVqP7OfcUuM8eJx0sOY17wc+girs5utpFZUU= github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs= github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/bincooo/claude-api v1.0.2 h1:qhN+s5vpMRU7e7IN+K+uyCmUDV/jtnihgfST8IC0G1o= @@ -568,12 +571,16 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8= google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/antage/eventsource.v1 v1.0.0-20150318155416-803f4c5af225 h1:xy+AV3uSExoRQc2qWXeZdbhFGwBFK/AmGlrBZEjbvuQ= +gopkg.in/antage/eventsource.v1 v1.0.0-20150318155416-803f4c5af225/go.mod h1:SiXNRpUllqhl+GIw2V/BtKI7BUlz+uxov9vBFtXHqh8= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= gopkg.in/ini.v1 v1.67.0 h1:Dgnx+6+nfE+IfzjUEISNeydPJh9AXNNsWbGP9KzCsOA= gopkg.in/ini.v1 v1.67.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= +gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= +gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= diff --git a/utils/sse.go b/utils/sse.go new file mode 100644 index 0000000..29e7314 --- /dev/null +++ b/utils/sse.go @@ -0,0 +1,122 @@ +package utils + +import ( + "bufio" + "bytes" + "chat/globals" + "crypto/tls" + "fmt" + "net/http" +) + +func SSEClient(method string, uri string, headers map[string]string, body interface{}, callback func(string) error) error { + http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true} + + client := &http.Client{} + req, err := http.NewRequest(method, uri, ConvertBody(body)) + if err != nil { + return nil + } + for key, value := range headers { + req.Header.Set(key, value) + } + + res, err := client.Do(req) + if err != nil { + return err + } + if res.StatusCode >= 400 { + return fmt.Errorf("request failed with status: %s", res.Status) + } + + defer res.Body.Close() + + events, err := CreateSSEInstance(res) + if err != nil { + return err + } + + select { + case ev := <-events: + if err := callback(ev.Data); err != nil { + return err + } + } + + return nil +} + +// Event represents a Server-Sent Event +type Event struct { + Name string + ID string + Data string +} + +func CreateSSEInstance(resp *http.Response) (chan Event, error) { + events := make(chan Event) + reader := bufio.NewReader(resp.Body) + + go loop(reader, events) + + return events, nil +} + +func loop(reader *bufio.Reader, events chan Event) { + ev := Event{} + + var buf bytes.Buffer + + for { + line, err := reader.ReadBytes('\n') + if err != nil { + globals.Info(fmt.Sprintf("[sse] error during read response body: %s", err)) + close(events) + } + + switch { + case ioPrefix(line, ":"): + // Comment, do nothing + case ioPrefix(line, "retry:"): + // Retry, do nothing for now + + // id of event + case ioPrefix(line, "id: "): + ev.ID = string(line[4:]) + case ioPrefix(line, "id:"): + ev.ID = string(line[3:]) + + // name of event + case ioPrefix(line, "event: "): + ev.Name = string(line[7 : len(line)-1]) + case ioPrefix(line, "event:"): + ev.Name = string(line[6 : len(line)-1]) + + // event data + case ioPrefix(line, "data: "): + buf.Write(line[6:]) + case ioPrefix(line, "data:"): + buf.Write(line[5:]) + + // end of event + case bytes.Equal(line, []byte("\n")): + b := buf.Bytes() + + if ioPrefix(b, "{") { + if err == nil { + ev.Data = string(b) + buf.Reset() + events <- ev + ev = Event{} + } + } + + default: + globals.Info(fmt.Sprintf("[sse] unknown line: %s", line)) + } + } +} + +func ioPrefix(s []byte, prefix string) bool { + return bytes.HasPrefix(s, []byte(prefix)) +}