126 lines
2.8 KiB
Go
126 lines
2.8 KiB
Go
// Streaming example demonstrating real-time streaming with SGLang Go SDK
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/sglang/sglang-go-grpc-sdk"
|
|
)
|
|
|
|
func main() {
|
|
// Get configuration from environment or command line
|
|
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
|
|
if endpoint == "" {
|
|
endpoint = "grpc://localhost:20000"
|
|
}
|
|
|
|
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
|
|
if tokenizerPath == "" {
|
|
tokenizerPath = "./examples/tokenizer"
|
|
}
|
|
|
|
// Create client
|
|
client, err := sglang.NewClient(sglang.ClientConfig{
|
|
Endpoint: endpoint,
|
|
TokenizerPath: tokenizerPath,
|
|
})
|
|
if err != nil {
|
|
log.Fatalf("Failed to create client: %v", err)
|
|
}
|
|
defer client.Close()
|
|
|
|
// Create streaming chat completion request
|
|
req := sglang.ChatCompletionRequest{
|
|
Model: "default",
|
|
Messages: []sglang.ChatMessage{
|
|
{
|
|
Role: "system",
|
|
Content: "You are a helpful assistant.",
|
|
},
|
|
{
|
|
Role: "user",
|
|
Content: "写一首春天的诗歌",
|
|
},
|
|
},
|
|
Stream: true,
|
|
Temperature: float32Ptr(0.7),
|
|
MaxCompletionTokens: intPtr(500),
|
|
SkipSpecialTokens: true,
|
|
Tools: nil, // Use nil instead of empty slice to avoid template errors
|
|
}
|
|
|
|
// Create streaming completion
|
|
ctx := context.Background()
|
|
stream, err := client.CreateChatCompletionStream(ctx, req)
|
|
if err != nil {
|
|
log.Fatalf("Failed to create stream: %v", err)
|
|
}
|
|
defer stream.Close()
|
|
|
|
fmt.Println("=== Streaming Response ===")
|
|
fmt.Println()
|
|
|
|
var fullContent strings.Builder
|
|
chunkCount := 0
|
|
startTime := time.Now()
|
|
var firstTokenTime time.Time
|
|
firstTokenReceived := false
|
|
|
|
for {
|
|
chunk, err := stream.Recv()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
log.Fatalf("Stream error: %v", err)
|
|
}
|
|
|
|
chunkCount++
|
|
|
|
// Extract content from delta
|
|
for _, choice := range chunk.Choices {
|
|
if choice.Delta.Content != "" {
|
|
fmt.Print(choice.Delta.Content)
|
|
fullContent.WriteString(choice.Delta.Content)
|
|
|
|
// Track first token time (TTFT)
|
|
if !firstTokenReceived {
|
|
firstTokenTime = time.Now()
|
|
firstTokenReceived = true
|
|
ttft := firstTokenTime.Sub(startTime)
|
|
fmt.Printf("\n[TTFT: %v]\n", ttft)
|
|
}
|
|
}
|
|
|
|
if choice.FinishReason != "" {
|
|
fmt.Printf("\n\n[Finished: %s]\n", choice.FinishReason)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Calculate metrics
|
|
if firstTokenReceived {
|
|
elapsed := time.Since(startTime)
|
|
tokensPerSecond := float64(fullContent.Len()) / elapsed.Seconds()
|
|
fmt.Printf("\n=== Metrics ===\n")
|
|
fmt.Printf("Total chunks: %d\n", chunkCount)
|
|
fmt.Printf("Total content length: %d characters\n", fullContent.Len())
|
|
fmt.Printf("Time elapsed: %v\n", elapsed)
|
|
fmt.Printf("Tokens per second: %.2f\n", tokensPerSecond)
|
|
}
|
|
}
|
|
|
|
func float32Ptr(f float32) *float32 {
|
|
return &f
|
|
}
|
|
|
|
func intPtr(i int) *int {
|
|
return &i
|
|
}
|