From 3b18fd4cf560382fae5fa1b4602a130c19268802 Mon Sep 17 00:00:00 2001 From: ybyang <10629930+whybeyoung@users.noreply.github.com> Date: Mon, 17 Nov 2025 02:12:10 +0800 Subject: [PATCH] [router] bindings for go (#13384) Signed-off-by: ybyang --- sgl-router/bindings/golang/.gitignore | 24 + sgl-router/bindings/golang/Cargo.toml | 48 ++ sgl-router/bindings/golang/Makefile | 103 +++ sgl-router/bindings/golang/README.md | 552 +++++++++++++ sgl-router/bindings/golang/client.go | 510 ++++++++++++ sgl-router/bindings/golang/client_test.go | 325 ++++++++ .../bindings/golang/examples/simple/main.go | 85 ++ .../bindings/golang/examples/simple/run.sh | 46 ++ .../golang/examples/streaming/main.go | 125 +++ .../bindings/golang/examples/streaming/run.sh | 46 ++ sgl-router/bindings/golang/go.mod | 3 + .../bindings/golang/integration_test.go | 228 ++++++ .../bindings/golang/internal/ffi/client.go | 228 ++++++ sgl-router/bindings/golang/src/client.rs | 279 +++++++ sgl-router/bindings/golang/src/error.rs | 50 ++ .../bindings/golang/src/grpc_converter.rs | 758 ++++++++++++++++++ sgl-router/bindings/golang/src/lib.rs | 88 ++ sgl-router/bindings/golang/src/memory.rs | 28 + sgl-router/bindings/golang/src/stream.rs | 288 +++++++ sgl-router/bindings/golang/src/tokenizer.rs | 379 +++++++++ sgl-router/bindings/golang/src/tool_parser.rs | 329 ++++++++ sgl-router/bindings/golang/src/utils.rs | 44 + .../grpc/regular/responses/handlers.rs | 5 +- 23 files changed, 4568 insertions(+), 3 deletions(-) create mode 100644 sgl-router/bindings/golang/.gitignore create mode 100644 sgl-router/bindings/golang/Cargo.toml create mode 100644 sgl-router/bindings/golang/Makefile create mode 100644 sgl-router/bindings/golang/README.md create mode 100644 sgl-router/bindings/golang/client.go create mode 100644 sgl-router/bindings/golang/client_test.go create mode 100644 sgl-router/bindings/golang/examples/simple/main.go create mode 100755 sgl-router/bindings/golang/examples/simple/run.sh create mode 100644 sgl-router/bindings/golang/examples/streaming/main.go create mode 100755 sgl-router/bindings/golang/examples/streaming/run.sh create mode 100644 sgl-router/bindings/golang/go.mod create mode 100644 sgl-router/bindings/golang/integration_test.go create mode 100644 sgl-router/bindings/golang/internal/ffi/client.go create mode 100644 sgl-router/bindings/golang/src/client.rs create mode 100644 sgl-router/bindings/golang/src/error.rs create mode 100644 sgl-router/bindings/golang/src/grpc_converter.rs create mode 100644 sgl-router/bindings/golang/src/lib.rs create mode 100644 sgl-router/bindings/golang/src/memory.rs create mode 100644 sgl-router/bindings/golang/src/stream.rs create mode 100644 sgl-router/bindings/golang/src/tokenizer.rs create mode 100644 sgl-router/bindings/golang/src/tool_parser.rs create mode 100644 sgl-router/bindings/golang/src/utils.rs diff --git a/sgl-router/bindings/golang/.gitignore b/sgl-router/bindings/golang/.gitignore new file mode 100644 index 000000000..9e1dd2429 --- /dev/null +++ b/sgl-router/bindings/golang/.gitignore @@ -0,0 +1,24 @@ +# Build artifacts +target/ +lib/ + +# Compiled binaries +examples/simple/simple +examples/streaming/streaming + +# Go build artifacts +*.o +*.a +*.so +*.dylib + +# IDE and editor files +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# Environment files +.env +.env.local diff --git a/sgl-router/bindings/golang/Cargo.toml b/sgl-router/bindings/golang/Cargo.toml new file mode 100644 index 000000000..eaedea36c --- /dev/null +++ b/sgl-router/bindings/golang/Cargo.toml @@ -0,0 +1,48 @@ +[package] +name = "sglang-router-golang" +version = "0.2.2" +edition = "2021" + +[lib] +name = "sglang_router_rs" +crate-type = ["cdylib"] + +[dependencies] +tokio = { version = "1.42.0", features = ["full"] } +serde_json = { version = "1.0", default-features = false, features = [ + "std", + "preserve_order", +] } +uuid = { version = "1.10", features = ["v4", "serde"] } +once_cell = "1.21.3" +futures-util = "0.3" +tracing = "0.1" + +[dependencies.sglang-router] +path = "../.." +default-features = true +package = "sglang-router" + +[features] +default = [] +vendored-openssl = ["sglang-router/vendored-openssl"] + +[profile.release] +opt-level = "z" # Optimize for size +lto = "fat" # Full LTO for smaller binaries +codegen-units = 1 # Better optimization, slower compile +strip = true # Strip debug symbols + +[profile.ci] +inherits = "release" +opt-level = 2 # Lighter optimization (still fast runtime, much faster compile) +lto = "thin" # Thin LTO - good balance +codegen-units = 16 # More parallelization for faster builds +strip = true + +[profile.dev] +opt-level = 0 +debug = 1 +split-debuginfo = "unpacked" +incremental = true +codegen-units = 256 diff --git a/sgl-router/bindings/golang/Makefile b/sgl-router/bindings/golang/Makefile new file mode 100644 index 000000000..a1b734046 --- /dev/null +++ b/sgl-router/bindings/golang/Makefile @@ -0,0 +1,103 @@ +# Makefile for sglang-router golang bindings +# This builds the Rust FFI library and provides convenience targets for Go development + +# Configuration +CARGO_BUILD_DIR ?= $(shell pwd)/target +BUILD_MODE ?= release +LIB_NAME = libsglang_router_rs + +# Detect OS +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Linux) + LIB_EXT = .so + LD_LIBRARY_PATH_VAR = LD_LIBRARY_PATH +endif +ifeq ($(UNAME_S),Darwin) + LIB_EXT = .dylib + LD_LIBRARY_PATH_VAR = DYLD_LIBRARY_PATH +endif + +# Paths +ROOT_DIR := $(shell pwd) +RUST_SRC_DIR := $(ROOT_DIR)/src +LIB_BUILD_DIR := $(CARGO_BUILD_DIR)/$(BUILD_MODE) +LIB_BUILD_PATH := $(LIB_BUILD_DIR)/$(LIB_NAME)$(LIB_EXT) +LIB_EXPORT_DIR := $(ROOT_DIR)/lib +LIB_EXPORT_PATH := $(LIB_EXPORT_DIR)/$(LIB_NAME)$(LIB_EXT) + +# Python LDFLAGS (needed for Rust FFI that depends on Python) +PYTHON_LDFLAGS := $(shell python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "") + +# CGO flags - use exported lib directory if available, otherwise build directory +LIB_DIR := $(if $(wildcard $(LIB_EXPORT_PATH)),$(LIB_EXPORT_DIR),$(LIB_BUILD_DIR)) +export CGO_LDFLAGS = -L$(LIB_DIR) -lsglang_router_rs $(PYTHON_LDFLAGS) -ldl +export $(LD_LIBRARY_PATH_VAR) := $(LIB_DIR):$($(LD_LIBRARY_PATH_VAR)) + +.PHONY: all build build-dev lib lib-clean clean test examples help run-simple run-streaming check-lib + +help: + @echo "Available targets:" + @echo " build - Build release version of Rust FFI library" + @echo " build-dev - Build debug version of Rust FFI library" + @echo " lib - Copy built library to ./lib directory" + @echo " lib-clean - Clean ./lib directory" + @echo " clean - Clean build artifacts" + @echo " test - Run Go tests" + @echo " examples - Build example programs" + @echo " run-simple - Run simple example" + @echo " run-streaming - Run streaming example" + +all: build + +build: + @echo "Building Rust FFI library (release mode)..." + @CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --release --manifest-path Cargo.toml + @echo "Library built at: $(LIB_BUILD_PATH)" + +build-dev: + @echo "Building Rust FFI library (debug mode)..." + @CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --manifest-path Cargo.toml + @echo "Library built at: $(LIB_BUILD_DIR)/debug/$(LIB_NAME)$(LIB_EXT)" + +lib: build + @echo "Copying library to ./lib directory..." + @mkdir -p $(LIB_EXPORT_DIR) + @cp $(LIB_BUILD_PATH) $(LIB_EXPORT_PATH) + @echo "Library exported at: $(LIB_EXPORT_PATH)" + +lib-clean: + @echo "Cleaning ./lib directory..." + @rm -rf $(LIB_EXPORT_DIR) + @echo "Lib directory cleaned" + +clean: lib-clean + @echo "Cleaning build artifacts..." + @CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo clean --manifest-path Cargo.toml + @echo "Clean complete" + +test: build + @echo "Running Go tests..." + @go test ./... + +examples: build + @echo "Building example programs..." + @cd examples/simple && go build -o simple main.go + @cd examples/streaming && go build -o streaming main.go + @echo "Examples built" + +run-simple: build + @echo "Running simple example..." + @cd examples/simple && bash run.sh + +run-streaming: build + @echo "Running streaming example..." + @cd examples/streaming && bash run.sh + +# Check if library exists (either in lib dir or build dir) +check-lib: + @if [ ! -f "$(LIB_EXPORT_PATH)" ] && [ ! -f "$(LIB_BUILD_PATH)" ]; then \ + echo "Error: Library not found at $(LIB_EXPORT_PATH) or $(LIB_BUILD_PATH)"; \ + echo "Run 'make build' or 'make lib' first"; \ + exit 1; \ + fi + @echo "Library found at: $(LIB_DIR)/$(LIB_NAME)$(LIB_EXT)" diff --git a/sgl-router/bindings/golang/README.md b/sgl-router/bindings/golang/README.md new file mode 100644 index 000000000..f140440f2 --- /dev/null +++ b/sgl-router/bindings/golang/README.md @@ -0,0 +1,552 @@ +# SGLang Go gRPC SDK + +A high-level Go SDK for interacting with SGLang gRPC API, designed with an OpenAI-style API for familiarity and ease of use. + +**Location**: `sgl-router/bindings/golang/` + +## Table of Contents + +- [Features](#features) +- [Installation](#installation) +- [Quick Start](#quick-start) + - [Basic Usage](#basic-usage) + - [Streaming Usage](#streaming-usage) +- [Examples](#examples) +- [Configuration](#configuration) +- [API Reference](#api-reference) +- [Testing](#testing) + - [Unit Tests](#unit-tests) + - [Integration Tests](#integration-tests) + - [Benchmarks](#benchmarks) +- [Documentation](#documentation) +- [Development](#development) +- [Troubleshooting](#troubleshooting) +- [License](#license) + +## Features + +- **OpenAI-style API**: Familiar interface similar to OpenAI Go SDK +- **Streaming Support**: Real-time streaming chat completions +- **Non-streaming Support**: Simple request/response API +- **Tool Calling**: Support for function calling and tool use +- **Type-safe**: Full Go type definitions for requests and responses +- **Comprehensive Testing**: 18+ unit and integration tests +- **Thread-safe**: All public methods are safe for concurrent use +- **Well-documented**: Full API documentation with examples + +## Installation + +```bash +go get github.com/sglang/sglang-go-grpc-sdk +``` + +### Build Requirements + +- Go 1.21 or later +- Rust toolchain (for building the FFI library) +- Python 3.x (for Python bindings in Rust FFI) +- Tokio runtime for async operations + +## Quick Start + +### Basic Usage (Non-streaming) + +```go +package main + +import ( + "context" + "fmt" + "log" + + "github.com/sglang/sglang-go-grpc-sdk" +) + +func main() { + // Create client + client, err := sglang.NewClient(sglang.ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + }) + if err != nil { + log.Fatal(err) + } + defer client.Close() + + // Create completion + resp, err := client.CreateChatCompletion(context.Background(), sglang.ChatCompletionRequest{ + Model: "default", + Messages: []sglang.ChatMessage{ + {Role: "user", Content: "Hello!"}, + }, + Stream: false, + }) + if err != nil { + log.Fatal(err) + } + + fmt.Println(resp.Choices[0].Message.Content) + fmt.Printf("Usage: Prompt=%d, Completion=%d, Total=%d\n", + resp.Usage.PromptTokens, + resp.Usage.CompletionTokens, + resp.Usage.TotalTokens) +} +``` + +### Streaming Usage + +```go +package main + +import ( + "context" + "fmt" + "io" + "log" + + "github.com/sglang/sglang-go-grpc-sdk" +) + +func main() { + // Create client + client, err := sglang.NewClient(sglang.ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + }) + if err != nil { + log.Fatal(err) + } + defer client.Close() + + // Create streaming completion + ctx := context.Background() + stream, err := client.CreateChatCompletionStream(ctx, sglang.ChatCompletionRequest{ + Model: "default", + Messages: []sglang.ChatMessage{ + {Role: "user", Content: "Tell me a story"}, + }, + Stream: true, + MaxCompletionTokens: intPtr(500), + }) + if err != nil { + log.Fatal(err) + } + defer stream.Close() + + // Read streaming response + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + log.Fatal(err) + } + + for _, choice := range chunk.Choices { + if choice.Delta.Content != "" { + fmt.Print(choice.Delta.Content) + } + } + } + fmt.Println() // newline +} + +// Helper functions for optional pointer fields +func intPtr(i int) *int { + return &i +} + +func float32Ptr(f float32) *float32 { + return &f +} +``` + +## Examples + +The SDK includes several examples in the `examples/` directory: + +- **simple**: Basic non-streaming chat completion example +- **streaming**: Real-time streaming with performance metrics + +### Running Examples + +```bash +# Run simple example +cd bindings/golang/examples/simple +bash run.sh + +# Run streaming example +cd bindings/golang/examples/streaming +bash run.sh + +# Or use Makefile from bindings/golang directory +cd bindings/golang +make run-simple +make run-streaming +``` + +Examples automatically detect the server endpoint and tokenizer path via environment variables or defaults. + +## Configuration + +### Environment Variables + +- `SGL_GRPC_ENDPOINT`: gRPC server endpoint (default: `grpc://localhost:20000`) +- `SGL_TOKENIZER_PATH`: Path to tokenizer directory (required) +- `CARGO_BUILD_DIR`: Rust build output directory (auto-detected if not set) + +### ClientConfig + +```go +type ClientConfig struct { + // Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000") + // Required field. Must include the scheme (grpc://) and port number. + Endpoint string + + // TokenizerPath is the path to the tokenizer directory containing + // tokenizer configuration files (e.g., tokenizer.json, vocab.json) + // Required field. + TokenizerPath string +} +``` + +## API Reference + +### Client Methods + +```go +type Client struct { + // Thread-safe client for SGLang gRPC API +} + +// Creates a new client with the given configuration +func NewClient(config ClientConfig) (*Client, error) + +// Closes the client and releases all resources +func (c *Client) Close() error + +// Creates a non-streaming chat completion +func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) + +// Creates a streaming chat completion +func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error) +``` + +### Request Types + +- `ChatCompletionRequest`: Main request type for chat completions + - Model, Messages, Stream, Temperature, TopP, MaxCompletionTokens, Tools, etc. +- `ChatMessage`: Individual message in a conversation + - Role, Content +- `Tool`: Tool/function definition for function calling + - Type, Function (name, description, parameters) + +### Response Types + +- `ChatCompletionResponse`: Non-streaming response + - ID, Model, Created, Choices, Usage +- `ChatCompletionStreamResponse`: Streaming response chunk + - Same structure as above but for incremental updates +- `Message`: Complete message with content and tool calls +- `ToolCall`: Tool call information with function and arguments +- `Usage`: Token usage statistics + - PromptTokens, CompletionTokens, TotalTokens + +## Testing + +The SDK includes comprehensive testing infrastructure with both unit and integration tests. + +### Unit Tests + +Unit tests are located in `client_test.go` and test individual components without requiring a server. + +#### Running Unit Tests + +```bash +# Run all unit tests +go test ./... + +# Run with verbose output +go test -v ./... + +# Run specific test +go test -run TestClientConfig + +# Run tests with race detector (detects concurrency issues) +go test -race ./... + +# Run with coverage analysis +go test -cover ./... + +# Generate detailed coverage report +go test -coverprofile=coverage.out ./... +go tool cover -html=coverage.out -o coverage.html +``` + +#### Unit Test Coverage + +- **Configuration validation** (`TestClientConfig`) - Validates ClientConfig requirements +- **Type structures** - Verifying all struct types work correctly +- **Response handling** - Testing response parsing and validation +- **Concurrent operations** (`TestConcurrentClientOperations`) - Thread-safety verification +- **Benchmarks** (`BenchmarkChatCompletionRequest`) - Performance measurement + +**Test Files**: +- `client_test.go` - 10 unit tests covering core functionality +- Tests cover: config validation, message types, request validation, close operations, response types, streaming, tools, concurrency, and context cancellation + +### Integration Tests + +Integration tests require a running SGLang server and test the full client-server interaction. + +#### Prerequisites + +1. Start an SGLang server: + +```bash +# Using Python (requires sglang package installed) +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-hf + +# Or using pre-built Docker image +docker run -p 20000:20000 lmsys/sglang:latest + +# Or build your own +sglang launch_server --model-path +``` + +2. Set required environment variables: + +```bash +# Set the gRPC endpoint (default: grpc://localhost:20000) +export SGL_GRPC_ENDPOINT=grpc://localhost:20000 + +# Set the tokenizer path (required) +export SGL_TOKENIZER_PATH=/path/to/tokenizer +``` + +#### Running Integration Tests + +```bash +# Run all integration tests +go test -tags=integration ./... + +# Run specific integration test +go test -tags=integration -run TestIntegrationNonStreamingCompletion + +# Run with verbose output +go test -tags=integration -v ./... + +# Run with race detector +go test -tags=integration -race ./... +``` + +#### Integration Test Coverage + +**Test File**: `integration_test.go` - 4 integration tests + +- `TestIntegrationNonStreamingCompletion` - Basic non-streaming request/response +- `TestIntegrationStreamingCompletion` - Streaming response handling +- `TestIntegrationConcurrentRequests` - Multiple simultaneous requests +- `TestIntegrationContextCancellation` - Context timeout and cancellation + +### Benchmarks + +Measure performance of SDK operations: + +```bash +# Run all benchmarks +go test -bench=. -benchmem ./... + +# Run specific benchmark +go test -bench=BenchmarkChatCompletionRequest -benchmem + +# Run for longer duration +go test -bench=. -benchtime=10s ./... +``` + +Current benchmarks: +- `BenchmarkChatCompletionRequest` - Measures request creation performance + +### CI/CD Integration + +Add to your GitHub Actions workflow: + +```yaml +- name: Run Go tests + run: | + go test -race -cover ./... + +- name: Run integration tests (on main branch) + if: github.ref == 'refs/heads/main' + env: + SGL_GRPC_ENDPOINT: grpc://localhost:20000 + SGL_TOKENIZER_PATH: /path/to/tokenizer + run: go test -tags=integration ./... +``` + +## Documentation + +### Code Documentation + +All public types and functions include comprehensive documentation: + +1. **Package-level documentation** in `client.go` with usage examples +2. **Type documentation** for all structs with field descriptions +3. **Function documentation** with: + - Purpose and behavior description + - Parameter documentation with types and constraints + - Return value documentation + - Error cases and handling + - Safety notes (for FFI functions) + - Usage examples + +### Key Documented Components + +- `Client` - Main client with thread-safety notes +- `ClientConfig` - Configuration requirements and validation rules +- `ChatCompletionRequest` - Request structure with field descriptions +- `ChatCompletionResponse` - Response structure and usage +- `ChatCompletionStreamResponse` - Streaming response format +- `Usage` - Token usage information structure +- `Tool`, `Function`, `ToolCall` - Tool call structures + +### Viewing Documentation + +Generate and view HTML documentation: + +```bash +# Install godoc (if not already installed) +go install golang.org/x/tools/cmd/godoc@latest + +# Generate and serve documentation +godoc -http=:6060 + +# Visit: http://localhost:6060/pkg/github.com/sglang/sglang-go-grpc-sdk/ +``` + +## Development + +### Building + +```bash +cd bindings/golang + +# Build the Go bindings (compiles Rust FFI library) +make build + +# Clean build +make clean && make build +``` + +### Code Quality + +Ensure code quality before committing: + +```bash +# Run Go vet (check for potential bugs) +go vet ./... + +# Format code +go fmt ./... + +# Run all tests with race detection +go test -race ./... +``` + +### Project Structure + +``` +bindings/golang/ +├── client.go # Main client implementation +├── client_test.go # Unit tests +├── integration_test.go # Integration tests +├── README.md # This file +├── Makefile # Build automation +├── Cargo.toml # Rust FFI dependencies +├── examples/ # Example programs +│ ├── simple/ # Non-streaming example +│ └── streaming/ # Streaming example +├── src/ # Rust FFI source +│ ├── client.rs # Client FFI +│ ├── stream.rs # Stream handling +│ ├── grpc_converter.rs # Response conversion +│ └── ... +└── internal/ # Internal packages + └── ffi/ # FFI bindings +``` + +## Troubleshooting + +### Connection Errors + +**Error**: `connection refused` or `failed to dial` + +**Solution**: +1. Ensure SGLang server is running: `python -m sglang.launch_server` +2. Check endpoint: `echo $SGL_GRPC_ENDPOINT` +3. Verify port is not blocked: `nc -zv localhost 20000` + +### Tokenizer Not Found + +**Error**: `tokenizer path not found` or `tokenizer configuration missing` + +**Solution**: +1. Set `SGL_TOKENIZER_PATH` environment variable +2. Verify path contains required files: `ls $SGL_TOKENIZER_PATH` +3. Files should include: `tokenizer.json`, `vocab.json`, `config.json` + +### Build Failures + +**Error**: `library 'sglang_router_rs' not found` + +**Solution**: +1. Rebuild Rust library: `cd sgl-router/bindings/golang && make build` +2. Or manually with cargo: `cd sgl-router/bindings/golang && cargo build --release` +3. Set `CARGO_BUILD_DIR` if using non-standard build location +4. Ensure Rust toolchain is installed: `rustup toolchain list` + +### Tests Hanging + +**Error**: Tests seem to hang indefinitely + +**Solution**: +1. Use timeout for hanging tests: `timeout 30s go test ./...` +2. Run with verbose output to see which test hangs: `go test -v ./...` +3. Ensure server is responsive: `grpcurl -plaintext localhost:20000 list` + +### Memory Issues + +**Error**: Out of memory during tests + +**Solution**: +```bash +# Run with memory limit for long-running tests +GODEBUG=madvdontneed=1 go test -timeout 5m ./... + +# Monitor memory during tests +watch -n1 'ps aux | grep test' +``` + +## Contributing + +When adding new features: + +1. Add comprehensive documentation to public types/functions +2. Include usage examples for complex APIs +3. Add unit tests covering happy path and error cases +4. Add integration tests if server interaction required +5. Ensure code passes `go vet` and `go test -race` +6. Update this README if adding new features + +## License + +See LICENSE file for details. + +--- + +**Need Help?** +- Check examples in `examples/` directory +- Run tests to see working code: `go test -v ./...` +- Review function documentation: `godoc` or inline comments +- Check troubleshooting section above diff --git a/sgl-router/bindings/golang/client.go b/sgl-router/bindings/golang/client.go new file mode 100644 index 000000000..110135da5 --- /dev/null +++ b/sgl-router/bindings/golang/client.go @@ -0,0 +1,510 @@ +// Package sglang provides a Go SDK for SGLang gRPC API. +// +// SGLang is a fast language model serving framework. This package provides a Go client +// library for interacting with SGLang's gRPC API, following the style of OpenAI's Go SDK. +// +// Basic usage: +// +// client, err := sglang.NewClient(sglang.ClientConfig{ +// Endpoint: "grpc://localhost:20000", +// TokenizerPath: "/path/to/tokenizer", +// }) +// if err != nil { +// log.Fatal(err) +// } +// defer client.Close() +// +// resp, err := client.CreateChatCompletion(ctx, sglang.ChatCompletionRequest{ +// Model: "default", +// Messages: []sglang.ChatMessage{ +// {Role: "user", Content: "Hello"}, +// }, +// }) +// +// For streaming responses, use CreateChatCompletionStream instead. +package sglang + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "sync" + + "github.com/sglang/sglang-go-grpc-sdk/internal/ffi" +) + +// Client is the main client for interacting with SGLang gRPC API. +// It manages the connection to the SGLang server and handles both streaming +// and non-streaming chat completions. +// +// Thread-safe: All public methods are safe for concurrent use. +type Client struct { + endpoint string + tokenizerPath string + clientHandle *ffi.SglangClientHandle + mu sync.RWMutex +} + +// ClientConfig holds configuration for creating a new client. +type ClientConfig struct { + // Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000"). + // Required field. Must include the scheme (grpc://) and port number. + Endpoint string + + // TokenizerPath is the path to the tokenizer directory containing + // tokenizer configuration files (e.g., tokenizer.json, vocab.json). + // Required field. + TokenizerPath string +} + +// NewClient creates a new SGLang client with the given configuration. +// +// The client maintains a long-lived connection to the SGLang server and should +// be reused for multiple requests. Call Close() to release resources. +// +// Returns an error if: +// - Endpoint is empty +// - TokenizerPath is empty +// - Connection to the server fails +func NewClient(config ClientConfig) (*Client, error) { + if config.Endpoint == "" { + return nil, errors.New("endpoint is required") + } + if config.TokenizerPath == "" { + return nil, errors.New("tokenizer path is required") + } + + clientHandle, err := ffi.NewClient(config.Endpoint, config.TokenizerPath) + if err != nil { + return nil, fmt.Errorf("failed to create client: %w", err) + } + + return &Client{ + endpoint: config.Endpoint, + tokenizerPath: config.TokenizerPath, + clientHandle: clientHandle, + }, nil +} + +// Close closes the client and releases all resources. +// +// After Close() is called, the client cannot be used for further requests. +// Calling Close() multiple times is safe and idempotent. +func (c *Client) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + + if c.clientHandle != nil { + c.clientHandle.Free() + c.clientHandle = nil + } + return nil +} + +// ChatCompletionRequest represents a request for chat completion. +// It follows the OpenAI API style for familiar usage. +type ChatCompletionRequest struct { + // Model specifies the model to use for completion (e.g., "default") + Model string `json:"model"` + // Messages is the list of messages in the conversation + Messages []ChatMessage `json:"messages"` + Temperature *float32 `json:"temperature,omitempty"` + TopP *float32 `json:"top_p,omitempty"` + TopK *int `json:"top_k,omitempty"` + MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"` + Stream bool `json:"stream"` + Tools []Tool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Stop interface{} `json:"stop,omitempty"` + StopTokenIDs []int `json:"stop_token_ids,omitempty"` + SkipSpecialTokens bool `json:"skip_special_tokens,omitempty"` + FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"` + PresencePenalty *float32 `json:"presence_penalty,omitempty"` + ResponseFormat *ResponseFormat `json:"response_format,omitempty"` + Seed *int `json:"seed,omitempty"` + Logprobs bool `json:"logprobs,omitempty"` + TopLogprobs *int `json:"top_logprobs,omitempty"` + User string `json:"user,omitempty"` +} + +// ChatMessage represents a single message in a chat conversation +type ChatMessage struct { + Role string `json:"role"` + Content interface{} `json:"content"` + Name string `json:"name,omitempty"` +} + +// Tool represents a tool/function that can be called +type Tool struct { + Type string `json:"type"` + Function Function `json:"function"` +} + +// Function represents a function definition +type Function struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Parameters map[string]interface{} `json:"parameters"` +} + +// ResponseFormat represents the response format +type ResponseFormat struct { + Type string `json:"type"` +} + +// ChatCompletionResponse represents a non-streaming chat completion response +type ChatCompletionResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + SystemFingerprint string `json:"system_fingerprint,omitempty"` + Choices []Choice `json:"choices"` + Usage Usage `json:"usage"` +} + +// Choice represents a choice in the completion response +type Choice struct { + Index int `json:"index"` + Message Message `json:"message"` + FinishReason string `json:"finish_reason"` +} + +// Message represents a message in the response +type Message struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` +} + +// ToolCall represents a tool call in the response +type ToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function FunctionCall `json:"function"` +} + +// FunctionCall represents a function call +type FunctionCall struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +// Usage represents token usage information +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// ChatCompletionStreamResponse represents a streaming chat completion response +type ChatCompletionStreamResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + SystemFingerprint string `json:"system_fingerprint,omitempty"` + Choices []StreamChoice `json:"choices"` + Usage *Usage `json:"usage,omitempty"` +} + +// StreamChoice represents a choice in a streaming response +type StreamChoice struct { + Index int `json:"index"` + Delta MessageDelta `json:"delta"` + FinishReason string `json:"finish_reason,omitempty"` +} + +// MessageDelta represents incremental message updates +type MessageDelta struct { + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` +} + +// CreateChatCompletion creates a non-streaming chat completion with context support. +// +// Context Support: +// The ctx parameter is fully supported for cancellation and timeouts: +// - If ctx is cancelled, the request will be interrupted on the next stream.Recv() call +// - If ctx times out, the request will return context.DeadlineExceeded +// +// Example with timeout: +// +// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) +// defer cancel() +// resp, err := client.CreateChatCompletion(ctx, req) +// +// Note: Internally, this creates a stream and collects all chunks, +// so context monitoring happens at the chunk level. +func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) { + // For non-streaming, we'll collect all chunks and return the final response + req.Stream = true // We still use streaming internally, but collect all chunks + + // Prepare request: if Tools is empty, set to nil for proper JSON serialization + if len(req.Tools) == 0 { + req.Tools = nil + } + + stream, err := c.CreateChatCompletionStream(ctx, req) + if err != nil { + return nil, err + } + defer stream.Close() + + var fullContent strings.Builder + var fullToolCalls []ToolCall + var finishReason string + var usage Usage + var responseID string + var created int64 + var model string + var systemFingerprint string + + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + + if chunk.ID != "" { + responseID = chunk.ID + } + if chunk.Created > 0 { + created = chunk.Created + } + if chunk.Model != "" { + model = chunk.Model + } + if chunk.SystemFingerprint != "" { + systemFingerprint = chunk.SystemFingerprint + } + + for _, choice := range chunk.Choices { + if choice.Delta.Content != "" { + fullContent.WriteString(choice.Delta.Content) + } + if len(choice.Delta.ToolCalls) > 0 { + fullToolCalls = append(fullToolCalls, choice.Delta.ToolCalls...) + } + // Always update finish_reason if present (even if empty string, but should not be empty) + // The last chunk (Complete message) should have finish_reason set + if choice.FinishReason != "" { + finishReason = choice.FinishReason + } + } + + // Extract usage from chunk if available (usually in the last chunk) + // Always update usage if present, as the last chunk should have the final usage + if chunk.Usage != nil { + usage = *chunk.Usage + } + } + + // Build final response + message := Message{ + Role: "assistant", + Content: fullContent.String(), + } + if len(fullToolCalls) > 0 { + message.ToolCalls = fullToolCalls + } + + // Ensure finish_reason is set (defensive check) + // If finish_reason is still empty, default to "stop" + if finishReason == "" { + finishReason = "stop" + } + + return &ChatCompletionResponse{ + ID: responseID, + Object: "chat.completion", + Created: created, + Model: model, + SystemFingerprint: systemFingerprint, + Choices: []Choice{ + { + Index: 0, + Message: message, + FinishReason: finishReason, + }, + }, + Usage: usage, + }, nil +} + +// ChatCompletionStream represents a streaming chat completion +type ChatCompletionStream struct { + stream *ffi.SglangStreamHandle + mu sync.Mutex + done bool // Track if stream has been marked as done + ctx context.Context // Context for cancellation support + cancel context.CancelFunc // Cancel function to stop monitoring goroutine + closed chan struct{} // Signal when stream is closed +} + +// Recv receives the next chunk from the stream. +// +// Supports context cancellation: if the context passed to CreateChatCompletionStream +// is cancelled, Recv will return context.Canceled error on the next call. +func (s *ChatCompletionStream) Recv() (*ChatCompletionStreamResponse, error) { + s.mu.Lock() + defer s.mu.Unlock() + + // Check if context was cancelled + select { + case <-s.ctx.Done(): + return nil, s.ctx.Err() // Returns context.Canceled or context.DeadlineExceeded + default: + } + + if s.stream == nil { + return nil, io.EOF + } + + // If stream was already marked as done, immediately return EOF + // This prevents calling ReadNext() again after isDone=1 + if s.done { + return nil, io.EOF + } + + // Loop to handle empty responses (Ok(None) from Rust) + // Keep reading until we get actual data or stream ends + for { + responseJSON, isDone, err := s.stream.ReadNext() + if err != nil { + return nil, err + } + + // Mark stream as done if ReadNext indicates completion + if isDone { + s.done = true + } + + // If we have a response, parse and return it + if responseJSON != "" { + var response ChatCompletionStreamResponse + if err := json.Unmarshal([]byte(responseJSON), &response); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + return &response, nil + } + + // If stream is done but no response, return EOF + if isDone { + return nil, io.EOF + } + + // Empty response and stream not done - loop to read next chunk + // This handles Ok(None) cases where Rust returns no data but stream continues + } +} + +// Close closes the stream and cancels any pending operations. +func (s *ChatCompletionStream) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + + // Cancel the context to signal the monitoring goroutine to stop + if s.cancel != nil { + s.cancel() + } + + // Signal that stream is closed + select { + case <-s.closed: + // Already closed + default: + close(s.closed) + } + + // Free the stream to mark it as completed + // This prevents AbortOnDropStream from sending abort when dropped + if s.stream != nil { + s.stream.Free() + s.stream = nil + } + return nil +} + +// CreateChatCompletionStream creates a streaming chat completion with context cancellation support. +// +// Context Support: +// The ctx parameter is now fully supported for cancellation and timeouts: +// - If ctx is cancelled, stream.Recv() will return context.Canceled on the next call +// - If ctx times out (WithTimeout), stream.Recv() will return context.DeadlineExceeded +// - Calling stream.Close() also cancels the context +// +// Example with timeout: +// +// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) +// defer cancel() +// stream, err := client.CreateChatCompletionStream(ctx, req) +// // Stream will auto-close if 30 seconds elapse +// +// Example with cancellation: +// +// ctx, cancel := context.WithCancel(context.Background()) +// stream, err := client.CreateChatCompletionStream(ctx, req) +// go func() { +// time.Sleep(5*time.Second) +// cancel() // Cancel after 5 seconds +// }() +func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error) { + c.mu.RLock() + defer c.mu.RUnlock() + + if c.clientHandle == nil { + return nil, errors.New("client is closed") + } + + // Marshal request to JSON, then ensure tools field is always present. + // Due to omitempty tag, empty Tools slice will be omitted from JSON. + // We need to ensure tools field is always present as [] when empty (not omitted), + // matching the behavior of complete_sdk example. + reqJSON, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + // Unmarshal into map and ensure tools field is present + var reqMap map[string]interface{} + if err := json.Unmarshal(reqJSON, &reqMap); err != nil { + return nil, fmt.Errorf("failed to unmarshal request to map: %w", err) + } + + // Add empty tools array if not present + if _, exists := reqMap["tools"]; !exists { + reqMap["tools"] = []interface{}{} + } + + // Marshal back to JSON + reqJSON, err = json.Marshal(reqMap) + if err != nil { + return nil, fmt.Errorf("failed to marshal request map to JSON: %w", err) + } + + // Create stream + streamHandle, err := c.clientHandle.ChatCompletionStream(string(reqJSON)) + if err != nil { + return nil, fmt.Errorf("failed to create stream: %w", err) + } + + // Create a child context from the provided context for cancellation support + streamCtx, cancel := context.WithCancel(ctx) + + stream := &ChatCompletionStream{ + stream: streamHandle, + ctx: streamCtx, + cancel: cancel, + closed: make(chan struct{}), + } + + return stream, nil +} diff --git a/sgl-router/bindings/golang/client_test.go b/sgl-router/bindings/golang/client_test.go new file mode 100644 index 000000000..9b7736691 --- /dev/null +++ b/sgl-router/bindings/golang/client_test.go @@ -0,0 +1,325 @@ +package sglang + +import ( + "context" + "testing" +) + +// TestClientConfig tests ClientConfig validation +func TestClientConfig(t *testing.T) { + tests := []struct { + name string + config ClientConfig + wantErr bool + }{ + { + name: "valid config", + config: ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + }, + wantErr: false, + }, + { + name: "missing endpoint", + config: ClientConfig{ + Endpoint: "", + TokenizerPath: "/path/to/tokenizer", + }, + wantErr: true, + }, + { + name: "missing tokenizer path", + config: ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "", + }, + wantErr: true, + }, + { + name: "both missing", + config: ClientConfig{ + Endpoint: "", + TokenizerPath: "", + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := NewClient(tt.config) + if (err != nil) != tt.wantErr { + t.Errorf("NewClient() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +// TestChatMessageTypes tests ChatMessage struct and its variants +func TestChatMessageTypes(t *testing.T) { + msg := ChatMessage{ + Role: "user", + Content: "Hello", + } + + if msg.Role != "user" { + t.Errorf("Expected role 'user', got '%s'", msg.Role) + } + if msg.Content != "Hello" { + t.Errorf("Expected content 'Hello', got '%s'", msg.Content) + } +} + +// TestChatCompletionRequestValidation tests ChatCompletionRequest validation +func TestChatCompletionRequestValidation(t *testing.T) { + // Test valid request + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "test"}, + }, + Stream: false, + } + + if req.Model == "" { + t.Error("Expected model to be set") + } + + if len(req.Messages) == 0 { + t.Error("Expected messages to be non-empty") + } + + if req.Messages[0].Role != "user" { + t.Errorf("Expected first message role 'user', got '%s'", req.Messages[0].Role) + } +} + +// TestClientClose tests that Close can be called multiple times safely +func TestClientClose(t *testing.T) { + // Create a mock client (note: in real tests, you might want to skip this + // if it requires actual server connection) + config := ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + } + + // Skip if connection fails (expected in unit test environment) + client, err := NewClient(config) + if err != nil { + t.Skip("Skipping client close test: server not available") + } + + // First close should succeed + if err := client.Close(); err != nil { + t.Errorf("First Close() failed: %v", err) + } + + // Second close should also succeed (idempotent) + if err := client.Close(); err != nil { + t.Errorf("Second Close() failed: %v", err) + } +} + +// TestChatCompletionResponseTypes tests response type structures +func TestChatCompletionResponseTypes(t *testing.T) { + resp := ChatCompletionResponse{ + ID: "test-id", + Model: "default", + Created: 1234567890, + Choices: []Choice{ + { + Message: Message{ + Role: "assistant", + Content: "Hello", + }, + FinishReason: "stop", + }, + }, + Usage: Usage{ + PromptTokens: 10, + CompletionTokens: 20, + TotalTokens: 30, + }, + } + + if resp.ID != "test-id" { + t.Errorf("Expected ID 'test-id', got '%s'", resp.ID) + } + + if len(resp.Choices) != 1 { + t.Errorf("Expected 1 choice, got %d", len(resp.Choices)) + } + + if resp.Choices[0].Message.Content != "Hello" { + t.Errorf("Expected content 'Hello', got '%s'", resp.Choices[0].Message.Content) + } + + if resp.Usage.TotalTokens != 30 { + t.Errorf("Expected total tokens 30, got %d", resp.Usage.TotalTokens) + } +} + +// TestStreamingResponseTypes tests streaming response structures +func TestStreamingResponseTypes(t *testing.T) { + chunk := ChatCompletionStreamResponse{ + ID: "stream-id", + Created: 1234567890, + Choices: []StreamChoice{ + { + Index: 0, + Delta: MessageDelta{ + Content: "Hello", + }, + FinishReason: "", + }, + }, + } + + if chunk.ID != "stream-id" { + t.Errorf("Expected ID 'stream-id', got '%s'", chunk.ID) + } + + if len(chunk.Choices) == 0 { + t.Error("Expected at least one choice") + } + + if chunk.Choices[0].Delta.Content != "Hello" { + t.Errorf("Expected delta content 'Hello', got '%s'", chunk.Choices[0].Delta.Content) + } +} + +// TestToolCallStructure tests Tool and ToolCall structures +func TestToolCallStructure(t *testing.T) { + tool := Tool{ + Type: "function", + Function: Function{ + Name: "get_weather", + Description: "Get the weather", + Parameters: map[string]interface{}{ + "location": "string", + }, + }, + } + + if tool.Type != "function" { + t.Errorf("Expected tool type 'function', got '%s'", tool.Type) + } + + if tool.Function.Name != "get_weather" { + t.Errorf("Expected function name 'get_weather', got '%s'", tool.Function.Name) + } + + toolCall := ToolCall{ + ID: "call-123", + Type: "function", + Function: FunctionCall{ + Name: "get_weather", + Arguments: `{"location": "San Francisco"}`, + }, + } + + if toolCall.ID != "call-123" { + t.Errorf("Expected tool call ID 'call-123', got '%s'", toolCall.ID) + } +} + +// TestConcurrentClientOperations tests thread safety +// This is a basic test that just verifies concurrent calls don't panic +func TestConcurrentClientOperations(t *testing.T) { + config := ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + } + + client, err := NewClient(config) + if err != nil { + t.Skip("Skipping concurrent operations test: server not available") + } + defer client.Close() + + // Try concurrent Close calls (should not panic or race) + done := make(chan bool, 2) + + go func() { + client.Close() + done <- true + }() + + go func() { + client.Close() + done <- true + }() + + <-done + <-done +} + +// BenchmarkChatCompletionRequest benchmarks request creation +func BenchmarkChatCompletionRequest(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "test message"}, + }, + Stream: false, + Temperature: floatPtr(0.7), + MaxCompletionTokens: intPtr(100), + } + } +} + +// Helper functions for benchmarks +func floatPtr(f float32) *float32 { + return &f +} + +func intPtr(i int) *int { + return &i +} + +// TestContextCancellation tests that cancelled context is handled gracefully. +// +// NOTE: Currently, the FFI layer is blocking and doesn't actively monitor context cancellation. +// This test verifies that the client at least returns an error rather than panicking or +// hanging indefinitely when a pre-cancelled context is passed. +// +// Future: When FFI supports context cancellation (via signals or async operations), +// this test should be updated to assert that the error is context.Canceled or wrapped +// context cancellation error. +func TestContextCancellation(t *testing.T) { + config := ClientConfig{ + Endpoint: "grpc://localhost:20000", + TokenizerPath: "/path/to/tokenizer", + } + + client, err := NewClient(config) + if err != nil { + t.Skip("Skipping context cancellation test: server not available") + } + defer client.Close() + + // Create a pre-cancelled context + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "test"}, + }, + } + + // Attempt request with cancelled context + // Since FFI is blocking, we expect either: + // 1. An error from the server/network + // 2. The call to complete normally (FFI doesn't check context) + // What we DON'T expect is a panic or indefinite hang + _, err = client.CreateChatCompletion(ctx, req) + if err != nil { + t.Logf("Request with cancelled context returned error: %v", err) + } else { + t.Logf("Request with cancelled context completed (FFI may not support context cancellation)") + } +} diff --git a/sgl-router/bindings/golang/examples/simple/main.go b/sgl-router/bindings/golang/examples/simple/main.go new file mode 100644 index 000000000..9983362c1 --- /dev/null +++ b/sgl-router/bindings/golang/examples/simple/main.go @@ -0,0 +1,85 @@ +// Simple example demonstrating basic usage of SGLang Go SDK +package main + +import ( + "context" + "fmt" + "log" + "os" + + "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 chat completion request + req := sglang.ChatCompletionRequest{ + Model: "default", + Messages: []sglang.ChatMessage{ + { + Role: "system", + Content: "You are a helpful assistant.", + }, + { + Role: "user", + Content: "写一首歌关于夏天", + }, + }, + Stream: false, + Temperature: float32Ptr(0.7), + MaxCompletionTokens: intPtr(200), + SkipSpecialTokens: true, + Tools: nil, // Use nil instead of empty slice to avoid template errors + } + + // Create completion + ctx := context.Background() + resp, err := client.CreateChatCompletion(ctx, req) + if err != nil { + log.Fatalf("Failed to create completion: %v", err) + } + + // Print response + fmt.Println("=== Response ===") + fmt.Printf("ID: %s\n", resp.ID) + fmt.Printf("Model: %s\n", resp.Model) + fmt.Printf("Created: %d\n", resp.Created) + fmt.Println("\nContent:") + for _, choice := range resp.Choices { + fmt.Println(choice.Message.Content) + } + fmt.Printf("\nFinish Reason: %s\n", resp.Choices[0].FinishReason) + fmt.Printf("\nUsage: Prompt=%d, Completion=%d, Total=%d\n", + resp.Usage.PromptTokens, + resp.Usage.CompletionTokens, + resp.Usage.TotalTokens, + ) +} + +func float32Ptr(f float32) *float32 { + return &f +} + +func intPtr(i int) *int { + return &i +} diff --git a/sgl-router/bindings/golang/examples/simple/run.sh b/sgl-router/bindings/golang/examples/simple/run.sh new file mode 100755 index 000000000..9153f2e2c --- /dev/null +++ b/sgl-router/bindings/golang/examples/simple/run.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +# Simple example runner +# Usage: ./run.sh [tokenizer_path] [endpoint] + +# Set library path for Rust FFI library +# The library should be in ./lib directory (created by 'make lib') +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib" + +# Check if lib directory exists +if [ ! -d "$LIB_DIR" ]; then + echo "Error: Library directory not found at $LIB_DIR" + echo "Please run 'make lib' first to build and export the library" + exit 1 +fi + +# Get Python LDFLAGS (needed for Rust FFI that depends on Python) +PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "") + +# Set CGO_LDFLAGS to link with the Rust library +export CGO_LDFLAGS="-L${LIB_DIR} -lsglang_router_rs ${PYTHON_LDFLAGS} -ldl" + +# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH +if [[ "$OSTYPE" == "darwin"* ]]; then + export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}" +else + export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}" +fi + +# Default configuration (can be overridden by environment variables or command line arguments) +# Tokenizer path: ../tokenizer (relative to this script) +DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}" +DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}" + +TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}" +ENDPOINT="${2:-${DEFAULT_ENDPOINT}}" + +echo "Running simple example..." +echo "Library path: ${LIB_DIR}" +echo "Tokenizer: $TOKENIZER_PATH" +echo "Endpoint: $ENDPOINT" +echo "" + +cd "$(dirname "${BASH_SOURCE[0]}")" +SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go diff --git a/sgl-router/bindings/golang/examples/streaming/main.go b/sgl-router/bindings/golang/examples/streaming/main.go new file mode 100644 index 000000000..fa27909c4 --- /dev/null +++ b/sgl-router/bindings/golang/examples/streaming/main.go @@ -0,0 +1,125 @@ +// 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 +} diff --git a/sgl-router/bindings/golang/examples/streaming/run.sh b/sgl-router/bindings/golang/examples/streaming/run.sh new file mode 100755 index 000000000..49911cc0c --- /dev/null +++ b/sgl-router/bindings/golang/examples/streaming/run.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +# Streaming example runner +# Usage: ./run.sh [tokenizer_path] [endpoint] + +# Set library path for Rust FFI library +# The library should be in ./lib directory (created by 'make lib') +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib" + +# Check if lib directory exists +if [ ! -d "$LIB_DIR" ]; then + echo "Error: Library directory not found at $LIB_DIR" + echo "Please run 'make lib' first to build and export the library" + exit 1 +fi + +# Get Python LDFLAGS (needed for Rust FFI that depends on Python) +PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "") + +# Set CGO_LDFLAGS to link with the Rust library +export CGO_LDFLAGS="-L${LIB_DIR} -lsglang_router_rs ${PYTHON_LDFLAGS} -ldl" + +# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH +if [[ "$OSTYPE" == "darwin"* ]]; then + export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}" +else + export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}" +fi + +# Default configuration (can be overridden by environment variables or command line arguments) +# Tokenizer path: ../tokenizer (relative to this script) +DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}" +DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}" + +TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}" +ENDPOINT="${2:-${DEFAULT_ENDPOINT}}" + +echo "Running streaming example..." +echo "Library path: ${LIB_DIR}" +echo "Tokenizer: $TOKENIZER_PATH" +echo "Endpoint: $ENDPOINT" +echo "" + +cd "$(dirname "${BASH_SOURCE[0]}")" +SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go diff --git a/sgl-router/bindings/golang/go.mod b/sgl-router/bindings/golang/go.mod new file mode 100644 index 000000000..9b820fcf3 --- /dev/null +++ b/sgl-router/bindings/golang/go.mod @@ -0,0 +1,3 @@ +module github.com/sglang/sglang-go-grpc-sdk + +go 1.21 diff --git a/sgl-router/bindings/golang/integration_test.go b/sgl-router/bindings/golang/integration_test.go new file mode 100644 index 000000000..5ed9713f5 --- /dev/null +++ b/sgl-router/bindings/golang/integration_test.go @@ -0,0 +1,228 @@ +//go:build integration +// +build integration + +// integration_test.go contains integration tests that require a running SGLang server +// +// To run these tests: +// 1. Start an SGLang server: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-hf +// 2. Run: go test -tags=integration -run TestIntegration + +package sglang + +import ( + "context" + "io" + "os" + "testing" + "time" +) + +// getTestConfig returns test configuration from environment or defaults +func getTestConfig(t *testing.T) ClientConfig { + endpoint := os.Getenv("SGL_GRPC_ENDPOINT") + if endpoint == "" { + endpoint = "grpc://localhost:20000" + } + + tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH") + if tokenizerPath == "" { + t.Skip("SGL_TOKENIZER_PATH not set") + } + + return ClientConfig{ + Endpoint: endpoint, + TokenizerPath: tokenizerPath, + } +} + +// TestIntegrationNonStreamingCompletion tests non-streaming chat completion +func TestIntegrationNonStreamingCompletion(t *testing.T) { + config := getTestConfig(t) + + client, err := NewClient(config) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer client.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "Say 'Hello, World!' only"}, + }, + Stream: false, + Temperature: float32Ptr(0.0), + MaxCompletionTokens: intPtr(50), + } + + resp, err := client.CreateChatCompletion(ctx, req) + if err != nil { + t.Fatalf("CreateChatCompletion failed: %v", err) + } + + if resp.ID == "" { + t.Error("Response ID is empty") + } + + if len(resp.Choices) == 0 { + t.Error("Response has no choices") + } + + if resp.Choices[0].Message.Content == "" { + t.Error("Response content is empty") + } + + if resp.Usage == nil || resp.Usage.TotalTokens == 0 { + t.Error("Usage information is missing or invalid") + } + + t.Logf("Response: %s", resp.Choices[0].Message.Content) + t.Logf("Usage: %+v", resp.Usage) +} + +// TestIntegrationStreamingCompletion tests streaming chat completion +func TestIntegrationStreamingCompletion(t *testing.T) { + config := getTestConfig(t) + + client, err := NewClient(config) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer client.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "Count from 1 to 5"}, + }, + Stream: true, + Temperature: float32Ptr(0.0), + MaxCompletionTokens: intPtr(100), + } + + stream, err := client.CreateChatCompletionStream(ctx, req) + if err != nil { + t.Fatalf("CreateChatCompletionStream failed: %v", err) + } + defer stream.Close() + + chunkCount := 0 + totalContent := "" + + for { + chunk, err := stream.Recv() + if err == io.EOF { + // io.EOF is expected at end of stream + break + } + if err != nil { + t.Fatalf("Stream error: %v", err) + } + + chunkCount++ + + for _, choice := range chunk.Choices { + if choice.Delta.Content != "" { + totalContent += choice.Delta.Content + } + } + } + + if chunkCount == 0 { + t.Error("Received no chunks from stream") + } + + if totalContent == "" { + t.Error("Received no content from stream") + } + + t.Logf("Received %d chunks with content: %s", chunkCount, totalContent) +} + +// TestIntegrationConcurrentRequests tests multiple concurrent requests +func TestIntegrationConcurrentRequests(t *testing.T) { + config := getTestConfig(t) + + client, err := NewClient(config) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer client.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + numRequests := 3 + done := make(chan error, numRequests) + + for i := 0; i < numRequests; i++ { + go func(idx int) { + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "Say 'test'"}, + }, + Stream: false, + MaxCompletionTokens: intPtr(50), + } + + _, err := client.CreateChatCompletion(ctx, req) + done <- err + }(i) + } + + // Collect results + for i := 0; i < numRequests; i++ { + if err := <-done; err != nil { + t.Errorf("Request %d failed: %v", i, err) + } + } + + t.Logf("All %d concurrent requests completed successfully", numRequests) +} + +// TestIntegrationContextCancellation tests that context cancellation is handled +func TestIntegrationContextCancellation(t *testing.T) { + config := getTestConfig(t) + + client, err := NewClient(config) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer client.Close() + + // Create a context that cancels immediately + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + req := ChatCompletionRequest{ + Model: "default", + Messages: []ChatMessage{ + {Role: "user", Content: "test"}, + }, + Stream: false, + } + + // Should handle cancelled context gracefully + _, err = client.CreateChatCompletion(ctx, req) + if err == nil { + t.Error("Expected error from cancelled context") + } + + t.Logf("Cancelled context handled: %v", err) +} + +// Helper functions +func float32Ptr(f float32) *float32 { + return &f +} + +func intPtr(i int) *int { + return &i +} diff --git a/sgl-router/bindings/golang/internal/ffi/client.go b/sgl-router/bindings/golang/internal/ffi/client.go new file mode 100644 index 000000000..13439e3c9 --- /dev/null +++ b/sgl-router/bindings/golang/internal/ffi/client.go @@ -0,0 +1,228 @@ +// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface). +// +// This package wraps the Rust FFI layer of SGLang, providing low-level access to: +// - Client creation and connection management +// - Chat completion streaming +// - Stream reading and response conversion +// - Memory management for C strings +// +// Internal use only: This package is intended for internal use by the sglang package. +// End users should use the public sglang package instead. +package ffi + +/* +#cgo LDFLAGS: -lsglang_router_rs -ldl +#include +#include + +// Error codes +typedef enum { + SGL_ERROR_SUCCESS = 0, + SGL_ERROR_INVALID_ARGUMENT = 1, + SGL_ERROR_TOKENIZATION_ERROR = 2, + SGL_ERROR_PARSING_ERROR = 3, + SGL_ERROR_MEMORY_ERROR = 4, + SGL_ERROR_UNKNOWN = 99 +} SglErrorCode; + +// Opaque handles +typedef void* SglangClientHandle; +typedef void* SglangStreamHandle; + +// Client SDK functions +SglangClientHandle* sgl_client_create(const char* endpoint, const char* tokenizer_path, char** error_out); +void sgl_client_free(SglangClientHandle* handle); +SglErrorCode sgl_client_chat_completion_stream(SglangClientHandle* client_handle, const char* request_json, SglangStreamHandle** stream_handle_out, char** error_out); +SglErrorCode sgl_stream_read_next(SglangStreamHandle* stream_handle, char** response_json_out, int* is_done_out, char** error_out); +void sgl_stream_free(SglangStreamHandle* handle); +void sgl_free_string(char* s); +*/ +import "C" + +import ( + "fmt" + "unsafe" +) + +// ErrorCode represents FFI error codes returned by Rust functions. +// +// These codes indicate the result of FFI operations. Use Error() to get a human-readable +// error message. +type ErrorCode int + +const ( + // ErrorSuccess indicates the operation completed successfully + ErrorSuccess ErrorCode = 0 + // ErrorInvalidArgument indicates invalid arguments were passed to the FFI function + ErrorInvalidArgument ErrorCode = 1 + // ErrorTokenizationError indicates an error during tokenization + ErrorTokenizationError ErrorCode = 2 + // ErrorParsingError indicates an error parsing the response or request + ErrorParsingError ErrorCode = 3 + // ErrorMemoryError indicates a memory allocation error + ErrorMemoryError ErrorCode = 4 + // ErrorUnknown indicates an unclassified error + ErrorUnknown ErrorCode = 99 +) + +// Error implements the error interface for ErrorCode. +func (e ErrorCode) Error() string { + switch e { + case ErrorSuccess: + return "success" + case ErrorInvalidArgument: + return "invalid argument" + case ErrorTokenizationError: + return "tokenization error" + case ErrorParsingError: + return "parsing error" + case ErrorMemoryError: + return "memory error" + case ErrorUnknown: + return "unknown error" + default: + return fmt.Sprintf("unknown error code: %d", e) + } +} + +// SglangClientHandle wraps the Rust client SDK FFI handle. +// +// This struct maintains a connection to the SGLang gRPC server and is used +// to create streams and manage the underlying Rust client resources. +type SglangClientHandle struct { + handle *C.SglangClientHandle +} + +// NewClient creates a new SGLang client handle via FFI. +// +// This function initializes the Rust client with the given endpoint and tokenizer path. +// +// Parameters: +// - endpoint: gRPC endpoint URL (e.g., "grpc://localhost:20000") +// - tokenizerPath: Path to tokenizer directory +// +// Returns: +// - *SglangClientHandle: A new client handle +// - error: An error if client creation failed +func NewClient(endpoint, tokenizerPath string) (*SglangClientHandle, error) { + cEndpoint := C.CString(endpoint) + defer C.free(unsafe.Pointer(cEndpoint)) + + cTokenizerPath := C.CString(tokenizerPath) + defer C.free(unsafe.Pointer(cTokenizerPath)) + + var errorPtr *C.char + handle := C.sgl_client_create(cEndpoint, cTokenizerPath, &errorPtr) + + if handle == nil { + errorMsg := "" + if errorPtr != nil { + errorMsg = C.GoString(errorPtr) + C.sgl_free_string(errorPtr) + } + if errorMsg == "" { + errorMsg = "failed to create client" + } + return nil, fmt.Errorf("%s", errorMsg) + } + + return &SglangClientHandle{handle: handle}, nil +} + +// Free releases the client handle +func (h *SglangClientHandle) Free() { + if h.handle != nil { + C.sgl_client_free(h.handle) + h.handle = nil + } +} + +// ChatCompletionStream creates a streaming chat completion request +func (h *SglangClientHandle) ChatCompletionStream(requestJSON string) (*SglangStreamHandle, error) { + if h.handle == nil { + return nil, fmt.Errorf("client handle is nil") + } + + cRequestJSON := C.CString(requestJSON) + defer C.free(unsafe.Pointer(cRequestJSON)) + + var streamHandle *C.SglangStreamHandle + var errorPtr *C.char + + result := C.sgl_client_chat_completion_stream( + h.handle, + cRequestJSON, + &streamHandle, + &errorPtr, + ) + + if ErrorCode(result) != ErrorSuccess { + errorMsg := "" + if errorPtr != nil { + errorMsg = C.GoString(errorPtr) + C.sgl_free_string(errorPtr) + } + if errorMsg == "" { + errorMsg = fmt.Sprintf("error code %d", result) + } + return nil, fmt.Errorf("%s", errorMsg) + } + + if streamHandle == nil { + return nil, fmt.Errorf("stream handle is nil") + } + + return &SglangStreamHandle{handle: streamHandle}, nil +} + +// SglangStreamHandle wraps the Rust stream FFI handle +type SglangStreamHandle struct { + handle *C.SglangStreamHandle +} + +// ReadNext reads the next chunk from the stream +// Returns: (responseJSON, isDone, error) +func (h *SglangStreamHandle) ReadNext() (string, bool, error) { + if h.handle == nil { + return "", true, fmt.Errorf("stream handle is nil") + } + + var responseJSON *C.char + var isDone C.int + var errorPtr *C.char + + result := C.sgl_stream_read_next( + h.handle, + &responseJSON, + &isDone, + &errorPtr, + ) + + if ErrorCode(result) != ErrorSuccess { + errorMsg := "" + if errorPtr != nil { + errorMsg = C.GoString(errorPtr) + C.sgl_free_string(errorPtr) + } + if errorMsg == "" { + errorMsg = fmt.Sprintf("error code %d", result) + } + return "", isDone == 1, fmt.Errorf("%s", errorMsg) + } + + responseStr := "" + if responseJSON != nil { + responseStr = C.GoString(responseJSON) + C.sgl_free_string(responseJSON) + } + + return responseStr, isDone == 1, nil +} + +// Free releases the stream handle +func (h *SglangStreamHandle) Free() { + if h.handle != nil { + C.sgl_stream_free(h.handle) + h.handle = nil + } +} diff --git a/sgl-router/bindings/golang/src/client.rs b/sgl-router/bindings/golang/src/client.rs new file mode 100644 index 000000000..87543ecec --- /dev/null +++ b/sgl-router/bindings/golang/src/client.rs @@ -0,0 +1,279 @@ +//! Client SDK FFI functions + +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char}; +use std::ptr; +use std::sync::Arc; +use tokio::runtime::Runtime; +use once_cell::sync::Lazy; +use uuid::Uuid; + +use sglang_router::tokenizer::create_tokenizer_from_file; +use sglang_router::tokenizer::traits::Tokenizer; +use sglang_router::grpc_client::sglang_scheduler::SglangSchedulerClient; +use sglang_router::protocols::chat::ChatCompletionRequest; +use sglang_router::routers::grpc::utils::{process_chat_messages, generate_tool_constraints}; + +use super::error::{SglErrorCode, set_error_message}; +use super::grpc_converter::sgl_grpc_response_converter_create; +use super::tokenizer::TokenizerHandle; +use super::stream::SglangStreamHandle; + +/// Global tokio runtime for async operations +static RUNTIME: Lazy = Lazy::new(|| { + Runtime::new().expect("Failed to create tokio runtime for client FFI") +}); + +/// Handle for complete client SDK (gRPC client + tokenizer) +/// This handle manages the connection to sglang and provides a complete SDK interface +pub struct SglangClientHandle { + pub(crate) client: Arc, + pub(crate) tokenizer: Arc, +} + +/// Handle for streaming request (includes prompt token count) +#[allow(dead_code)] +pub struct StreamRequestState { + pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request +} + +/// Create a new SGLang client handle +/// +/// # Arguments +/// * `endpoint` - gRPC endpoint (e.g., "grpc://localhost:20000") +/// * `tokenizer_path` - Path to tokenizer directory +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * Pointer to SglangClientHandle on success, null on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_client_create( + endpoint: *const c_char, + tokenizer_path: *const c_char, + error_out: *mut *mut c_char, +) -> *mut SglangClientHandle { + if endpoint.is_null() || tokenizer_path.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return ptr::null_mut(); + } + + let endpoint_str = match CStr::from_ptr(endpoint).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in endpoint"); + return ptr::null_mut(); + } + }; + + let tokenizer_path_str = match CStr::from_ptr(tokenizer_path).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in tokenizer_path"); + return ptr::null_mut(); + } + }; + + // Create tokenizer + let tokenizer = match create_tokenizer_from_file(tokenizer_path_str) { + Ok(t) => t, + Err(e) => { + set_error_message(error_out, &format!("Failed to create tokenizer: {}", e)); + return ptr::null_mut(); + } + }; + + // Create gRPC client + let client = match RUNTIME.block_on(async { + SglangSchedulerClient::connect(endpoint_str).await + }) { + Ok(c) => Arc::new(c), + Err(e) => { + set_error_message(error_out, &format!("Failed to connect to endpoint: {}", e)); + return ptr::null_mut(); + } + }; + + Box::into_raw(Box::new(SglangClientHandle { + client, + tokenizer, + })) +} + +/// Free a client handle +#[no_mangle] +pub unsafe extern "C" fn sgl_client_free(handle: *mut SglangClientHandle) { + if !handle.is_null() { + let _ = Box::from_raw(handle); + } +} + +/// Send a chat completion request and start streaming +/// +/// # Arguments +/// * `client_handle` - Client handle +/// * `request_json` - OpenAI ChatCompletionRequest as JSON string +/// * `stream_handle_out` - Pointer to receive stream handle +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_client_chat_completion_stream( + client_handle: *mut SglangClientHandle, + request_json: *const c_char, + stream_handle_out: *mut *mut SglangStreamHandle, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if client_handle.is_null() || request_json.is_null() || stream_handle_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let request_str = match CStr::from_ptr(request_json).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in request_json"); + return SglErrorCode::InvalidArgument; + } + }; + + let client_ref = &*client_handle; + let client = Arc::clone(&client_ref.client); + let tokenizer = Arc::clone(&client_ref.tokenizer); + + // Parse OpenAI ChatCompletionRequest + let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) { + Ok(req) => req, + Err(e) => { + set_error_message(error_out, &format!("Failed to parse request JSON: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + // Process messages and apply chat template + let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) { + Ok(msgs) => msgs, + Err(e) => { + set_error_message(error_out, &format!("Failed to process messages: {}", e)); + return SglErrorCode::TokenizationError; + } + }; + + // Tokenize + let token_ids = match tokenizer.encode(&processed_messages.text) { + Ok(encoding) => encoding.token_ids().to_vec(), + Err(e) => { + set_error_message(error_out, &format!("Failed to tokenize: {}", e)); + return SglErrorCode::TokenizationError; + } + }; + let prompt_tokens = token_ids.len() as i32; // Save prompt token count + + // Generate tool constraints if needed + let tool_constraint = if let Some(tools) = chat_request.tools.as_ref() { + match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) { + Ok(Some((constraint_type, constraint_value))) => Some((constraint_type, constraint_value)), + Ok(None) => None, + Err(e) => { + set_error_message(error_out, &format!("Failed to generate tool constraints: {}", e)); + return SglErrorCode::ParsingError; + } + } + } else { + None + }; + + // Build GenerateRequest + let request_id = format!("chatcmpl-{}", Uuid::new_v4()); + let proto_request = match client.build_generate_request_from_chat( + request_id.clone(), + &chat_request, + processed_messages.text, + token_ids, + processed_messages.multimodal_inputs, + tool_constraint, + ) { + Ok(req) => req, + Err(e) => { + set_error_message(error_out, &format!("Failed to build generate request: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + // Send request and get stream + let stream = match RUNTIME.block_on(async { + client.generate(proto_request).await + }) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to send request: {}", e)); + return SglErrorCode::UnknownError; + } + }; + + // Create response converter + let tools_json = chat_request.tools.as_ref() + .and_then(|t| serde_json::to_string(t).ok()) + .map(|s| CString::new(s).unwrap().into_raw()); + let tool_choice_json = chat_request.tool_choice.as_ref() + .and_then(|tc| serde_json::to_string(tc).ok()) + .map(|s| CString::new(s).unwrap().into_raw()); + let stop_json = chat_request.stop.as_ref() + .and_then(|s| serde_json::to_string(s).ok()) + .map(|s| CString::new(s).unwrap().into_raw()); + let stop_token_ids_json = chat_request.stop_token_ids.as_ref() + .and_then(|ids| serde_json::to_string(ids).ok()) + .map(|s| CString::new(s).unwrap().into_raw()); + + // Create tokenizer handle for converter (we'll create a temporary one) + let tokenizer_handle = Box::into_raw(Box::new(TokenizerHandle { + tokenizer: Arc::clone(&tokenizer), + })); + + let converter = sgl_grpc_response_converter_create( + tokenizer_handle, + CString::new(chat_request.model.clone()).unwrap().as_ptr(), + CString::new(request_id.clone()).unwrap().as_ptr(), + tools_json.unwrap_or(ptr::null_mut()), + tool_choice_json.unwrap_or(ptr::null_mut()), + stop_json.unwrap_or(ptr::null_mut()), + stop_token_ids_json.unwrap_or(ptr::null_mut()), + if chat_request.skip_special_tokens { 1 } else { 0 }, + error_out, + ); + + // Free temporary tokenizer handle (converter now owns the tokenizer) + let _ = Box::from_raw(tokenizer_handle); + + if converter.is_null() { + return SglErrorCode::MemoryError; + } + + // Clean up temporary CStrings + if let Some(ptr) = tools_json { + let _ = CString::from_raw(ptr); + } + if let Some(ptr) = tool_choice_json { + let _ = CString::from_raw(ptr); + } + if let Some(ptr) = stop_json { + let _ = CString::from_raw(ptr); + } + if let Some(ptr) = stop_token_ids_json { + let _ = CString::from_raw(ptr); + } + + // Create converter handle and set initial_prompt_tokens immediately + let mut converter_handle = *Box::from_raw(converter); + converter_handle.initial_prompt_tokens = Some(prompt_tokens); + + // Create stream handle with prompt_tokens + *stream_handle_out = Box::into_raw(Box::new(SglangStreamHandle { + stream: Arc::new(tokio::sync::Mutex::new(stream)), + converter: Arc::new(tokio::sync::Mutex::new(converter_handle)), + client: Arc::clone(&client), + prompt_tokens, + })); + + SglErrorCode::Success +} diff --git a/sgl-router/bindings/golang/src/error.rs b/sgl-router/bindings/golang/src/error.rs new file mode 100644 index 000000000..3e43945ef --- /dev/null +++ b/sgl-router/bindings/golang/src/error.rs @@ -0,0 +1,50 @@ +//! Error handling for FFI functions + +use std::ffi::CString; +use std::os::raw::c_char; +use std::ptr; + +/// Error codes returned by FFI functions +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SglErrorCode { + Success = 0, + InvalidArgument = 1, + TokenizationError = 2, + ParsingError = 3, + MemoryError = 4, + UnknownError = 99, +} + +/// Helper to set error message in FFI output parameter +pub fn set_error_message(error_out: *mut *mut c_char, message: &str) { + unsafe { + if !error_out.is_null() { + if let Ok(cstr) = CString::new(message) { + *error_out = cstr.into_raw(); + } else { + *error_out = ptr::null_mut(); + } + } + } +} + +/// Helper to set error message from format string +pub fn set_error_message_fmt(error_out: *mut *mut c_char, fmt: std::fmt::Arguments) { + if !error_out.is_null() { + let msg = format!("{}", fmt); + set_error_message(error_out, &msg); + } +} + +/// Helper to clear error message +pub fn clear_error_message(error_out: *mut *mut c_char) { + unsafe { + if !error_out.is_null() { + *error_out = ptr::null_mut(); + } + } +} + +// Helper functions for error handling +// Note: Some helper functions are kept for potential future use diff --git a/sgl-router/bindings/golang/src/grpc_converter.rs b/sgl-router/bindings/golang/src/grpc_converter.rs new file mode 100644 index 000000000..c8ec7ecf3 --- /dev/null +++ b/sgl-router/bindings/golang/src/grpc_converter.rs @@ -0,0 +1,758 @@ +//! gRPC response converter FFI functions + +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char, c_int}; +use std::ptr; +use std::sync::Arc; +use std::collections::HashMap; +use serde_json::Value; +use tokio::runtime::Runtime; +use once_cell::sync::Lazy; + +use sglang_router::tokenizer::traits::Tokenizer; +use sglang_router::tokenizer::stream::DecodeStream; +use sglang_router::tool_parser::ToolParser; +use sglang_router::protocols::common::{Tool, ToolChoice, ToolChoiceValue, ToolCallDelta, FunctionCallDelta, Usage, StringOrArray}; +use sglang_router::tokenizer::stop::StopSequenceDecoder; +use sglang_router::grpc_client::sglang_proto as proto; + +use super::error::{SglErrorCode, set_error_message, clear_error_message}; +use super::tokenizer::TokenizerHandle; +use super::utils::generate_tool_call_id; + +/// Global parser factory (initialized once) +// Use the re-exported ParserFactory from tool_parser module +static PARSER_FACTORY: Lazy = Lazy::new(|| { + // ParserFactory is re-exported from tool_parser::factory, so we can use it directly + sglang_router::tool_parser::ParserFactory::default() +}); + +/// Global tokio runtime for async operations +static RUNTIME: Lazy = Lazy::new(|| { + Runtime::new().expect("Failed to create tokio runtime for gRPC converter FFI") +}); + +/// Handle for gRPC response converter (maintains state for streaming) +#[repr(C)] +pub struct GrpcResponseConverterHandle { + pub(crate) tokenizer: Arc, + pub(crate) tool_parser: Option>>>, + pub(crate) stop_decoder: Option>>, + pub(crate) model: String, + pub(crate) request_id: String, + pub(crate) created: u64, + pub(crate) system_fingerprint: Option, + pub(crate) tools: Option>, + pub(crate) tool_choice: Option, + pub(crate) history_tool_calls_count: usize, + pub(crate) stream_buffers: HashMap, // Per-index text buffers + pub(crate) decode_streams: HashMap, // Per-index incremental decoders + pub(crate) has_tool_calls: HashMap, // Track if tool calls were emitted + pub(crate) is_first_chunk: HashMap, // Track first chunk per index + pub(crate) prompt_tokens: HashMap, // Track prompt tokens per index (from chunks) + pub(crate) completion_tokens: HashMap, // Track completion tokens per index (cumulative) + pub(crate) initial_prompt_tokens: Option, // Initial prompt tokens from request (if available) + pub(crate) skip_special_tokens: bool, // Whether to skip special tokens when decoding +} + +/// Create a gRPC response converter handle +/// +/// # Arguments +/// * `tokenizer_handle` - Tokenizer handle (must be valid) +/// * `model` - Model name +/// * `request_id` - Request ID +/// * `tools_json` - Optional JSON array of tools +/// * `tool_choice_json` - Optional JSON object for tool_choice +/// * `stop` - Optional stop sequences (JSON array) +/// * `stop_token_ids` - Optional stop token IDs (JSON array) +/// * `skip_special_tokens` - Whether to skip special tokens +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * Pointer to GrpcResponseConverterHandle on success, null on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_grpc_response_converter_create( + tokenizer_handle: *mut TokenizerHandle, + model: *const c_char, + request_id: *const c_char, + tools_json: *const c_char, + tool_choice_json: *const c_char, + stop: *const c_char, + stop_token_ids: *const c_char, + skip_special_tokens: c_int, + error_out: *mut *mut c_char, +) -> *mut GrpcResponseConverterHandle { + if tokenizer_handle.is_null() || model.is_null() || request_id.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return ptr::null_mut(); + } + + let model_str = match CStr::from_ptr(model).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in model"); + return ptr::null_mut(); + } + }; + + let request_id_str = match CStr::from_ptr(request_id).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in request_id"); + return ptr::null_mut(); + } + }; + + let handle_ref = &*tokenizer_handle; + let tokenizer = Arc::clone(&handle_ref.tokenizer); + + // Parse tools if provided + let tools: Option> = if !tools_json.is_null() { + match CStr::from_ptr(tools_json).to_str() { + Ok(s) => serde_json::from_str::>(s).ok(), + Err(_) => None, + } + } else { + None + }; + + // Parse tool_choice if provided + let tool_choice: Option = if !tool_choice_json.is_null() { + match CStr::from_ptr(tool_choice_json).to_str() { + Ok(s) => serde_json::from_str::(s).ok(), + Err(_) => None, + } + } else { + None + }; + + // Parse stop sequences + let stop: Option = if !stop.is_null() { + let stop_str = match CStr::from_ptr(stop).to_str() { + Ok(s) => s, + Err(_) => return ptr::null_mut(), + }; + serde_json::from_str::(stop_str).ok() + } else { + None + }; + + // Parse stop token IDs + let stop_token_ids: Option> = if !stop_token_ids.is_null() { + let ids_str = match CStr::from_ptr(stop_token_ids).to_str() { + Ok(s) => s, + Err(_) => return ptr::null_mut(), + }; + serde_json::from_str::>(ids_str).ok() + } else { + None + }; + + // Create stop decoder if needed + let stop_decoder = if stop.is_some() || stop_token_ids.is_some() { + Some(Arc::new(tokio::sync::Mutex::new( + sglang_router::routers::grpc::utils::create_stop_decoder( + &tokenizer, + stop.as_ref(), + stop_token_ids.as_ref(), + skip_special_tokens != 0, + false, // no_stop_trim + ), + ))) + } else { + None + }; + + // Create tool parser if tools are provided + let tool_parser = if tools.is_some() { + PARSER_FACTORY.registry().create_for_model(model_str) + .map(|p| Arc::new(tokio::sync::Mutex::new(p))) + } else { + None + }; + + // Get system fingerprint from model (simplified) + let system_fingerprint = Some("fp_placeholder".to_string()); // TODO: Get actual fingerprint + + Box::into_raw(Box::new(GrpcResponseConverterHandle { + tokenizer, + tool_parser, + stop_decoder, + model: model_str.to_string(), + request_id: request_id_str.to_string(), + created: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(), + system_fingerprint, + tools, + tool_choice, + history_tool_calls_count: 0, + stream_buffers: HashMap::new(), + decode_streams: HashMap::new(), + has_tool_calls: HashMap::new(), + is_first_chunk: HashMap::new(), + prompt_tokens: HashMap::new(), + completion_tokens: HashMap::new(), + initial_prompt_tokens: None, // Will be set from stream handle + skip_special_tokens: skip_special_tokens != 0, + })) +} + +/// Convert a gRPC GenerateResponse chunk to OpenAI format +/// +/// # Arguments +/// * `handle` - Converter handle +/// * `response_json` - JSON string of proto.GenerateResponse +/// * `result_json_out` - Pointer to receive OpenAI format JSON (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_grpc_response_converter_convert_chunk( + handle: *mut GrpcResponseConverterHandle, + response_json: *const c_char, + result_json_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || response_json.is_null() || result_json_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let response_str = match CStr::from_ptr(response_json).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in response_json"); + return SglErrorCode::InvalidArgument; + } + }; + + // Parse proto.GenerateResponse from JSON + let json_value: Value = match serde_json::from_str(response_str) { + Ok(v) => v, + Err(e) => { + set_error_message(error_out, &format!("Failed to parse response JSON: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + // Build proto::GenerateResponse from JSON value + let mut proto_response = proto::GenerateResponse { + request_id: json_value.get("request_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + response: None, + }; + + // Parse the response oneof field + if let Some(chunk_json) = json_value.get("chunk") { + let chunk = proto::GenerateStreamChunk { + token_ids: chunk_json.get("token_ids") + .and_then(|v| v.as_array()) + .map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect()) + .unwrap_or_default(), + prompt_tokens: chunk_json.get("prompt_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + completion_tokens: chunk_json.get("completion_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + cached_tokens: chunk_json.get("cached_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + output_logprobs: None, + hidden_states: vec![], + input_logprobs: None, + index: 0, + }; + proto_response.response = Some(proto::generate_response::Response::Chunk(chunk)); + } else if let Some(complete_json) = json_value.get("complete") { + let complete = proto::GenerateComplete { + output_ids: complete_json.get("output_ids") + .and_then(|v| v.as_array()) + .map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect()) + .unwrap_or_default(), + finish_reason: complete_json.get("finish_reason") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + prompt_tokens: complete_json.get("prompt_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + completion_tokens: complete_json.get("completion_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + cached_tokens: complete_json.get("cached_tokens") + .and_then(|v| v.as_i64()) + .map(|n| n as i32) + .unwrap_or(0), + output_logprobs: None, + all_hidden_states: vec![], + input_logprobs: None, + matched_stop: None, + index: 0, + }; + proto_response.response = Some(proto::generate_response::Response::Complete(complete)); + } else if let Some(error_json) = json_value.get("error") { + let error = proto::GenerateError { + message: error_json.get("message") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + http_status_code: error_json.get("http_status_code") + .and_then(|v| v.as_str()) + .unwrap_or("500") + .to_string(), + details: error_json.get("details") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + }; + proto_response.response = Some(proto::generate_response::Response::Error(error)); + } else { + set_error_message(error_out, "Response JSON must contain 'chunk', 'complete', or 'error' field"); + return SglErrorCode::ParsingError; + } + + let handle_ref = &mut *handle; + let tokenizer = Arc::clone(&handle_ref.tokenizer); + let model = handle_ref.model.clone(); + let request_id = handle_ref.request_id.clone(); + let created = handle_ref.created; + let system_fingerprint = handle_ref.system_fingerprint.clone(); + + // Use tokio runtime to run async code + let result = RUNTIME.block_on(async { + convert_proto_chunk_to_openai( + proto_response, + handle_ref, + &tokenizer, + &model, + &request_id, + created, + system_fingerprint.as_deref(), + ) + .await + }); + + match result { + Ok(Some(openai_response)) => { + // Serialize to JSON + let result_str = match serde_json::to_string(&openai_response) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to serialize response: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + let result_cstr = match CString::new(result_str) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + + *result_json_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Ok(None) => { + // No response to send (e.g., empty chunk) + let empty = CString::new("").unwrap(); + *result_json_out = empty.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &format!("Conversion error: {}", e)); + SglErrorCode::ParsingError + } + } +} + +/// Helper function to convert proto chunk to OpenAI format +pub(crate) async fn convert_proto_chunk_to_openai( + proto_response: proto::GenerateResponse, + handle: &mut GrpcResponseConverterHandle, + tokenizer: &Arc, + model: &str, + request_id: &str, + created: u64, + system_fingerprint: Option<&str>, +) -> Result, String> { + use sglang_router::grpc_client::sglang_proto::generate_response::Response::*; + use sglang_router::protocols::chat::{ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice}; + + match proto_response.response { + Some(Chunk(chunk)) => { + let index = chunk.index; + + // Mark as not first chunk if we've seen this index before + let is_first = handle.is_first_chunk.entry(index).or_insert(true); + let first_chunk = *is_first; + *is_first = false; + + // Track token counts from chunks (cumulative values from proto) + // These are cumulative values, so we always use the latest value + // For prompt_tokens, if chunk value is 0, preserve existing value or use initial_prompt_tokens + // This prevents overwriting valid prompt_tokens with 0 + if chunk.prompt_tokens > 0 { + handle.prompt_tokens.insert(index, chunk.prompt_tokens); + } else { + // If chunk.prompt_tokens is 0, try to preserve existing value or use initial_prompt_tokens + if !handle.prompt_tokens.contains_key(&index) { + // No existing value, try to use initial_prompt_tokens + if let Some(initial_prompt) = handle.initial_prompt_tokens { + handle.prompt_tokens.insert(index, initial_prompt); + } + } + // If existing value exists, keep it (don't overwrite with 0) + } + // For completion_tokens, always update (even if 0) as it's cumulative + handle.completion_tokens.insert(index, chunk.completion_tokens); + + // Process tokens through stop decoder if available, otherwise use incremental decoder + let chunk_text = if let Some(ref stop_decoder) = handle.stop_decoder { + let mut decoder_guard = stop_decoder.lock().await; + let mut text = String::new(); + for &token_id in &chunk.token_ids { + match decoder_guard.process_token(token_id).unwrap_or_else(|_| { + sglang_router::tokenizer::stop::SequenceDecoderOutput::Held + }) { + sglang_router::tokenizer::stop::SequenceDecoderOutput::Text(t) => { + text.push_str(&t); + } + sglang_router::tokenizer::stop::SequenceDecoderOutput::StoppedWithText(t) => { + text.push_str(&t); + break; + } + sglang_router::tokenizer::stop::SequenceDecoderOutput::Stopped => { + break; + } + sglang_router::tokenizer::stop::SequenceDecoderOutput::Held => {} + } + } + text + } else { + // Use incremental decoder to handle multi-byte character boundaries + let decode_stream = handle.decode_streams.entry(index).or_insert_with(|| { + DecodeStream::new( + Arc::clone(&tokenizer), + &[], // No prompt tokens for completion + handle.skip_special_tokens, + ) + }); + + // Process tokens incrementally + let mut text_parts = Vec::new(); + for &token_id in &chunk.token_ids { + if let Ok(Some(text)) = decode_stream.step(token_id) { + text_parts.push(text); + } + } + text_parts.join("") + }; + + if chunk_text.is_empty() { + return Ok(None); + } + + // Send first chunk with role + if first_chunk { + let first_response = ChatCompletionStreamResponse { + id: request_id.to_string(), + object: "chat.completion.chunk".to_string(), + created, + model: model.to_string(), + system_fingerprint: system_fingerprint.map(|s| s.to_string()), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: Some("assistant".to_string()), + content: None, + tool_calls: None, + reasoning_content: None, + }, + logprobs: None, + finish_reason: None, + matched_stop: None, + }], + usage: None, + }; + return Ok(Some(first_response)); + } + + // Update stream buffer + let stream_buffer = handle.stream_buffers.entry(index).or_default(); + stream_buffer.push_str(&chunk_text); + + // Handle tool calls if tools are provided + if let (Some(ref tools), Some(ref tool_parser)) = (handle.tools.as_ref(), handle.tool_parser.as_ref()) { + let tool_choice_enabled = !matches!( + handle.tool_choice, + Some(ToolChoice::Value(ToolChoiceValue::None)) + ); + + if tool_choice_enabled { + let mut parser_guard = tool_parser.lock().await; + match parser_guard.parse_incremental(&chunk_text, tools).await { + Ok(streaming_result) => { + if !streaming_result.calls.is_empty() { + handle.has_tool_calls.insert(index, true); + // Convert tool call items to OpenAI format + let tool_call_deltas: Vec<_> = streaming_result + .calls + .into_iter() + .map(|item| { + let id = if let Some(ref name) = item.name { + generate_tool_call_id( + model, + name, + item.tool_index, + handle.history_tool_calls_count, + ) + } else { + format!("call_{}", item.tool_index) + }; + + ToolCallDelta { + index: item.tool_index as u32, + id: Some(id), + tool_type: if item.name.is_some() { + Some("function".to_string()) + } else { + None + }, + function: Some(FunctionCallDelta { + name: item.name, + arguments: if !item.parameters.is_empty() { + Some(item.parameters) + } else { + None + }, + }), + } + }) + .collect(); + + let tool_response = ChatCompletionStreamResponse { + id: request_id.to_string(), + object: "chat.completion.chunk".to_string(), + created, + model: model.to_string(), + system_fingerprint: system_fingerprint.map(|s| s.to_string()), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: Some("assistant".to_string()), + content: None, + tool_calls: Some(tool_call_deltas), + reasoning_content: None, + }, + logprobs: None, + finish_reason: None, + matched_stop: None, + }], + usage: None, + }; + return Ok(Some(tool_response)); + } + } + Err(e) => { + // Log error but continue with regular content + tracing::warn!("Tool parser error: {}", e); + } + } + } + } + + // Regular content emission + let content_response = ChatCompletionStreamResponse { + id: request_id.to_string(), + object: "chat.completion.chunk".to_string(), + created, + model: model.to_string(), + system_fingerprint: system_fingerprint.map(|s| s.to_string()), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: Some("assistant".to_string()), + content: Some(chunk_text), + tool_calls: None, + reasoning_content: None, + }, + logprobs: None, + finish_reason: None, + matched_stop: None, + }], + usage: None, + }; + + Ok(Some(content_response)) + } + Some(Complete(complete)) => { + let index = complete.index; + + // Flush any remaining text + // Flush any remaining text from decode stream + let mut final_text = handle.stream_buffers.remove(&index).unwrap_or_default(); + if let Some(ref mut decode_stream) = handle.decode_streams.get_mut(&index) { + if let Ok(Some(remaining)) = decode_stream.flush() { + final_text.push_str(&remaining); + } + } + handle.decode_streams.remove(&index); + + // Determine finish reason - ensure it's never empty + // If finish_reason is empty, try to infer from other fields or use default + let finish_reason = if handle.has_tool_calls.get(&index).copied().unwrap_or(false) + && (complete.finish_reason == "stop" || complete.finish_reason.is_empty()) + { + "tool_calls".to_string() + } else if complete.finish_reason.is_empty() || complete.finish_reason.trim().is_empty() { + // If finish_reason is empty, try to infer from completion_tokens or use default + if complete.completion_tokens > 0 { + // If we have completion tokens, likely stopped normally + "stop".to_string() + } else if !complete.output_ids.is_empty() { + // If we have output_ids, likely stopped normally + "stop".to_string() + } else { + // Default fallback - always ensure we have a value + "stop".to_string() + } + } else { + complete.finish_reason.clone() + }; + + // Ensure finish_reason is never empty (defensive check) + let finish_reason = if finish_reason.is_empty() || finish_reason.trim().is_empty() { + "stop".to_string() + } else { + finish_reason + }; + + // Extract matched_stop + let matched_stop = match &complete.matched_stop { + Some(proto::generate_complete::MatchedStop::MatchedTokenId(token_id)) => { + Some(Value::Number(serde_json::Number::from(*token_id))) + } + Some(proto::generate_complete::MatchedStop::MatchedStopStr(stop_str)) => { + Some(Value::String(stop_str.clone())) + } + None => None, + }; + + // Build usage - prefer values from complete message, but fallback to accumulated values from chunks + // Complete message should have the final values, but sometimes they might be 0 or missing + // Always use the latest cumulative value from chunks if available, otherwise use complete message value + let mut prompt_tokens = handle.prompt_tokens.get(&index) + .copied() + .filter(|&v| v > 0) + .unwrap_or(complete.prompt_tokens); + let mut completion_tokens = handle.completion_tokens.get(&index) + .copied() + .filter(|&v| v > 0) + .unwrap_or(complete.completion_tokens); + + // Always try to use initial_prompt_tokens if prompt_tokens is 0 or missing + // This is the most reliable source for prompt tokens since we calculate it from the request + if prompt_tokens == 0 { + if let Some(initial_prompt) = handle.initial_prompt_tokens { + prompt_tokens = initial_prompt; + } + } + + // If completion_tokens is 0, try to infer from output_ids or accumulated chunks + if completion_tokens == 0 { + // Try to use completion_tokens from complete message even if 0 + // Or calculate from output_ids + if complete.completion_tokens > 0 { + completion_tokens = complete.completion_tokens; + } else if !complete.output_ids.is_empty() { + completion_tokens = complete.output_ids.len() as i32; + } else if let Some(&last_completion) = handle.completion_tokens.get(&index) { + completion_tokens = last_completion; + } + } + + // Final fallback: if both are still 0, try to use initial_prompt_tokens for prompt + // and calculate completion from output_ids + if prompt_tokens == 0 && completion_tokens == 0 { + // Try to infer from output_ids if available + let output_ids_len = complete.output_ids.len() as i32; + if output_ids_len > 0 { + completion_tokens = output_ids_len; + // Always try to use initial_prompt_tokens for prompt + if let Some(initial_prompt) = handle.initial_prompt_tokens { + prompt_tokens = initial_prompt; + } + } + } + + // Final defensive check: ensure prompt_tokens is set if we have initial_prompt_tokens + if prompt_tokens == 0 { + if let Some(initial_prompt) = handle.initial_prompt_tokens { + prompt_tokens = initial_prompt; + } + } + + // Always create usage, even if values are 0 (defensive) + let usage = Some(Usage { + prompt_tokens: prompt_tokens.max(0) as u32, + completion_tokens: completion_tokens.max(0) as u32, + total_tokens: (prompt_tokens.max(0) + completion_tokens.max(0)) as u32, + completion_tokens_details: None, + }); + + let finish_response = ChatCompletionStreamResponse { + id: request_id.to_string(), + object: "chat.completion.chunk".to_string(), + created, + model: model.to_string(), + system_fingerprint: system_fingerprint.map(|s| s.to_string()), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: Some("assistant".to_string()), + content: if !final_text.is_empty() { + Some(final_text) + } else { + None + }, + tool_calls: None, + reasoning_content: None, + }, + logprobs: None, + finish_reason: Some(finish_reason), + matched_stop, + }], + usage, + }; + + Ok(Some(finish_response)) + } + Some(Error(error)) => { + Err(format!("Server error: {} (status: {})", error.message, error.http_status_code)) + } + None => Ok(None), + } +} + +/// Free a gRPC response converter handle +#[no_mangle] +pub unsafe extern "C" fn sgl_grpc_response_converter_free(handle: *mut GrpcResponseConverterHandle) { + if !handle.is_null() { + let _ = Box::from_raw(handle); + } +} diff --git a/sgl-router/bindings/golang/src/lib.rs b/sgl-router/bindings/golang/src/lib.rs new file mode 100644 index 000000000..5bbe027ae --- /dev/null +++ b/sgl-router/bindings/golang/src/lib.rs @@ -0,0 +1,88 @@ +//! FFI module for exposing sgl-router preprocessing and postprocessing functions +//! to C-compatible languages (e.g., Golang via cgo) +//! +//! This module provides C-compatible function signatures for: +//! - Tokenizer operations (encode, decode, chat template) +//! - Tool parser operations (parse tool calls) +//! - Tool constraint generation +//! - gRPC client SDK (complete request-response flow) +//! +//! # Safety +//! All functions marked with `#[no_mangle]` and `extern "C"` must be called +//! with valid pointers and follow the documented memory management rules. + +// Re-export error types +pub use error::{SglErrorCode, set_error_message, set_error_message_fmt, clear_error_message}; + +// Re-export memory management functions +pub use memory::{sgl_free_string, sgl_free_token_ids}; + +// Re-export tokenizer functions +pub use tokenizer::{ + TokenizerHandle, + sgl_tokenizer_create_from_file, + sgl_tokenizer_encode, + sgl_tokenizer_apply_chat_template, + sgl_tokenizer_apply_chat_template_with_tools, + sgl_tokenizer_decode, + sgl_tokenizer_free, +}; + +// Re-export tool parser functions +pub use tool_parser::{ + ToolParserHandle, + sgl_tool_parser_create, + sgl_tool_parser_parse_complete, + sgl_tool_parser_parse_incremental, + sgl_tool_parser_reset, + sgl_tool_parser_free, +}; + +// Re-export gRPC converter functions +pub use grpc_converter::{ + GrpcResponseConverterHandle, + sgl_grpc_response_converter_create, + sgl_grpc_response_converter_convert_chunk, + sgl_grpc_response_converter_free, +}; + +// Re-export client SDK functions +pub use client::{ + SglangClientHandle, + sgl_client_create, + sgl_client_free, +}; + +// Re-export stream functions +pub use stream::{ + SglangStreamHandle, + sgl_stream_read_next, + sgl_stream_free, +}; + +// Re-export client stream function (defined in client.rs but used by stream) +pub use client::sgl_client_chat_completion_stream; + +// Re-export utility functions +pub use utils::sgl_generate_tool_constraints; + +// Sub-modules +mod error; +mod memory; +mod tokenizer; +mod tool_parser; +mod grpc_converter; +mod client; +mod stream; +mod utils; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_error_codes() { + assert_eq!(SglErrorCode::Success as i32, 0); + assert_eq!(SglErrorCode::InvalidArgument as i32, 1); + } +} diff --git a/sgl-router/bindings/golang/src/memory.rs b/sgl-router/bindings/golang/src/memory.rs new file mode 100644 index 000000000..9aac9b333 --- /dev/null +++ b/sgl-router/bindings/golang/src/memory.rs @@ -0,0 +1,28 @@ +//! Memory management for FFI functions + +use std::ffi::CString; +use std::os::raw::c_char; + +/// Free a C string allocated by Rust +/// +/// # Safety +/// This function must only be called with pointers returned by other FFI functions. +/// Calling with arbitrary pointers or multiple times on the same pointer is undefined behavior. +#[no_mangle] +pub unsafe extern "C" fn sgl_free_string(s: *mut c_char) { + if !s.is_null() { + let _ = CString::from_raw(s); + } +} + +/// Free token IDs array allocated by Rust +/// +/// # Safety +/// This function must only be called with pointers returned by `sgl_tokenizer_encode`. +/// The `count` parameter must match the length of the array. +#[no_mangle] +pub unsafe extern "C" fn sgl_free_token_ids(ptr: *mut u32, count: usize) { + if !ptr.is_null() && count > 0 { + let _ = Vec::from_raw_parts(ptr, count, count); + } +} diff --git a/sgl-router/bindings/golang/src/stream.rs b/sgl-router/bindings/golang/src/stream.rs new file mode 100644 index 000000000..8bd57ac6c --- /dev/null +++ b/sgl-router/bindings/golang/src/stream.rs @@ -0,0 +1,288 @@ +//! Stream handling FFI functions +//! +//! This module provides FFI (Foreign Function Interface) functions for managing +//! streaming responses from the SGLang gRPC API. It handles: +//! +//! - Creating and managing stream handles +//! - Reading chunks from streams and converting them to OpenAI format +//! - Managing automatic abort on stream drop (via AbortOnDropStream) +//! - Thread-safe access to streams and response converters +//! +//! # Safety +//! +//! All FFI functions are marked `unsafe` as per Rust FFI conventions. Callers must: +//! - Pass valid pointers +//! - Ensure proper pointer lifetime management +//! - Call corresponding free functions for cleanup + +use std::ffi::CString; +use std::os::raw::{c_char, c_int}; +use std::ptr; +use std::sync::Arc; +use tokio::runtime::Runtime; +use once_cell::sync::Lazy; +use futures_util::StreamExt; + +use sglang_router::grpc_client::{sglang_proto as proto, sglang_scheduler::{SglangSchedulerClient, AbortOnDropStream}}; + +use super::error::{SglErrorCode, set_error_message}; +use super::grpc_converter::{GrpcResponseConverterHandle, convert_proto_chunk_to_openai}; + +/// Global tokio runtime for async operations +static RUNTIME: Lazy = Lazy::new(|| { + Runtime::new().expect("Failed to create tokio runtime for stream FFI") +}); + +/// Handle for an active streaming request. +/// +/// This struct manages the stream and response converter for a single request. +/// It is wrapped in Arc and Mutex for thread-safe concurrent access. +/// +/// # Fields +/// +/// * `stream` - The gRPC stream wrapped in AbortOnDropStream for automatic cleanup +/// * `converter` - Response converter that transforms proto messages to OpenAI format +/// * `client` - The underlying gRPC client connection +/// * `prompt_tokens` - Number of prompt tokens from the original request +pub struct SglangStreamHandle { + pub(crate) stream: Arc>, + pub(crate) converter: Arc>, + #[allow(dead_code)] + pub(crate) client: Arc, + #[allow(dead_code)] + pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request +} + +/// Read next chunk from stream and convert to OpenAI format. +/// +/// This function reads the next chunk from the gRPC stream, converts it from the +/// internal protocol format to OpenAI-compatible JSON format, and returns it via +/// the output parameters. +/// +/// # Arguments +/// +/// * `stream_handle` - Mutable pointer to the stream handle +/// * `response_json_out` - Pointer to receive OpenAI format JSON string +/// - Caller must free this with `sgl_free_string` +/// - May be NULL if no data available +/// * `is_done_out` - Pointer to receive completion status +/// - 0 = stream has more data +/// - 1 = stream is complete +/// * `error_out` - Optional pointer to receive error message +/// - Only set if function returns an error code +/// - Must be freed with `sgl_free_string` if not NULL +/// +/// # Returns +/// +/// * `SglErrorCode::Success` - Successfully read a chunk or reached end of stream +/// * Other error codes - See `SglErrorCode` for details +/// +/// # Safety +/// +/// - All pointers must be valid and properly aligned +/// - `stream_handle` must point to a valid `SglangStreamHandle` +/// - Output pointers must be writable +/// +/// # Notes +/// +/// - Complete messages are identified by the presence of `proto::GenerateResponse::Complete` +/// - When is_done=1, this may be the last readable chunk or the stream may be ending +/// - Subsequent calls after is_done=1 will mark the stream as complete internally +#[no_mangle] +pub unsafe extern "C" fn sgl_stream_read_next( + stream_handle: *mut SglangStreamHandle, + response_json_out: *mut *mut c_char, + is_done_out: *mut c_int, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if stream_handle.is_null() || response_json_out.is_null() || is_done_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let handle_ref = &*stream_handle; + let stream = Arc::clone(&handle_ref.stream); + let converter = Arc::clone(&handle_ref.converter); + + // Read next chunk from stream + let chunk_result = RUNTIME.block_on(async { + let mut stream_guard = stream.lock().await; + stream_guard.next().await + }); + + match chunk_result { + Some(Ok(proto_response)) => { + // Convert proto response to OpenAI format + // We need to get the converter lock first + let conversion_result = RUNTIME.block_on(async { + let mut converter_guard = converter.lock().await; + + // Clone necessary fields for conversion + let tokenizer = Arc::clone(&converter_guard.tokenizer); + let model = converter_guard.model.clone(); + let request_id = converter_guard.request_id.clone(); + let created = converter_guard.created; + let system_fingerprint = converter_guard.system_fingerprint.clone(); + + // Call the conversion function + convert_proto_chunk_to_openai( + proto_response.clone(), + &mut *converter_guard, + &tokenizer, + &model, + &request_id, + created, + system_fingerprint.as_deref(), + ) + .await + }); + + match conversion_result { + Ok(Some(openai_response)) => { + // Serialize to JSON + let result_str = match serde_json::to_string(&openai_response) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to serialize response: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + let result_cstr = match CString::new(result_str) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + + // Check if this is a complete response (stream done) + let is_complete = matches!(proto_response.response, Some(proto::generate_response::Response::Complete(_)) | Some(proto::generate_response::Response::Error(_))); + + *response_json_out = result_cstr.into_raw(); + *is_done_out = if is_complete { 1 } else { 0 }; + + if is_complete { + // Mark stream as completed + // Ensure mark_completed() completes and is visible before returning + // Use yield_now to ensure Release ordering is fully propagated + RUNTIME.block_on(async { + let stream_guard = stream.lock().await; + stream_guard.mark_completed(); + // Keep the guard until mark_completed() is fully executed + drop(stream_guard); + // Yield to ensure Release ordering is propagated before returning + // This prevents race condition where Free() is called immediately + // and Drop might not see the mark_completed() write + tokio::task::yield_now().await; + }); + } + + SglErrorCode::Success + } + Ok(None) => { + // No response to send (e.g., empty chunk) + // Don't mark as completed - stream might continue + // Just return null and let caller read more + *response_json_out = ptr::null_mut(); + *is_done_out = 0; // Keep stream open, not done yet + SglErrorCode::Success + } + Err(e) => { + // Conversion error - don't mark as completed + // Let the stream end naturally or return error without stopping stream + set_error_message(error_out, &format!("Conversion error: {}", e)); + *response_json_out = ptr::null_mut(); + *is_done_out = 0; // Don't mark as done - let caller decide + SglErrorCode::ParsingError + } + } + } + Some(Err(e)) => { + // Stream error - mark as completed to prevent abort + RUNTIME.block_on(async { + let stream_guard = stream.lock().await; + stream_guard.mark_completed(); + drop(stream_guard); + // Yield to ensure Release ordering is propagated + tokio::task::yield_now().await; + }); + + set_error_message(error_out, &format!("Stream error: {}", e)); + *is_done_out = 1; + SglErrorCode::UnknownError + } + None => { + // Stream ended naturally (no more chunks) + // Mark stream as completed before returning to prevent abort + RUNTIME.block_on(async { + let stream_guard = stream.lock().await; + stream_guard.mark_completed(); + drop(stream_guard); + // Yield to ensure Release ordering is propagated + tokio::task::yield_now().await; + }); + + *response_json_out = ptr::null_mut(); + *is_done_out = 1; + SglErrorCode::Success + } + } +} + +/// Free a stream handle and release all associated resources. +/// +/// This function must be called exactly once for each stream handle returned by +/// `sgl_client_chat_completion_stream`. It marks the stream as completed internally +/// to prevent abort signals from being sent when resources are cleaned up. +/// +/// # Arguments +/// +/// * `handle` - Mutable pointer to the stream handle to free +/// - If NULL, this function does nothing +/// +/// # Safety +/// +/// - Must be called only once per handle +/// - Handle must not be used after calling this function +/// - After this call, the stream is no longer valid +/// +/// # Notes +/// +/// - This function internally calls `mark_completed()` before freeing to ensure +/// the stream cleanup doesn't trigger an abort RPC to the server +/// - Memory fences are used to ensure visibility across threads +#[no_mangle] +pub unsafe extern "C" fn sgl_stream_free(handle: *mut SglangStreamHandle) { + if !handle.is_null() { + let handle_ref = Box::from_raw(handle); + + // Mark stream as completed to prevent abort on drop + // By this point, the stream should already be completed by ReadNext() + // but we call it again to be safe + RUNTIME.block_on(async { + let stream_guard = handle_ref.stream.lock().await; + stream_guard.mark_completed(); + // Keep guard alive to ensure mark_completed() write completes + drop(stream_guard); + // Yield to ensure the atomic write is visible + tokio::task::yield_now().await; + }); + + // Use a strong memory fence to ensure mark_completed()'s Release write + // is visible before we drop the last Arc reference + std::sync::atomic::fence(std::sync::atomic::Ordering::SeqCst); + + // Now drop all references - if mark_completed() was called successfully, + // the drop won't send an abort + drop(handle_ref.stream); + + // Free converter + let converter = Arc::try_unwrap(handle_ref.converter) + .ok() + .map(|m| m.into_inner()); + if let Some(conv) = converter { + super::grpc_converter::sgl_grpc_response_converter_free(Box::into_raw(Box::new(conv))); + } + } +} diff --git a/sgl-router/bindings/golang/src/tokenizer.rs b/sgl-router/bindings/golang/src/tokenizer.rs new file mode 100644 index 000000000..191054c29 --- /dev/null +++ b/sgl-router/bindings/golang/src/tokenizer.rs @@ -0,0 +1,379 @@ +//! Tokenizer FFI functions + +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char, c_int}; +use std::ptr; +use std::sync::Arc; +use serde_json::Value; + +use sglang_router::tokenizer::{ + create_tokenizer_from_file, + traits::Tokenizer as TokenizerTrait, + chat_template::ChatTemplateParams, + huggingface::HuggingFaceTokenizer, +}; + +use super::error::{SglErrorCode, set_error_message, clear_error_message}; + +/// Opaque handle for a tokenizer instance +#[repr(C)] +pub struct TokenizerHandle { + pub(crate) tokenizer: Arc, +} + +/// Create a tokenizer from a file path +/// +/// # Arguments +/// * `path` - Path to tokenizer.json file (null-terminated C string) +/// * `error_out` - Optional pointer to receive error message (must be freed with sgl_free_string) +/// +/// # Returns +/// * Pointer to TokenizerHandle on success, null on failure +/// +/// # Safety +/// The returned handle must be freed with `sgl_tokenizer_free`. +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_create_from_file( + path: *const c_char, + error_out: *mut *mut c_char, +) -> *mut TokenizerHandle { + if path.is_null() { + set_error_message(error_out, "path cannot be null"); + return ptr::null_mut(); + } + + let path_str = match CStr::from_ptr(path).to_str() { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Invalid UTF-8 in path: {}", e)); + return ptr::null_mut(); + } + }; + + match create_tokenizer_from_file(path_str) { + Ok(tokenizer) => { + clear_error_message(error_out); + Box::into_raw(Box::new(TokenizerHandle { + tokenizer, + })) + } + Err(e) => { + set_error_message(error_out, &e.to_string()); + ptr::null_mut() + } + } +} + +/// Encode text to token IDs +/// +/// # Arguments +/// * `handle` - Tokenizer handle (must not be null) +/// * `text` - Input text (null-terminated C string) +/// * `token_ids_out` - Pointer to receive array of token IDs (must be freed with sgl_free_token_ids) +/// * `token_count_out` - Pointer to receive token count +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +/// +/// # Safety +/// The token_ids_out array must be freed with sgl_free_token_ids() after use. +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_encode( + handle: *mut TokenizerHandle, + text: *const c_char, + token_ids_out: *mut *mut u32, + token_count_out: *mut usize, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || text.is_null() || token_ids_out.is_null() || token_count_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let text_str = match CStr::from_ptr(text).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in text"); + return SglErrorCode::InvalidArgument; + } + }; + + let tokenizer = &(*handle).tokenizer; + match tokenizer.encode(text_str) { + Ok(encoding) => { + let token_ids = encoding.token_ids(); + let count = token_ids.len(); + + // Allocate memory for token IDs using Vec, then leak to give ownership to C + let vec = token_ids.to_vec(); + let ptr = vec.as_ptr() as *mut u32; + let _ = std::mem::ManuallyDrop::new(vec); + + *token_ids_out = ptr; + *token_count_out = count; + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &e.to_string()); + SglErrorCode::TokenizationError + } + } +} + +/// Apply chat template to messages with tools support +/// +/// # Arguments +/// * `handle` - Tokenizer handle +/// * `messages_json` - JSON string of messages array +/// * `tools_json` - Optional JSON string of tools array (null or empty string for no tools) +/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_apply_chat_template_with_tools( + handle: *mut TokenizerHandle, + messages_json: *const c_char, + tools_json: *const c_char, + result_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || messages_json.is_null() || result_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let messages_str = match CStr::from_ptr(messages_json).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in messages_json"); + return SglErrorCode::InvalidArgument; + } + }; + + // Parse JSON messages + let messages: Vec = match serde_json::from_str(messages_str) { + Ok(msgs) => msgs, + Err(e) => { + set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e)); + return SglErrorCode::InvalidArgument; + } + }; + + // Parse tools JSON if provided + let tools: Option> = if tools_json.is_null() { + None + } else { + let tools_str = match CStr::from_ptr(tools_json).to_str() { + Ok(s) => { + if s.is_empty() { + None + } else { + match serde_json::from_str::>(s) { + Ok(t) => Some(t), + Err(e) => { + set_error_message(error_out, &format!("Failed to parse tools JSON: {}", e)); + return SglErrorCode::InvalidArgument; + } + } + } + } + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in tools_json"); + return SglErrorCode::InvalidArgument; + } + }; + tools_str + }; + + // Get the tokenizer from handle + let handle_ref = &*handle; + let tokenizer = &handle_ref.tokenizer; + + // Try to downcast to HuggingFaceTokenizer + if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::() { + // Apply chat template with tools + let empty_docs: [Value; 0] = []; + let tools_slice = tools.as_ref().map(|t| t.as_slice()); + let params = ChatTemplateParams { + add_generation_prompt: true, + tools: tools_slice, + documents: Some(&empty_docs), + template_kwargs: None, + }; + + match hf_tokenizer.apply_chat_template(&messages, params) { + Ok(result) => { + let result_cstr = match CString::new(result) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + *result_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &format!("Failed to apply chat template: {}", e)); + SglErrorCode::TokenizationError + } + } + } else { + set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers"); + SglErrorCode::TokenizationError + } +} + +/// Apply chat template to messages +/// +/// # Arguments +/// * `handle` - Tokenizer handle +/// * `messages_json` - JSON string of messages array +/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_apply_chat_template( + handle: *mut TokenizerHandle, + messages_json: *const c_char, + result_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || messages_json.is_null() || result_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let messages_str = match CStr::from_ptr(messages_json).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in messages_json"); + return SglErrorCode::InvalidArgument; + } + }; + + // Parse JSON messages + let messages: Vec = match serde_json::from_str(messages_str) { + Ok(msgs) => msgs, + Err(e) => { + set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e)); + return SglErrorCode::InvalidArgument; + } + }; + + // Get the tokenizer from handle + let handle_ref = &*handle; + let tokenizer = &handle_ref.tokenizer; + + // Try to downcast to HuggingFaceTokenizer + if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::() { + // Apply chat template with default parameters + // Use empty arrays instead of None to avoid template errors + // Set add_generation_prompt to true so the model knows to start generating + let empty_tools: [Value; 0] = []; + let empty_docs: [Value; 0] = []; + let params = ChatTemplateParams { + add_generation_prompt: true, // Important: tells the model to start generating + tools: Some(&empty_tools), + documents: Some(&empty_docs), + template_kwargs: None, + }; + + match hf_tokenizer.apply_chat_template(&messages, params) { + Ok(result) => { + let result_cstr = match CString::new(result) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + *result_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &format!("Failed to apply chat template: {}", e)); + SglErrorCode::TokenizationError + } + } + } else { + set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers"); + SglErrorCode::TokenizationError + } +} + +/// Decode token IDs to text +/// +/// # Arguments +/// * `handle` - Tokenizer handle +/// * `token_ids` - Array of token IDs +/// * `token_count` - Number of tokens +/// * `skip_special_tokens` - Whether to skip special tokens +/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_decode( + handle: *mut TokenizerHandle, + token_ids: *const u32, + token_count: usize, + skip_special_tokens: c_int, + result_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || token_ids.is_null() || result_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + if token_count == 0 { + let empty = CString::new("").unwrap(); + *result_out = empty.into_raw(); + clear_error_message(error_out); + return SglErrorCode::Success; + } + + // Convert C array to Rust slice + let token_slice = std::slice::from_raw_parts(token_ids, token_count); + + let tokenizer = &(*handle).tokenizer; + match tokenizer.decode(token_slice, skip_special_tokens != 0) { + Ok(text) => { + let result_cstr = match CString::new(text) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + *result_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &e.to_string()); + SglErrorCode::TokenizationError + } + } +} + +/// Free a tokenizer handle +/// +/// # Safety +/// This function must only be called once per handle, and the handle must not be used after calling. +#[no_mangle] +pub unsafe extern "C" fn sgl_tokenizer_free(handle: *mut TokenizerHandle) { + if !handle.is_null() { + let _ = Box::from_raw(handle); + } +} diff --git a/sgl-router/bindings/golang/src/tool_parser.rs b/sgl-router/bindings/golang/src/tool_parser.rs new file mode 100644 index 000000000..1889ac162 --- /dev/null +++ b/sgl-router/bindings/golang/src/tool_parser.rs @@ -0,0 +1,329 @@ +//! Tool parser FFI functions + +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char}; +use std::ptr; +use std::sync::Arc; +use std::collections::HashMap; +use serde_json::{json, Value}; +use tokio::runtime::Runtime; +use once_cell::sync::Lazy; + +use sglang_router::tool_parser::{ParserFactory, ToolParser}; +use sglang_router::protocols::common::Tool; + +use super::error::{SglErrorCode, set_error_message, clear_error_message}; +use super::utils::generate_tool_call_id; + +/// Global parser factory (initialized once) +static PARSER_FACTORY: Lazy = Lazy::new(|| ParserFactory::new()); + +/// Global tokio runtime for async operations +static RUNTIME: Lazy = Lazy::new(|| { + Runtime::new().expect("Failed to create tokio runtime for tool parser FFI") +}); + +/// Opaque handle for a tool parser instance +/// Note: For streaming, we need mutable access, so we use Arc> internally +/// Note: This is an opaque handle, C code doesn't access fields directly +pub struct ToolParserHandle { + parser: Arc>>, + model: String, // Store model name for ID generation + history_tool_calls_count: usize, // Track tool call count for ID generation + tool_index_to_id: HashMap, // Map tool_index to ID for incremental updates +} + +/// Create a tool parser +/// +/// # Arguments +/// * `parser_type` - Parser type name (e.g., "json", "llama", "mistral") or model name (e.g., "gpt-4") +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * Pointer to ToolParserHandle on success, null on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tool_parser_create( + parser_type: *const c_char, + error_out: *mut *mut c_char, +) -> *mut ToolParserHandle { + if parser_type.is_null() { + set_error_message(error_out, "parser_type cannot be null"); + return ptr::null_mut(); + } + + let type_str = match CStr::from_ptr(parser_type).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in parser_type"); + return ptr::null_mut(); + } + }; + + // Create parser using factory + // The factory will determine the parser type based on model name or use the provided type + let parser = if let Some(parser_box) = PARSER_FACTORY.registry().create_for_model(type_str) { + parser_box + } else if let Some(parser_box) = PARSER_FACTORY.registry().create_parser(type_str) { + parser_box + } else { + set_error_message(error_out, &format!("Unknown parser type: {}", type_str)); + return ptr::null_mut(); + }; + + Box::into_raw(Box::new(ToolParserHandle { + parser: Arc::new(tokio::sync::Mutex::new(parser)), + model: type_str.to_string(), + history_tool_calls_count: 0, + tool_index_to_id: HashMap::new(), + })) +} + +/// Parse complete tool calls from text +/// +/// # Arguments +/// * `handle` - Tool parser handle +/// * `text` - Input text to parse +/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tool_parser_parse_complete( + handle: *mut ToolParserHandle, + text: *const c_char, + result_json_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || text.is_null() || result_json_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let text_str = match CStr::from_ptr(text).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in text"); + return SglErrorCode::InvalidArgument; + } + }; + + let handle_ref = &*handle; + let parser = Arc::clone(&handle_ref.parser); + let model = handle_ref.model.clone(); + let history_count = handle_ref.history_tool_calls_count; + + // Use tokio runtime to run async code + let result = RUNTIME.block_on(async { + let parser_guard = parser.lock().await; + parser_guard.parse_complete(text_str).await + }); + + match result { + Ok((normal_text, tool_calls)) => { + // Convert Rust ToolCall to OpenAI format + let openai_tool_calls: Vec = tool_calls + .into_iter() + .enumerate() + .map(|(index, tc)| { + // Generate ID for this tool call + let id = generate_tool_call_id(&model, &tc.function.name, index, history_count); + json!({ + "id": id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }) + .collect(); + + // Build result JSON + let result_json = json!({ + "normal_text": normal_text, + "tool_calls": openai_tool_calls + }); + + let result_str = match serde_json::to_string(&result_json) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to serialize JSON: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + let result_cstr = match CString::new(result_str) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + + *result_json_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &format!("Parse error: {}", e)); + SglErrorCode::ParsingError + } + } +} + +/// Parse tool calls incrementally from streaming chunks +/// +/// # Arguments +/// * `handle` - Tool parser handle +/// * `chunk` - New text chunk from stream +/// * `tools_json` - JSON array of available tools (for validation, can be null/empty) +/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string) +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_tool_parser_parse_incremental( + handle: *mut ToolParserHandle, + chunk: *const c_char, + tools_json: *const c_char, + result_json_out: *mut *mut c_char, + error_out: *mut *mut c_char, +) -> SglErrorCode { + if handle.is_null() || chunk.is_null() || result_json_out.is_null() { + set_error_message(error_out, "Invalid arguments: null pointer"); + return SglErrorCode::InvalidArgument; + } + + let chunk_str = match CStr::from_ptr(chunk).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in chunk"); + return SglErrorCode::InvalidArgument; + } + }; + + // Parse tools JSON if provided + let tools: Vec = if !tools_json.is_null() { + let tools_str = match CStr::from_ptr(tools_json).to_str() { + Ok(s) => s, + Err(_) => { + set_error_message(error_out, "Invalid UTF-8 in tools_json"); + return SglErrorCode::InvalidArgument; + } + }; + match serde_json::from_str::>(tools_str) { + Ok(t) => t, + Err(_) => vec![], // If parsing fails, use empty tools + } + } else { + vec![] + }; + + let handle_ref = &*handle; + let parser = Arc::clone(&handle_ref.parser); + let model = handle_ref.model.clone(); + let history_count = handle_ref.history_tool_calls_count; + + // Use tokio runtime to run async code + let result = RUNTIME.block_on(async { + let mut parser_guard = parser.lock().await; + parser_guard.parse_incremental(chunk_str, &tools).await + }); + + match result { + Ok(streaming_result) => { + // Convert StreamingParseResult to OpenAI format + let handle_mut = &mut *handle; + let openai_tool_calls: Vec = streaming_result + .calls + .into_iter() + .map(|item| { + // For incremental parsing, we may not have complete tool calls yet + // Generate or reuse ID based on tool_index + let id = if let Some(ref name) = item.name { + // New tool call with name - generate ID and store it + let id = generate_tool_call_id(&model, name, item.tool_index, history_count); + handle_mut.tool_index_to_id.insert(item.tool_index, id.clone()); + id + } else { + // Parameter update - reuse existing ID for this tool_index + handle_mut.tool_index_to_id + .get(&item.tool_index) + .cloned() + .unwrap_or_else(|| format!("call_{}", item.tool_index)) + }; + + json!({ + "id": id, + "type": "function", + "function": { + "name": item.name.unwrap_or_default(), + "arguments": item.parameters + } + }) + }) + .collect(); + + // Build result JSON + let result_json = json!({ + "normal_text": streaming_result.normal_text, + "tool_calls": openai_tool_calls + }); + + let result_str = match serde_json::to_string(&result_json) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to serialize JSON: {}", e)); + return SglErrorCode::ParsingError; + } + }; + + let result_cstr = match CString::new(result_str) { + Ok(s) => s, + Err(e) => { + set_error_message(error_out, &format!("Failed to create result string: {}", e)); + return SglErrorCode::MemoryError; + } + }; + + *result_json_out = result_cstr.into_raw(); + clear_error_message(error_out); + SglErrorCode::Success + } + Err(e) => { + set_error_message(error_out, &format!("Parse incremental error: {}", e)); + SglErrorCode::ParsingError + } + } +} + +/// Reset the parser state for reuse +#[no_mangle] +pub unsafe extern "C" fn sgl_tool_parser_reset(handle: *mut ToolParserHandle) { + if handle.is_null() { + return; + } + + let handle_ref = &mut *handle; + let parser = Arc::clone(&handle_ref.parser); + + // Reset parser state + RUNTIME.block_on(async { + let mut parser_guard = parser.lock().await; + parser_guard.reset(); + }); + + // Reset history count and tool index mapping + handle_ref.history_tool_calls_count = 0; + handle_ref.tool_index_to_id.clear(); +} + +/// Free a tool parser handle +#[no_mangle] +pub unsafe extern "C" fn sgl_tool_parser_free(handle: *mut ToolParserHandle) { + if !handle.is_null() { + let _ = Box::from_raw(handle); + } +} diff --git a/sgl-router/bindings/golang/src/utils.rs b/sgl-router/bindings/golang/src/utils.rs new file mode 100644 index 000000000..11bc82d94 --- /dev/null +++ b/sgl-router/bindings/golang/src/utils.rs @@ -0,0 +1,44 @@ +//! Utility functions for FFI + +use uuid::Uuid; + +/// Helper function to generate tool call ID (matches router implementation) +pub fn generate_tool_call_id( + model: &str, + function_name: &str, + index: usize, + history_tool_calls_count: usize, +) -> String { + if model.to_lowercase().contains("kimi") { + // KimiK2 format: functions.{name}:{global_index} + format!("functions.{}:{}", function_name, history_tool_calls_count + index) + } else { + // Standard OpenAI format: call_{24-char-uuid} + format!("call_{}", &Uuid::new_v4().simple().to_string()[..24]) + } +} + +/// Generate tool constraints (placeholder implementation) +/// +/// # Arguments +/// * `tools_json` - JSON array of tools +/// * `tool_choice_json` - JSON object representing tool_choice +/// * `constraint_type_out` - Pointer to receive constraint type (e.g., "json_schema") +/// * `constraint_schema_out` - Pointer to receive constraint schema JSON +/// * `error_out` - Optional pointer to receive error message +/// +/// # Returns +/// * SglErrorCode::Success on success, error code on failure +#[no_mangle] +pub unsafe extern "C" fn sgl_generate_tool_constraints( + _tools_json: *const std::os::raw::c_char, + _tool_choice_json: *const std::os::raw::c_char, + _constraint_type_out: *mut *mut std::os::raw::c_char, + _constraint_schema_out: *mut *mut std::os::raw::c_char, + error_out: *mut *mut std::os::raw::c_char, +) -> super::error::SglErrorCode { + // Implementation would parse JSON and call generate_tool_constraints + // This is a placeholder + super::error::set_error_message(error_out, "Tool constraint generation not yet implemented in FFI"); + super::error::SglErrorCode::UnknownError +} diff --git a/sgl-router/src/routers/grpc/regular/responses/handlers.rs b/sgl-router/src/routers/grpc/regular/responses/handlers.rs index 97108b104..967d06c62 100644 --- a/sgl-router/src/routers/grpc/regular/responses/handlers.rs +++ b/sgl-router/src/routers/grpc/regular/responses/handlers.rs @@ -389,9 +389,8 @@ async fn process_and_transform_sse_stream( // Include reasoning_tokens if present if let Some(details) = &u.completion_tokens_details { if let Some(reasoning_tokens) = details.reasoning_tokens { - usage_obj["output_tokens_details"] = json!({ - "reasoning_tokens": reasoning_tokens - }); + usage_obj["output_tokens_details"] = + json!({ "reasoning_tokens": reasoning_tokens }); } }