添加 /ai 接口
This commit is contained in:
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1 +1,2 @@
|
||||
*.db
|
||||
*.db
|
||||
.env
|
||||
@@ -6,7 +6,7 @@ COPY go.mod go.sum ./
|
||||
RUN go env -w GO111MODULE=on
|
||||
RUN go env -w GOPROXY=https://goproxy.cn,direct
|
||||
RUN go mod download
|
||||
COPY main.go ./
|
||||
COPY main.go .env ./
|
||||
|
||||
RUN go build -o api
|
||||
|
||||
|
||||
16
go.mod
16
go.mod
@@ -6,6 +6,7 @@ require (
|
||||
github.com/gin-contrib/cors v1.7.2
|
||||
github.com/gin-gonic/gin v1.10.0
|
||||
github.com/glebarez/sqlite v1.11.0
|
||||
github.com/openai/openai-go/v2 v2.7.0
|
||||
gorm.io/gorm v1.25.12
|
||||
)
|
||||
|
||||
@@ -23,9 +24,10 @@ require (
|
||||
github.com/go-playground/validator/v10 v10.20.0 // indirect
|
||||
github.com/goccy/go-json v0.10.2 // indirect
|
||||
github.com/google/go-cmp v0.5.9 // indirect
|
||||
github.com/google/uuid v1.4.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/joho/godotenv v1.5.1 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.2.7 // indirect
|
||||
github.com/kr/pretty v0.3.1 // indirect
|
||||
@@ -35,13 +37,17 @@ require (
|
||||
github.com/modern-go/reflect2 v1.0.2 // indirect
|
||||
github.com/pelletier/go-toml/v2 v2.2.2 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/tidwall/gjson v1.14.4 // indirect
|
||||
github.com/tidwall/match v1.1.1 // indirect
|
||||
github.com/tidwall/pretty v1.2.1 // indirect
|
||||
github.com/tidwall/sjson v1.2.5 // indirect
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||
github.com/ugorji/go/codec v1.2.12 // indirect
|
||||
golang.org/x/arch v0.8.0 // indirect
|
||||
golang.org/x/crypto v0.23.0 // indirect
|
||||
golang.org/x/net v0.25.0 // indirect
|
||||
golang.org/x/sys v0.20.0 // indirect
|
||||
golang.org/x/text v0.15.0 // indirect
|
||||
golang.org/x/crypto v0.32.0 // indirect
|
||||
golang.org/x/net v0.34.0 // indirect
|
||||
golang.org/x/sys v0.29.0 // indirect
|
||||
golang.org/x/text v0.21.0 // indirect
|
||||
google.golang.org/protobuf v1.34.1 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
modernc.org/libc v1.22.5 // indirect
|
||||
|
||||
34
go.sum
34
go.sum
@@ -39,12 +39,14 @@ github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ=
|
||||
github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo=
|
||||
github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
|
||||
github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
|
||||
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
|
||||
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
|
||||
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
|
||||
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
|
||||
@@ -64,6 +66,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
|
||||
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
|
||||
github.com/openai/openai-go/v2 v2.7.0 h1:/8MSFCXcasin7AyuWQ2au6FraXL71gzAs+VfbMv+J3k=
|
||||
github.com/openai/openai-go/v2 v2.7.0/go.mod h1:jrJs23apqJKKbT+pqtFgNKpRju/KP9zpUTZhz3GElQE=
|
||||
github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
|
||||
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
|
||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
||||
@@ -86,6 +90,16 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
|
||||
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
|
||||
github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
|
||||
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
|
||||
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
|
||||
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
|
||||
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
|
||||
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
|
||||
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
|
||||
github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
|
||||
@@ -93,16 +107,16 @@ github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZ
|
||||
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
|
||||
golang.org/x/arch v0.8.0 h1:3wRIsP3pM4yUptoR96otTUOXI367OS0+c9eeRi9doIc=
|
||||
golang.org/x/arch v0.8.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys=
|
||||
golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=
|
||||
golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc=
|
||||
golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
|
||||
golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/text v0.15.0 h1:h1V/4gjBv8v9cjcR6+AR5+/cIYK5N/WAgiv4xlsEtAk=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
|
||||
golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
|
||||
google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
|
||||
299
main.go
299
main.go
@@ -1,90 +1,209 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-contrib/cors"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// 预设代码
|
||||
type PresetCode struct {
|
||||
ID uint `gorm:"primarykey" json:"id"`
|
||||
Query string `gorm:"unique" json:"query"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
|
||||
type PresetCodeInput struct {
|
||||
Code string `json:"code" binding:"required"`
|
||||
Query string `json:"query" binding:"required"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
db, err := gorm.Open(sqlite.Open("database.db"), &gorm.Config{})
|
||||
|
||||
if err != nil {
|
||||
panic("fail to get data")
|
||||
}
|
||||
|
||||
db.AutoMigrate(&PresetCode{})
|
||||
|
||||
r := gin.Default()
|
||||
|
||||
config := cors.DefaultConfig()
|
||||
config.AllowMethods = []string{"GET", "POST", "DELETE", "PUT"}
|
||||
config.AllowOrigins = []string{
|
||||
"https://code.xuyue.cc",
|
||||
"http://code.xuyue.cc",
|
||||
"http://localhost:3000",
|
||||
}
|
||||
|
||||
r.Use(cors.New(config))
|
||||
|
||||
r.GET("/", func(c *gin.Context) {
|
||||
var codes []PresetCode
|
||||
db.Order("id desc").Find(&codes)
|
||||
c.JSON(http.StatusOK, gin.H{"data": codes})
|
||||
})
|
||||
|
||||
r.GET("/query/:query", func(c *gin.Context) {
|
||||
var code PresetCode
|
||||
db.Where("query = ?", c.Param("query")).First(&code)
|
||||
if code.Code != "" {
|
||||
c.JSON(http.StatusOK, gin.H{"data": code})
|
||||
} else {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Record not found!"})
|
||||
}
|
||||
})
|
||||
|
||||
r.POST("/", func(c *gin.Context) {
|
||||
var input PresetCodeInput
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
code := PresetCode{Code: input.Code, Query: input.Query}
|
||||
if err := db.Create(&code).Error; err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"data": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": code})
|
||||
})
|
||||
|
||||
r.DELETE("/:id", func(c *gin.Context) {
|
||||
var code PresetCode
|
||||
if err := db.Where("id = ?", c.Param("id")).First(&code).Error; err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Record not found!"})
|
||||
return
|
||||
}
|
||||
|
||||
db.Delete(&code)
|
||||
c.JSON(http.StatusOK, gin.H{"data": true})
|
||||
})
|
||||
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
|
||||
r.Run()
|
||||
}
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-contrib/cors"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/joho/godotenv"
|
||||
"github.com/openai/openai-go/v2"
|
||||
"github.com/openai/openai-go/v2/option"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// 预设代码
|
||||
type PresetCode struct {
|
||||
ID uint `gorm:"primarykey" json:"id"`
|
||||
Query string `gorm:"unique" json:"query"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
|
||||
type PresetCodeInput struct {
|
||||
Code string `json:"code" binding:"required"`
|
||||
Query string `json:"query" binding:"required"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
err := godotenv.Load()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
db, err := gorm.Open(sqlite.Open("database.db"), &gorm.Config{})
|
||||
|
||||
if err != nil {
|
||||
panic("fail to get data")
|
||||
}
|
||||
|
||||
db.AutoMigrate(&PresetCode{})
|
||||
|
||||
r := gin.Default()
|
||||
|
||||
config := cors.DefaultConfig()
|
||||
config.AllowMethods = []string{"GET", "POST", "DELETE", "PUT"}
|
||||
config.AllowOrigins = []string{
|
||||
"https://code.xuyue.cc",
|
||||
"http://code.xuyue.cc",
|
||||
"http://localhost:3000",
|
||||
}
|
||||
|
||||
r.Use(cors.New(config))
|
||||
|
||||
r.GET("/", func(c *gin.Context) {
|
||||
var codes []PresetCode
|
||||
db.Order("id desc").Find(&codes)
|
||||
c.JSON(http.StatusOK, gin.H{"data": codes})
|
||||
})
|
||||
|
||||
r.GET("/query/:query", func(c *gin.Context) {
|
||||
var code PresetCode
|
||||
db.Where("query = ?", c.Param("query")).First(&code)
|
||||
if code.Code != "" {
|
||||
c.JSON(http.StatusOK, gin.H{"data": code})
|
||||
} else {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Record not found!"})
|
||||
}
|
||||
})
|
||||
|
||||
r.POST("/", func(c *gin.Context) {
|
||||
var input PresetCodeInput
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
code := PresetCode{Code: input.Code, Query: input.Query}
|
||||
if err := db.Create(&code).Error; err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"data": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": code})
|
||||
})
|
||||
|
||||
r.DELETE("/:id", func(c *gin.Context) {
|
||||
var code PresetCode
|
||||
if err := db.Where("id = ?", c.Param("id")).First(&code).Error; err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Record not found!"})
|
||||
return
|
||||
}
|
||||
|
||||
db.Delete(&code)
|
||||
c.JSON(http.StatusOK, gin.H{"data": true})
|
||||
})
|
||||
|
||||
r.POST("/ai", func(c *gin.Context) {
|
||||
var payload struct {
|
||||
Code string `json:"code"`
|
||||
ErrorInfo string `json:"error_info"`
|
||||
Language string `json:"language"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&payload); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
systemPrompt := "你是编程老师，擅长分析代码和错误信息，一般出错在语法和格式，请指出错误在第几行，并给出中文的、简要的解决方法。用 markdown 格式返回。"
|
||||
userPrompt := fmt.Sprintf("编程语言：%s\n代码：\n```%s\n```\n错误信息：\n```%s\n```", payload.Language, payload.Code, payload.ErrorInfo)
|
||||
|
||||
apiKey := os.Getenv("API_KEY")
|
||||
|
||||
if apiKey == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "API_KEY is not set"})
|
||||
return
|
||||
}
|
||||
|
||||
client := openai.NewClient(
|
||||
option.WithBaseURL("https://api.deepseek.com"),
|
||||
option.WithAPIKey(apiKey),
|
||||
)
|
||||
|
||||
ctx := c.Request.Context()
|
||||
stream := client.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
|
||||
Messages: []openai.ChatCompletionMessageParamUnion{
|
||||
openai.SystemMessage(systemPrompt),
|
||||
openai.UserMessage(userPrompt),
|
||||
},
|
||||
Seed: openai.Int(0),
|
||||
Model: "deepseek-chat",
|
||||
})
|
||||
|
||||
if err := stream.Err(); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
defer stream.Close()
|
||||
|
||||
header := c.Writer.Header()
|
||||
header.Set("Content-Type", "text/event-stream")
|
||||
|
||||
if _, ok := c.Writer.(http.Flusher); !ok {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"})
|
||||
return
|
||||
}
|
||||
|
||||
c.Status(http.StatusOK)
|
||||
|
||||
replacer := strings.NewReplacer("\r\n", "\n", "\r", "\n")
|
||||
|
||||
writeSSE := func(w io.Writer, event string, payload gin.H) {
|
||||
sanitized := make(gin.H, len(payload))
|
||||
for k, v := range payload {
|
||||
sanitized[k] = v
|
||||
}
|
||||
if data, ok := payload["data"].(string); ok {
|
||||
sanitized["data"] = replacer.Replace(data)
|
||||
}
|
||||
bytes, err := json.Marshal(sanitized)
|
||||
if err != nil {
|
||||
log.Printf("failed to marshal sse payload: %v", err)
|
||||
return
|
||||
}
|
||||
if event != "" {
|
||||
fmt.Fprintf(w, "event: %s\n", event)
|
||||
}
|
||||
fmt.Fprintf(w, "data: %s\n\n", bytes)
|
||||
}
|
||||
|
||||
sentDone := false
|
||||
|
||||
c.Stream(func(w io.Writer) bool {
|
||||
if sentDone {
|
||||
return false
|
||||
}
|
||||
|
||||
for stream.Next() {
|
||||
chunk := stream.Current()
|
||||
var builder strings.Builder
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != "" {
|
||||
builder.WriteString(choice.Delta.Content)
|
||||
}
|
||||
}
|
||||
if builder.Len() == 0 {
|
||||
continue
|
||||
}
|
||||
writeSSE(w, "chunk", gin.H{"data": builder.String()})
|
||||
return true
|
||||
}
|
||||
|
||||
sentDone = true
|
||||
if err := stream.Err(); err != nil {
|
||||
writeSSE(w, "error", gin.H{"message": err.Error()})
|
||||
}
|
||||
writeSSE(w, "done", gin.H{"data": ""})
|
||||
return true
|
||||
})
|
||||
})
|
||||
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
|
||||
r.Run()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user