iop/apps/node/internal/adapters/vllm/vllm.go

806 lines
21 KiB
Go

// Package vllm provides an Adapter for OpenAI-compatible inference engines
// such as vLLM and SGLang via the /v1/chat/completions SSE endpoint.
package vllm
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
"go.uber.org/zap"
"iop/apps/node/internal/runtime"
"iop/packages/go/config"
)
const Name = "vllm"
const (
runtimeMetadataOpenAIToolCalls = "openai_tool_calls"
runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
)
type Vllm struct {
instanceName string
endpoint string
capacity int
maxQueue int
queueTimeoutMS int
requestTimeoutMS int
client *http.Client
logger *zap.Logger
}
// New creates a vLLM adapter. The optional instanceName is the registry
// instance key used to disambiguate multiple vLLM instances on the same node.
func New(cfg config.VllmConf, logger *zap.Logger, instanceName ...string) *Vllm {
endpoint := strings.TrimRight(cfg.Endpoint, "/")
name := Name
if len(instanceName) > 0 && instanceName[0] != "" {
name = instanceName[0]
}
return &Vllm{
instanceName: name,
endpoint: endpoint,
capacity: cfg.Capacity,
maxQueue: cfg.MaxQueue,
queueTimeoutMS: cfg.QueueTimeoutMS,
requestTimeoutMS: cfg.RequestTimeoutMS,
client: &http.Client{},
logger: logger,
}
}
func (v *Vllm) Name() string { return Name }
func (v *Vllm) TunnelProvider(ctx context.Context, req runtime.ProviderTunnelRequest, sink runtime.ProviderTunnelSink) error {
var seq int64 = 0
if v.endpoint == "" {
err := fmt.Errorf("vllm adapter: endpoint is required")
_ = emitTunnelError(ctx, sink, req, seq, err)
return err
}
urlStr := joinURL(v.endpoint, req.Path)
httpReq, err := http.NewRequestWithContext(ctx, req.Method, urlStr, bytes.NewReader(req.Body))
if err != nil {
err = fmt.Errorf("build request: %w", err)
_ = emitTunnelError(ctx, sink, req, seq, err)
return err
}
if req.Method == http.MethodPost {
httpReq.Header.Set("Content-Type", "application/json")
}
for k, val := range req.Headers {
httpReq.Header.Set(k, val)
}
resp, err := v.client.Do(httpReq)
if err != nil {
err = fmt.Errorf("request failed: %w", err)
_ = emitTunnelError(ctx, sink, req, seq, err)
return err
}
defer resp.Body.Close()
respHeaders := make(map[string]string)
for k, values := range resp.Header {
if len(values) > 0 {
respHeaders[k] = values[0]
}
}
err = sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: seq,
Kind: runtime.ProviderTunnelFrameKindResponseStart,
StatusCode: resp.StatusCode,
Headers: respHeaders,
Timestamp: time.Now(),
})
if err != nil {
return err
}
seq++
buf := make([]byte, 4096)
for {
n, readErr := resp.Body.Read(buf)
if n > 0 {
err = sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: seq,
Kind: runtime.ProviderTunnelFrameKindBody,
Body: append([]byte(nil), buf[:n]...),
Timestamp: time.Now(),
})
if err != nil {
return err
}
seq++
}
if readErr != nil {
if readErr == io.EOF {
break
}
if ctx.Err() != nil {
_ = emitTunnelError(ctx, sink, req, seq, ctx.Err())
return ctx.Err()
}
err = fmt.Errorf("read response body: %w", readErr)
_ = emitTunnelError(ctx, sink, req, seq, err)
return err
}
}
err = sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: seq,
Kind: runtime.ProviderTunnelFrameKindEnd,
End: true,
Timestamp: time.Now(),
})
return err
}
func emitTunnelError(ctx context.Context, sink runtime.ProviderTunnelSink, req runtime.ProviderTunnelRequest, seq int64, err error) error {
return sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: seq,
Kind: runtime.ProviderTunnelFrameKindError,
Error: err.Error(),
Timestamp: time.Now(),
})
}
func (v *Vllm) Capabilities(ctx context.Context) (runtime.Capabilities, error) {
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
targets, err := v.fetchTargets(probeCtx)
status := runtime.ProviderStatusAvailable
if err != nil {
status = runtime.ProviderStatusUnavailable
}
return runtime.Capabilities{
AdapterName: Name,
InstanceKey: v.instanceName,
Targets: targets,
MaxConcurrency: effectiveCapacity(v.capacity, 8),
MaxQueue: v.maxQueue,
QueueTimeoutMS: v.queueTimeoutMS,
RequestTimeoutMS: v.requestTimeoutMS,
ProviderStatus: runtime.NormalizeProviderStatus(status),
}, nil
}
func effectiveCapacity(capacity, defaultVal int) int {
if capacity <= 0 {
return defaultVal
}
return capacity
}
func (v *Vllm) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
if v.endpoint == "" {
return fmt.Errorf("vllm adapter: endpoint is required")
}
model := strings.TrimSpace(spec.Target)
if model == "" {
model = stringInput(spec.Input, "model")
}
if model == "" {
return fmt.Errorf("vllm adapter: target/model is required")
}
messages := messagesFromInput(spec.Input)
if len(messages) == 0 {
return fmt.Errorf("vllm adapter: messages are required")
}
if err := sink.Emit(ctx, runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeStart,
Timestamp: time.Now(),
}); err != nil {
return err
}
chatReq := vllmChatRequest{
Model: model,
Messages: messages,
Stream: true,
}
if v, ok := spec.Input["tools"]; ok {
chatReq.Tools = v
}
if v, ok := spec.Input["tool_choice"]; ok {
chatReq.ToolChoice = v
}
body, err := json.Marshal(chatReq)
if err != nil {
return fmt.Errorf("vllm adapter: marshal request: %w", err)
}
textToolFallback := false
v.logger.Info("vllm adapter executing",
zap.String("run_id", spec.RunID),
zap.String("target", model),
zap.String("endpoint", v.endpoint),
)
resp, err := v.doChatCompletion(ctx, body)
if err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm request failed: %v", err))
return fmt.Errorf("vllm adapter: request: %w", err)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
msg := readLimited(resp.Body, 4096)
_ = resp.Body.Close()
if isAutoToolChoiceUnsupportedError(msg) {
if forced, ok := forcedToolChoiceForSingleTool(chatReq.Tools); ok {
chatReq.ToolChoice = forced
body, err = json.Marshal(chatReq)
if err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm retry marshal failed: %v", err))
return fmt.Errorf("vllm adapter: retry marshal: %w", err)
}
resp, err = v.doChatCompletion(ctx, body)
if err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm retry request failed: %v", err))
return fmt.Errorf("vllm adapter: retry request: %w", err)
}
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
goto streamResponse
}
msg = readLimited(resp.Body, 4096)
_ = resp.Body.Close()
}
}
if fallbackReq, ok := textToolFallbackChatRequest(chatReq, msg); ok {
body, err = json.Marshal(fallbackReq)
if err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm text tool fallback marshal failed: %v", err))
return fmt.Errorf("vllm adapter: text tool fallback marshal: %w", err)
}
resp, err = v.doChatCompletion(ctx, body)
if err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm text tool fallback request failed: %v", err))
return fmt.Errorf("vllm adapter: text tool fallback request: %w", err)
}
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
textToolFallback = true
goto streamResponse
}
msg = readLimited(resp.Body, 4096)
_ = resp.Body.Close()
}
if msg == "" {
msg = resp.Status
}
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm returned %s: %s", resp.Status, msg))
return fmt.Errorf("vllm adapter: non-2xx response: %s", resp.Status)
}
streamResponse:
defer resp.Body.Close()
scanner := bufio.NewScanner(resp.Body)
outputTokens := 0
finishReason := ""
var usage *runtime.UsageStats
var toolCalls openAIToolCallAccumulator
for scanner.Scan() {
line := scanner.Text()
if !strings.HasPrefix(line, "data: ") {
continue
}
payload := strings.TrimPrefix(line, "data: ")
if payload == "[DONE]" {
return sink.Emit(ctx, completeEvent(spec.RunID, finishReason, usage, outputTokens, toolCalls.ToolCalls(), textToolFallback))
}
var chunk vllmChatChunk
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
continue
}
if len(chunk.Choices) == 0 {
continue
}
if chunk.Usage != nil {
usage = &runtime.UsageStats{
InputTokens: chunk.Usage.PromptTokens,
OutputTokens: chunk.Usage.CompletionTokens,
}
}
choice := chunk.Choices[0]
if choice.FinishReason != nil && *choice.FinishReason != "" {
finishReason = *choice.FinishReason
}
if len(choice.Delta.ToolCalls) > 0 {
toolCalls.AddDelta(choice.Delta.ToolCalls)
}
content := choice.Delta.Content
if content != "" {
outputTokens += len(strings.Fields(content))
if err := sink.Emit(ctx, runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeDelta,
Delta: content,
Timestamp: time.Now(),
}); err != nil {
return err
}
}
}
if err := scanner.Err(); err != nil {
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("vllm stream scan failed: %v", err))
return fmt.Errorf("vllm adapter: scan stream: %w", err)
}
_ = emitError(ctx, sink, spec.RunID, "vllm stream ended without [DONE]")
return fmt.Errorf("vllm adapter: stream ended without [DONE]")
}
// ProbeProvider checks the availability of the vLLM (or OpenAI-compatible, e.g., SGLang) endpoint
// and the presence of the target model using the OpenAI-compatible /v1/models endpoint.
func (v *Vllm) ProbeProvider(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
targets, err := v.fetchTargets(probeCtx)
result := runtime.ProviderProbeResult{
AdapterName: Name,
InstanceKey: v.instanceName,
Target: target,
}
if err != nil {
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
result.Detail = err.Error()
return result, nil
}
result.Targets = targets
if target == "" {
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
return result, nil
}
found := false
for _, t := range targets {
if t == target {
found = true
break
}
}
if found {
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
} else {
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
result.Detail = fmt.Sprintf("target model %q not found in provider models", target)
}
return result, nil
}
func (v *Vllm) fetchTargets(ctx context.Context) ([]string, error) {
if v.endpoint == "" {
return nil, fmt.Errorf("vllm adapter: endpoint is required")
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(v.endpoint, "/v1/models"), nil)
if err != nil {
return nil, fmt.Errorf("build request: %w", err)
}
resp, err := v.client.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("status code %d", resp.StatusCode)
}
var models vllmModelsResponse
if err := json.NewDecoder(resp.Body).Decode(&models); err != nil {
return nil, fmt.Errorf("decode response: %w", err)
}
out := make([]string, 0, len(models.Data))
for _, model := range models.Data {
if model.ID != "" {
out = append(out, model.ID)
}
}
return out, nil
}
func messagesFromInput(input map[string]any) []vllmMessage {
if raw, ok := input["messages"].([]any); ok {
out := make([]vllmMessage, 0, len(raw))
for _, item := range raw {
m, ok := item.(map[string]any)
if !ok {
continue
}
role, _ := m["role"].(string)
content, _ := m["content"].(string)
toolCallID, _ := m["tool_call_id"].(string)
toolCalls := anySlice(m["tool_calls"])
role = strings.TrimSpace(role)
content = strings.TrimSpace(content)
toolCallID = strings.TrimSpace(toolCallID)
if role == "" || (content == "" && len(toolCalls) == 0) {
continue
}
out = append(out, vllmMessage{
Role: role,
Content: content,
ToolCalls: toolCalls,
ToolCallID: toolCallID,
})
}
if len(out) > 0 {
return out
}
}
if prompt := strings.TrimSpace(stringInput(input, "prompt")); prompt != "" {
return []vllmMessage{{Role: "user", Content: prompt}}
}
return nil
}
func emitError(ctx context.Context, sink runtime.EventSink, runID, msg string) error {
return sink.Emit(ctx, runtime.RuntimeEvent{
RunID: runID,
Type: runtime.EventTypeError,
Error: msg,
Timestamp: time.Now(),
})
}
func (v *Vllm) doChatCompletion(ctx context.Context, body []byte) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(v.endpoint, "/v1/chat/completions"), bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("build request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
return v.client.Do(req)
}
func isAutoToolChoiceUnsupportedError(errorBody string) bool {
errorBody = strings.ToLower(errorBody)
return strings.Contains(errorBody, "auto") &&
strings.Contains(errorBody, "tool choice") &&
isToolCallingUnsupportedError(errorBody)
}
func isToolCallingUnsupportedError(errorBody string) bool {
errorBody = strings.ToLower(errorBody)
return strings.Contains(errorBody, "enable-auto-tool-choice") ||
strings.Contains(errorBody, "tool-call-parser")
}
func forcedToolChoiceForSingleTool(tools any) (map[string]any, bool) {
items, ok := tools.([]any)
if !ok || len(items) != 1 {
return nil, false
}
tool, ok := items[0].(map[string]any)
if !ok {
return nil, false
}
name := ""
if fn, ok := tool["function"].(map[string]any); ok {
name, _ = fn["name"].(string)
}
if name == "" {
name, _ = tool["name"].(string)
}
name = strings.TrimSpace(name)
if name == "" {
return nil, false
}
return map[string]any{
"type": "function",
"function": map[string]any{
"name": name,
},
}, true
}
func textToolFallbackChatRequest(req vllmChatRequest, errorBody string) (vllmChatRequest, bool) {
if !isToolCallingUnsupportedError(errorBody) {
return vllmChatRequest{}, false
}
instruction, ok := textToolFallbackInstruction(req.Tools)
if !ok {
return vllmChatRequest{}, false
}
next := req
next.Tools = nil
next.ToolChoice = nil
next.Messages = prependTextToolFallbackInstruction(req.Messages, instruction)
return next, true
}
func prependTextToolFallbackInstruction(messages []vllmMessage, instruction string) []vllmMessage {
system := vllmMessage{Role: "system", Content: instruction}
out := make([]vllmMessage, 0, len(messages)+1)
for _, msg := range messages {
if strings.EqualFold(strings.TrimSpace(msg.Role), "system") {
system.Content = joinTextToolFallbackSystemContent(system.Content, msg.Content)
continue
}
out = append(out, msg)
}
return append([]vllmMessage{system}, out...)
}
func joinTextToolFallbackSystemContent(first, next string) string {
first = strings.TrimSpace(first)
next = strings.TrimSpace(next)
switch {
case first == "":
return next
case next == "":
return first
default:
return first + "\n\n" + next
}
}
func textToolFallbackInstruction(tools any) (string, bool) {
items, ok := anyItems(tools)
if !ok || len(items) == 0 {
return "", false
}
encoded, err := json.Marshal(items)
if err != nil {
return "", false
}
return "Tool calls must be emitted as plain text because this backend does not support native OpenAI tool calling. When a tool is needed, respond with exactly one tool call and no markdown:\n" +
"<tool_call>\n<function=TOOL_NAME>\n<parameter=PARAMETER_NAME>JSON_VALUE</parameter>\n</function>\n</tool_call>\n" +
"run_commands executes from the client workspace root. Do not prepend cd to an absolute workspace path unless the user explicitly asks to operate in a different directory; prefer current-workspace commands such as git status.\n" +
"Use valid JSON for each parameter value and follow the supplied parameter schema. Available tools JSON: " + string(encoded), true
}
func anyItems(value any) ([]any, bool) {
switch v := value.(type) {
case []any:
return v, true
default:
encoded, err := json.Marshal(v)
if err != nil {
return nil, false
}
var out []any
if err := json.Unmarshal(encoded, &out); err != nil {
return nil, false
}
return out, true
}
}
func completeEvent(runID, finishReason string, usage *runtime.UsageStats, outputTokens int, toolCalls []any, textToolFallback bool) runtime.RuntimeEvent {
if usage == nil {
usage = &runtime.UsageStats{OutputTokens: outputTokens}
}
var metadata map[string]string
if finishReason != "" {
metadata = map[string]string{"finish_reason": finishReason}
}
if textToolFallback {
if metadata == nil {
metadata = make(map[string]string, 2)
}
metadata[runtimeMetadataOpenAITextToolFallback] = "true"
}
if len(toolCalls) > 0 {
if metadata == nil {
metadata = make(map[string]string, 2)
}
if finishReason == "" {
metadata["finish_reason"] = "tool_calls"
}
if encoded, err := json.Marshal(toolCalls); err == nil {
metadata[runtimeMetadataOpenAIToolCalls] = string(encoded)
}
}
return runtime.RuntimeEvent{
RunID: runID,
Type: runtime.EventTypeComplete,
Message: "vllm chat complete",
Usage: usage,
Metadata: metadata,
Timestamp: time.Now(),
}
}
func anySlice(v any) []any {
if items, ok := v.([]any); ok {
return items
}
return nil
}
func stringInput(input map[string]any, key string) string {
if input == nil {
return ""
}
if v, ok := input[key].(string); ok {
return v
}
return ""
}
func joinURL(baseURL, path string) string {
u, err := url.Parse(baseURL)
if err != nil {
return strings.TrimRight(baseURL, "/") + path
}
u.Path = strings.TrimRight(u.Path, "/") + path
return u.String()
}
func readLimited(r io.Reader, limit int64) string {
b, _ := io.ReadAll(io.LimitReader(r, limit))
return strings.TrimSpace(string(b))
}
type vllmChatRequest struct {
Model string `json:"model"`
Messages []vllmMessage `json:"messages"`
Stream bool `json:"stream"`
Tools any `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
}
type vllmMessage struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []any `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
type vllmChatChunk struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
ToolCalls []any `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
Usage *struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
} `json:"usage"`
}
type openAIToolCallAccumulator struct {
calls map[int]map[string]any
order []int
}
func (a *openAIToolCallAccumulator) AddDelta(raw []any) {
for fallbackIndex, item := range raw {
call, ok := item.(map[string]any)
if !ok {
continue
}
index := fallbackIndex
if parsed, ok := numericIndex(call["index"]); ok {
index = parsed
}
dst := a.ensure(index)
for key, value := range call {
switch key {
case "index":
continue
case "function":
mergeToolCallFunction(dst, value)
default:
if !emptyToolCallValue(value) {
dst[key] = value
}
}
}
}
}
func (a *openAIToolCallAccumulator) ToolCalls() []any {
if len(a.order) == 0 {
return nil
}
out := make([]any, 0, len(a.order))
for _, index := range a.order {
call := a.calls[index]
if len(call) == 0 {
continue
}
copyCall := make(map[string]any, len(call))
for key, value := range call {
copyCall[key] = value
}
out = append(out, copyCall)
}
return out
}
func (a *openAIToolCallAccumulator) ensure(index int) map[string]any {
if a.calls == nil {
a.calls = make(map[int]map[string]any)
}
if call, ok := a.calls[index]; ok {
return call
}
call := make(map[string]any)
a.calls[index] = call
a.order = append(a.order, index)
return call
}
func mergeToolCallFunction(call map[string]any, value any) {
fn, ok := value.(map[string]any)
if !ok {
return
}
current, _ := call["function"].(map[string]any)
if current == nil {
current = make(map[string]any)
call["function"] = current
}
for key, item := range fn {
if key == "arguments" {
if part, ok := item.(string); ok {
current["arguments"] = currentString(current["arguments"]) + part
continue
}
}
if !emptyToolCallValue(item) {
current[key] = item
}
}
}
func numericIndex(value any) (int, bool) {
switch v := value.(type) {
case int:
return v, true
case float64:
return int(v), true
default:
return 0, false
}
}
func emptyToolCallValue(value any) bool {
if value == nil {
return true
}
if text, ok := value.(string); ok {
return text == ""
}
return false
}
func currentString(value any) string {
if text, ok := value.(string); ok {
return text
}
return ""
}
type vllmModelsResponse struct {
Data []struct {
ID string `json:"id"`
} `json:"data"`
}