127 lines
3.1 KiB
Go
127 lines
3.1 KiB
Go
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
|
|
package ffi
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// BatchPostprocessor handles batch postprocessing of stream chunks to reduce FFI overhead
|
|
type BatchPostprocessor struct {
|
|
converter *GrpcResponseConverterHandle
|
|
buffer []string
|
|
batchSize int
|
|
flushInterval time.Duration
|
|
lastFlush time.Time
|
|
timer *time.Timer
|
|
}
|
|
|
|
// NewBatchPostprocessor creates a new batch postprocessor
|
|
func NewBatchPostprocessor(converter *GrpcResponseConverterHandle, batchSize int, flushInterval time.Duration) *BatchPostprocessor {
|
|
if batchSize <= 0 {
|
|
batchSize = 1
|
|
}
|
|
if flushInterval < 0 {
|
|
flushInterval = 0
|
|
}
|
|
|
|
return &BatchPostprocessor{
|
|
converter: converter,
|
|
buffer: make([]string, 0, batchSize),
|
|
batchSize: batchSize,
|
|
flushInterval: flushInterval,
|
|
lastFlush: time.Now(),
|
|
}
|
|
}
|
|
|
|
// AddChunk adds a chunk to the buffer and processes if batch is full
|
|
func (b *BatchPostprocessor) AddChunk(chunkJSON string) (results []string, shouldFlush bool, err error) {
|
|
if b.batchSize == 1 {
|
|
openaiJSON, _, err := PostprocessStreamChunk(b.converter, chunkJSON)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
return []string{openaiJSON}, false, nil
|
|
}
|
|
|
|
b.buffer = append(b.buffer, chunkJSON)
|
|
shouldProcess := len(b.buffer) >= b.batchSize
|
|
shouldFlushTimeout := b.flushInterval > 0 && time.Since(b.lastFlush) >= b.flushInterval
|
|
|
|
if shouldProcess || shouldFlushTimeout {
|
|
return b.processBatch()
|
|
}
|
|
|
|
return nil, false, nil
|
|
}
|
|
|
|
// Flush processes any remaining chunks in the buffer
|
|
func (b *BatchPostprocessor) Flush() (results []string, err error) {
|
|
if len(b.buffer) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
res, _, err := b.processBatch()
|
|
return res, err
|
|
}
|
|
|
|
// processBatch processes the current buffer and returns results
|
|
func (b *BatchPostprocessor) processBatch() (results []string, shouldFlush bool, err error) {
|
|
if len(b.buffer) == 0 {
|
|
return nil, false, nil
|
|
}
|
|
|
|
var sb strings.Builder
|
|
sb.Grow(len(b.buffer) * 200)
|
|
sb.WriteString(`[`)
|
|
for i, chunkJSONStr := range b.buffer {
|
|
if i > 0 {
|
|
sb.WriteString(`,`)
|
|
}
|
|
sb.WriteString(chunkJSONStr)
|
|
}
|
|
sb.WriteString(`]`)
|
|
bufferJSON := sb.String()
|
|
|
|
resultJSON, _, err := PostprocessStreamChunksBatch(
|
|
b.converter,
|
|
bufferJSON,
|
|
b.batchSize*2,
|
|
)
|
|
if err != nil {
|
|
return nil, false, fmt.Errorf("batch postprocessing failed: %w", err)
|
|
}
|
|
|
|
var resultArray []json.RawMessage
|
|
if err := json.Unmarshal([]byte(resultJSON), &resultArray); err != nil {
|
|
return nil, false, fmt.Errorf("failed to unmarshal results array: %w", err)
|
|
}
|
|
|
|
resultStrings := make([]string, 0, len(resultArray))
|
|
for _, rawMsg := range resultArray {
|
|
resultStrings = append(resultStrings, string(rawMsg))
|
|
}
|
|
|
|
b.buffer = b.buffer[:0]
|
|
b.lastFlush = time.Now()
|
|
|
|
if b.timer != nil {
|
|
b.timer.Stop()
|
|
b.timer = nil
|
|
}
|
|
|
|
return resultStrings, false, nil
|
|
}
|
|
|
|
// Reset clears the buffer and resets the postprocessor state
|
|
func (b *BatchPostprocessor) Reset() {
|
|
b.buffer = b.buffer[:0]
|
|
b.lastFlush = time.Now()
|
|
if b.timer != nil {
|
|
b.timer.Stop()
|
|
b.timer = nil
|
|
}
|
|
}
|