[SMG-GO] implement a Go SGLang Model Gateway - OpenAI Compatible API Server (#14770)
This commit is contained in:
ybyang
@@ -0,0 +1,556 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
sglang "github.com/sglang/sglang-go-grpc-sdk"
|
||||
"github.com/valyala/fasthttp"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"oai_server/models"
|
||||
"oai_server/service"
|
||||
"oai_server/utils"
|
||||
)
|
||||
|
||||
// ChatHandler handles chat completion requests
|
||||
type ChatHandler struct {
|
||||
logger *zap.Logger
|
||||
service *service.SGLangService
|
||||
}
|
||||
|
||||
// NewChatHandler creates a new chat handler
|
||||
func NewChatHandler(logger *zap.Logger, svc *service.SGLangService) *ChatHandler {
|
||||
return &ChatHandler{
|
||||
logger: logger,
|
||||
service: svc,
|
||||
}
|
||||
}
|
||||
|
||||
// recvResult holds the result of a RecvJSON() call
|
||||
type recvResult struct {
|
||||
chunkJSON string
|
||||
err error
|
||||
}
|
||||
|
||||
// HandleChatCompletion handles POST /v1/chat/completions
|
||||
func (h *ChatHandler) HandleChatCompletion(ctx *fasthttp.RequestCtx) {
|
||||
var req models.ChatRequest
|
||||
if err := json.Unmarshal(ctx.PostBody(), &req); err != nil {
|
||||
h.logger.Warn("Invalid chat completion request", zap.Error(err))
|
||||
utils.RespondError(ctx, 400, fmt.Sprintf("Invalid request: %v", err), "invalid_request_error")
|
||||
return
|
||||
}
|
||||
|
||||
path := string(ctx.Path())
|
||||
|
||||
defer func() {
|
||||
statusCode := ctx.Response.StatusCode()
|
||||
if statusCode == 0 {
|
||||
statusCode = 200
|
||||
}
|
||||
h.logHTTPResponse(statusCode, path)
|
||||
}()
|
||||
|
||||
// Convert to SGLang format
|
||||
messages := make([]sglang.ChatMessage, len(req.Messages))
|
||||
for i, msg := range req.Messages {
|
||||
role, roleOk := msg["role"]
|
||||
content, contentOk := msg["content"]
|
||||
|
||||
// Validate role
|
||||
if !roleOk || role == "" {
|
||||
h.logger.Warn("Missing or empty role in message", zap.Int("message_index", i))
|
||||
utils.RespondError(ctx, 400, "Message role is required and cannot be empty", "invalid_request_error")
|
||||
return
|
||||
}
|
||||
|
||||
// Ensure content is always a string (not null)
|
||||
// Chat template requires content field to be present, even if empty
|
||||
// If content is missing or null, use empty string
|
||||
contentStr := ""
|
||||
if contentOk && content != "" {
|
||||
contentStr = content
|
||||
}
|
||||
|
||||
messages[i] = sglang.ChatMessage{
|
||||
Role: role,
|
||||
Content: contentStr,
|
||||
}
|
||||
}
|
||||
|
||||
sglReq := sglang.ChatCompletionRequest{
|
||||
Model: req.Model,
|
||||
Messages: messages,
|
||||
Stream: req.Stream,
|
||||
}
|
||||
|
||||
if req.Temperature != nil {
|
||||
temp := float32(*req.Temperature)
|
||||
sglReq.Temperature = &temp
|
||||
}
|
||||
if req.TopP != nil {
|
||||
topP := float32(*req.TopP)
|
||||
sglReq.TopP = &topP
|
||||
}
|
||||
if req.MaxCompletionTokens != nil {
|
||||
sglReq.MaxCompletionTokens = req.MaxCompletionTokens
|
||||
} else if req.MaxTokens != nil {
|
||||
sglReq.MaxCompletionTokens = req.MaxTokens
|
||||
}
|
||||
|
||||
requestCtx := context.Background()
|
||||
|
||||
if req.Stream {
|
||||
h.handleStreamingCompletion(ctx, requestCtx, sglReq)
|
||||
} else {
|
||||
h.handleNonStreamingCompletion(ctx, requestCtx, sglReq)
|
||||
}
|
||||
}
|
||||
|
||||
// isBrokenPipeError checks if the error is a broken pipe error (client disconnected)
|
||||
func isBrokenPipeError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
errStr := err.Error()
|
||||
return strings.Contains(errStr, "broken pipe") ||
|
||||
strings.Contains(errStr, "connection reset by peer") ||
|
||||
strings.Contains(errStr, "connection closed") ||
|
||||
strings.Contains(errStr, "write: connection closed")
|
||||
}
|
||||
|
||||
// logHTTPResponse logs HTTP response with colored output
|
||||
func (h *ChatHandler) logHTTPResponse(statusCode int, path string) {
|
||||
var statusText string
|
||||
var colorCode string
|
||||
|
||||
switch {
|
||||
case statusCode >= 200 && statusCode < 300:
|
||||
colorCode = "\033[32m" // Green
|
||||
statusText = "OK"
|
||||
case statusCode >= 300 && statusCode < 400:
|
||||
colorCode = "\033[33m" // Yellow
|
||||
statusText = "Redirect"
|
||||
case statusCode >= 400 && statusCode < 500:
|
||||
colorCode = "\033[33m" // Yellow
|
||||
statusText = "Client Error"
|
||||
case statusCode >= 500:
|
||||
colorCode = "\033[31m" // Red
|
||||
statusText = "Server Error"
|
||||
default:
|
||||
colorCode = "\033[37m" // White
|
||||
statusText = "Unknown"
|
||||
}
|
||||
|
||||
resetCode := "\033[0m"
|
||||
msg := fmt.Sprintf("%s[%d %s]%s %s", colorCode, statusCode, statusText, resetCode, path)
|
||||
h.logger.Info(msg)
|
||||
}
|
||||
|
||||
func (h *ChatHandler) handleStreamingCompletion(ctx *fasthttp.RequestCtx, requestCtx context.Context, req sglang.ChatCompletionRequest) {
|
||||
|
||||
ctx.SetContentType("text/event-stream")
|
||||
ctx.Response.Header.Set("Cache-Control", "no-cache")
|
||||
ctx.Response.Header.Set("Connection", "keep-alive")
|
||||
ctx.Response.Header.Set("X-Accel-Buffering", "no")
|
||||
ctx.SetStatusCode(200)
|
||||
|
||||
var clientDisconnected bool
|
||||
// Flush timeout: prevent deadlock if client is slow or disconnected
|
||||
// This timeout should be longer than typical network latency but shorter than client timeout
|
||||
const flushTimeout = 5 * time.Second
|
||||
|
||||
ctx.SetBodyStreamWriter(func(w *bufio.Writer) {
|
||||
streamCtx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
stream, err := h.service.Client().CreateChatCompletionStream(streamCtx, req)
|
||||
if err != nil {
|
||||
h.logger.Error("Failed to create chat completion stream",
|
||||
zap.Error(err),
|
||||
zap.String("model", req.Model),
|
||||
)
|
||||
// Use sendSSEError to send error in consistent format
|
||||
errInfo, sendErr := h.sendSSEError(w, err)
|
||||
if sendErr != nil {
|
||||
h.logger.Warn("Failed to send SSE error", zap.Error(sendErr))
|
||||
} else if errInfo.IsTimeout {
|
||||
h.logger.Error("Stream creation timeout", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := stream.Close(); closeErr != nil {
|
||||
h.logger.Warn("Failed to close stream", zap.Error(closeErr))
|
||||
}
|
||||
}()
|
||||
|
||||
// Use a single dedicated goroutine to continuously call RecvJSON() and send results via channel
|
||||
recvChan := make(chan recvResult, 20)
|
||||
recvGoroutineDone := make(chan struct{})
|
||||
go func() {
|
||||
defer func() {
|
||||
close(recvChan)
|
||||
close(recvGoroutineDone)
|
||||
}()
|
||||
for {
|
||||
// Check context before calling RecvJSON() to avoid blocking if context is cancelled
|
||||
select {
|
||||
case <-streamCtx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Call RecvJSON() - this may block, but stream.Close() will unblock it
|
||||
// when context is cancelled (called from main loop)
|
||||
chunkJSON, err := stream.RecvJSON()
|
||||
|
||||
// Check context again after RecvJSON() returns
|
||||
select {
|
||||
case <-streamCtx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Send to channel (may block if channel is full)
|
||||
// If channel is full, this will block until main loop reads from it
|
||||
// This is acceptable because main loop should be actively reading
|
||||
select {
|
||||
case recvChan <- recvResult{chunkJSON: chunkJSON, err: err}:
|
||||
if err != nil {
|
||||
// EOF or other error, stop the goroutine
|
||||
return
|
||||
}
|
||||
case <-streamCtx.Done():
|
||||
// Context cancelled while sending, stop the goroutine
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
if clientDisconnected {
|
||||
cancel()
|
||||
// Close stream immediately to unblock RecvJSON() calls
|
||||
stream.Close()
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-streamCtx.Done():
|
||||
// Close stream to ensure RecvJSON() goroutine can exit
|
||||
stream.Close()
|
||||
return
|
||||
case result, ok := <-recvChan:
|
||||
if !ok {
|
||||
// Channel closed, stream ended
|
||||
return
|
||||
}
|
||||
if result.err == io.EOF {
|
||||
if !clientDisconnected {
|
||||
w.WriteString("data: [DONE]\n\n")
|
||||
// Flush with timeout to prevent deadlock
|
||||
flushDone := make(chan error, 1)
|
||||
go func() {
|
||||
flushDone <- w.Flush()
|
||||
}()
|
||||
flushCtx, flushCancel := context.WithTimeout(streamCtx, flushTimeout)
|
||||
defer flushCancel()
|
||||
select {
|
||||
case flushErr := <-flushDone:
|
||||
if flushErr != nil && !isBrokenPipeError(flushErr) {
|
||||
h.logger.Warn("Final flush error", zap.Error(flushErr))
|
||||
}
|
||||
case <-flushCtx.Done():
|
||||
if flushCtx.Err() == context.DeadlineExceeded {
|
||||
h.logger.Warn("Final flush timeout", zap.Duration("timeout", flushTimeout))
|
||||
}
|
||||
case <-streamCtx.Done():
|
||||
// Context cancelled, skip flush
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
if result.err != nil {
|
||||
if result.err == context.Canceled || result.err == context.DeadlineExceeded {
|
||||
return
|
||||
}
|
||||
// Send error to client before closing
|
||||
errInfo, sendErr := h.sendSSEError(w, result.err)
|
||||
if sendErr != nil {
|
||||
h.logger.Warn("Failed to send SSE error", zap.Error(sendErr))
|
||||
}
|
||||
if errInfo.IsTimeout {
|
||||
h.logger.Error("Stream timeout error", zap.Error(result.err))
|
||||
} else {
|
||||
h.logger.Error("Stream error", zap.Error(result.err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if result.chunkJSON == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
w.WriteString("data: ")
|
||||
w.WriteString(result.chunkJSON)
|
||||
w.WriteString("\n\n")
|
||||
|
||||
// Flush with timeout to prevent deadlock:
|
||||
// If Flush blocks indefinitely (slow client), RecvJSON goroutine may fill recvChan
|
||||
// and then block trying to send, causing deadlock
|
||||
// Note: bufio.Writer.Flush() doesn't have a timeout parameter, so we use
|
||||
// a goroutine + select pattern to implement timeout behavior
|
||||
flushDone := make(chan error, 1)
|
||||
go func() {
|
||||
flushDone <- w.Flush()
|
||||
}()
|
||||
|
||||
flushCtx, flushCancel := context.WithTimeout(streamCtx, flushTimeout)
|
||||
defer flushCancel()
|
||||
|
||||
select {
|
||||
case err := <-flushDone:
|
||||
if err != nil {
|
||||
if isBrokenPipeError(err) {
|
||||
clientDisconnected = true
|
||||
cancel()
|
||||
// Close stream immediately to unblock RecvJSON() calls
|
||||
stream.Close()
|
||||
return
|
||||
}
|
||||
h.logger.Warn("Flush error", zap.Error(err))
|
||||
}
|
||||
case <-flushCtx.Done():
|
||||
// Flush timeout: client may be slow or disconnected
|
||||
// Continue processing to avoid deadlock, but mark as disconnected
|
||||
if flushCtx.Err() == context.DeadlineExceeded {
|
||||
h.logger.Warn("Flush timeout, client may be slow or disconnected", zap.Duration("timeout", flushTimeout))
|
||||
}
|
||||
clientDisconnected = true
|
||||
cancel()
|
||||
stream.Close()
|
||||
return
|
||||
case <-streamCtx.Done():
|
||||
// Context cancelled, stop flushing
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (h *ChatHandler) handleNonStreamingCompletion(ctx *fasthttp.RequestCtx, requestCtx context.Context, req sglang.ChatCompletionRequest) {
|
||||
resp, err := h.service.Client().CreateChatCompletion(requestCtx, req)
|
||||
if err != nil {
|
||||
h.logger.Error("Failed to create chat completion",
|
||||
zap.Error(err),
|
||||
zap.String("model", req.Model),
|
||||
)
|
||||
utils.RespondError(ctx, 500, fmt.Sprintf("Failed to create completion: %v", err), "server_error")
|
||||
return
|
||||
}
|
||||
|
||||
// Convert to OpenAI format
|
||||
response := utils.BuildResponseBase(resp.ID, resp.Created, resp.Model)
|
||||
response["object"] = "chat.completion"
|
||||
|
||||
choices := make([]map[string]interface{}, len(resp.Choices))
|
||||
for i, choice := range resp.Choices {
|
||||
choiceMap := map[string]interface{}{
|
||||
"index": choice.Index,
|
||||
"message": map[string]interface{}{
|
||||
"role": choice.Message.Role,
|
||||
"content": choice.Message.Content,
|
||||
},
|
||||
"finish_reason": choice.FinishReason,
|
||||
}
|
||||
if len(choice.Message.ToolCalls) > 0 {
|
||||
toolCalls := make([]map[string]interface{}, len(choice.Message.ToolCalls))
|
||||
for j, tc := range choice.Message.ToolCalls {
|
||||
toolCalls[j] = map[string]interface{}{
|
||||
"id": tc.ID,
|
||||
"type": tc.Type,
|
||||
"function": map[string]interface{}{"name": tc.Function.Name, "arguments": tc.Function.Arguments},
|
||||
}
|
||||
}
|
||||
choiceMap["message"].(map[string]interface{})["tool_calls"] = toolCalls
|
||||
}
|
||||
choices[i] = choiceMap
|
||||
}
|
||||
response["choices"] = choices
|
||||
|
||||
// Usage is always present (not a pointer)
|
||||
response["usage"] = map[string]interface{}{
|
||||
"prompt_tokens": resp.Usage.PromptTokens,
|
||||
"completion_tokens": resp.Usage.CompletionTokens,
|
||||
"total_tokens": resp.Usage.TotalTokens,
|
||||
}
|
||||
|
||||
ctx.SetStatusCode(200)
|
||||
ctx.SetContentType("application/json")
|
||||
jsonData, _ := json.Marshal(response)
|
||||
ctx.Write(jsonData)
|
||||
}
|
||||
|
||||
// StreamErrorInfo holds parsed error information
|
||||
type StreamErrorInfo struct {
|
||||
Message string
|
||||
Type string
|
||||
Code int
|
||||
IsTimeout bool
|
||||
}
|
||||
|
||||
// parseStreamError parses error type and code
|
||||
func parseStreamError(err error) StreamErrorInfo {
|
||||
if err == nil {
|
||||
return StreamErrorInfo{}
|
||||
}
|
||||
|
||||
errorMsg := err.Error()
|
||||
// Check timeout error by message prefix
|
||||
isTimeout := strings.HasPrefix(errorMsg, "stream.Recv() timeout") || strings.Contains(errorMsg, "timeout after")
|
||||
|
||||
errorType := "server_error"
|
||||
errorCode := 500
|
||||
if isTimeout {
|
||||
errorType = "timeout_error"
|
||||
errorCode = 504
|
||||
}
|
||||
|
||||
return StreamErrorInfo{
|
||||
Message: errorMsg,
|
||||
Type: errorType,
|
||||
Code: errorCode,
|
||||
IsTimeout: isTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
// formatErrorJSON formats error as OpenAI JSON
|
||||
func formatErrorJSON(errInfo StreamErrorInfo) string {
|
||||
errorObj := map[string]interface{}{
|
||||
"error": map[string]interface{}{
|
||||
"message": errInfo.Message,
|
||||
"type": errInfo.Type,
|
||||
"code": errInfo.Code,
|
||||
},
|
||||
}
|
||||
jsonBytes, _ := json.Marshal(errorObj)
|
||||
return string(jsonBytes)
|
||||
}
|
||||
|
||||
// sendSSEError sends SSE error response. Callers should log errors.
|
||||
func (h *ChatHandler) sendSSEError(w *bufio.Writer, err error) (StreamErrorInfo, error) {
|
||||
errInfo := parseStreamError(err)
|
||||
errorJSON := formatErrorJSON(errInfo)
|
||||
|
||||
w.WriteString("data: ")
|
||||
w.WriteString(errorJSON)
|
||||
w.WriteString("\n\n")
|
||||
|
||||
if flushErr := w.Flush(); flushErr != nil && !isBrokenPipeError(flushErr) {
|
||||
h.logger.Warn("Failed to flush error response", zap.Error(flushErr))
|
||||
return errInfo, flushErr
|
||||
}
|
||||
|
||||
return errInfo, nil
|
||||
}
|
||||
|
||||
// HandleGenerate handles POST /generate (SGLang native API)
|
||||
func (h *ChatHandler) HandleGenerate(ctx *fasthttp.RequestCtx) {
|
||||
path := string(ctx.Path())
|
||||
|
||||
defer func() {
|
||||
statusCode := ctx.Response.StatusCode()
|
||||
if statusCode == 0 {
|
||||
statusCode = 200
|
||||
}
|
||||
h.logHTTPResponse(statusCode, path)
|
||||
}()
|
||||
|
||||
// Parse request body
|
||||
var req map[string]interface{}
|
||||
if err := json.Unmarshal(ctx.PostBody(), &req); err != nil {
|
||||
h.logger.Warn("Invalid generate request", zap.Error(err))
|
||||
utils.RespondError(ctx, 400, fmt.Sprintf("Invalid request: %v", err), "invalid_request_error")
|
||||
return
|
||||
}
|
||||
|
||||
// Extract text and sampling_params
|
||||
text, ok := req["text"].(string)
|
||||
if !ok || text == "" {
|
||||
utils.RespondError(ctx, 400, "Missing or invalid 'text' field", "invalid_request_error")
|
||||
return
|
||||
}
|
||||
|
||||
samplingParams, _ := req["sampling_params"].(map[string]interface{})
|
||||
if samplingParams == nil {
|
||||
samplingParams = make(map[string]interface{})
|
||||
}
|
||||
|
||||
// Convert to chat completion format for processing
|
||||
chatReq := sglang.ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []sglang.ChatMessage{{Role: "user", Content: text}},
|
||||
Stream: false,
|
||||
}
|
||||
|
||||
// Copy sampling params
|
||||
if maxNewTokens, ok := samplingParams["max_new_tokens"].(float64); ok {
|
||||
tokens := int(maxNewTokens)
|
||||
chatReq.MaxCompletionTokens = &tokens
|
||||
}
|
||||
if temp, ok := samplingParams["temperature"].(float64); ok {
|
||||
temp32 := float32(temp)
|
||||
chatReq.Temperature = &temp32
|
||||
}
|
||||
if topP, ok := samplingParams["top_p"].(float64); ok {
|
||||
topP32 := float32(topP)
|
||||
chatReq.TopP = &topP32
|
||||
}
|
||||
if topK, ok := samplingParams["top_k"].(float64); ok {
|
||||
topKInt := int(topK)
|
||||
chatReq.TopK = &topKInt
|
||||
}
|
||||
|
||||
requestCtx := context.Background()
|
||||
|
||||
// Use non-streaming completion for /generate endpoint
|
||||
resp, err := h.service.Client().CreateChatCompletion(requestCtx, chatReq)
|
||||
if err != nil {
|
||||
h.logger.Error("Failed to create completion",
|
||||
zap.Error(err),
|
||||
)
|
||||
utils.RespondError(ctx, 500, fmt.Sprintf("Failed to create completion: %v", err), "server_error")
|
||||
return
|
||||
}
|
||||
|
||||
// Convert to SGLang /generate response format
|
||||
// meta_info must match SGLang's expected format with completion_tokens at top level
|
||||
finishReason := resp.Choices[0].FinishReason
|
||||
if finishReason == "" {
|
||||
finishReason = "stop"
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"text": resp.Choices[0].Message.Content,
|
||||
"meta_info": map[string]interface{}{
|
||||
"id": resp.ID,
|
||||
"finish_reason": finishReason,
|
||||
"prompt_tokens": resp.Usage.PromptTokens,
|
||||
"completion_tokens": resp.Usage.CompletionTokens,
|
||||
"cached_tokens": 0, // Not available from chat completion API
|
||||
"weight_version": "", // Not available from chat completion API
|
||||
},
|
||||
}
|
||||
|
||||
ctx.SetStatusCode(200)
|
||||
ctx.SetContentType("application/json")
|
||||
jsonData, _ := json.Marshal(response)
|
||||
ctx.Write(jsonData)
|
||||
}
|
||||
Reference in New Issue
Block a user