mirror of
https://github.com/linux-do/new-api.git
synced 2025-09-17 16:06:38 +08:00
113 lines
3.1 KiB
Go
113 lines
3.1 KiB
Go
package service
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/dto"
|
|
)
|
|
|
|
func SetEventStreamHeaders(c *gin.Context) {
|
|
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
|
c.Writer.Header().Set("Cache-Control", "no-cache")
|
|
c.Writer.Header().Set("Connection", "keep-alive")
|
|
c.Writer.Header().Set("Transfer-Encoding", "chunked")
|
|
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
|
}
|
|
|
|
func StringData(c *gin.Context, str string) error {
|
|
//str = strings.TrimPrefix(str, "data: ")
|
|
//str = strings.TrimSuffix(str, "\r")
|
|
c.Render(-1, common.CustomEvent{Data: "data: " + str})
|
|
if flusher, ok := c.Writer.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
} else {
|
|
return errors.New("streaming error: flusher not found")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ObjectData(c *gin.Context, object interface{}) error {
|
|
jsonData, err := json.Marshal(object)
|
|
if err != nil {
|
|
return fmt.Errorf("error marshalling object: %w", err)
|
|
}
|
|
return StringData(c, string(jsonData))
|
|
}
|
|
|
|
func Done(c *gin.Context) {
|
|
_ = StringData(c, "[DONE]")
|
|
}
|
|
|
|
func WssString(c *gin.Context, ws *websocket.Conn, str string) error {
|
|
if ws == nil {
|
|
common.LogError(c, "websocket connection is nil")
|
|
return errors.New("websocket connection is nil")
|
|
}
|
|
common.LogInfo(c, fmt.Sprintf("sending message: %s", str))
|
|
return ws.WriteMessage(1, []byte(str))
|
|
}
|
|
|
|
func WssObject(c *gin.Context, ws *websocket.Conn, object interface{}) error {
|
|
jsonData, err := json.Marshal(object)
|
|
if err != nil {
|
|
return fmt.Errorf("error marshalling object: %w", err)
|
|
}
|
|
if ws == nil {
|
|
common.LogError(c, "websocket connection is nil")
|
|
return errors.New("websocket connection is nil")
|
|
}
|
|
common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData))
|
|
return ws.WriteMessage(1, jsonData)
|
|
}
|
|
|
|
func WssError(c *gin.Context, ws *websocket.Conn, openaiError dto.OpenAIError) {
|
|
errorObj := &dto.RealtimeEvent{
|
|
Type: "error",
|
|
EventId: GetLocalRealtimeID(c),
|
|
Error: &openaiError,
|
|
}
|
|
_ = WssObject(c, ws, errorObj)
|
|
}
|
|
|
|
func GetResponseID(c *gin.Context) string {
|
|
logID := c.GetString(common.RequestIdKey)
|
|
return fmt.Sprintf("chatcmpl-%s", logID)
|
|
}
|
|
|
|
func GetLocalRealtimeID(c *gin.Context) string {
|
|
logID := c.GetString(common.RequestIdKey)
|
|
return fmt.Sprintf("evt_%s", logID)
|
|
}
|
|
|
|
func GenerateStopResponse(id string, createAt int64, model string, finishReason string) *dto.ChatCompletionsStreamResponse {
|
|
return &dto.ChatCompletionsStreamResponse{
|
|
Id: id,
|
|
Object: "chat.completion.chunk",
|
|
Created: createAt,
|
|
Model: model,
|
|
SystemFingerprint: nil,
|
|
Choices: []dto.ChatCompletionsStreamResponseChoice{
|
|
{
|
|
FinishReason: &finishReason,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func GenerateFinalUsageResponse(id string, createAt int64, model string, usage dto.Usage) *dto.ChatCompletionsStreamResponse {
|
|
return &dto.ChatCompletionsStreamResponse{
|
|
Id: id,
|
|
Object: "chat.completion.chunk",
|
|
Created: createAt,
|
|
Model: model,
|
|
SystemFingerprint: nil,
|
|
Choices: make([]dto.ChatCompletionsStreamResponseChoice, 0),
|
|
Usage: &usage,
|
|
}
|
|
}
|