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

package vllm

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
)

const (
	runtimeMetadataOpenAIToolCalls        = "openai_tool_calls"
	runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
)

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 {
		if tools, ok := v.([]any); ok {
			chatReq.Tools = tools
		}
	}
	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,
			}
			if d := chunk.Usage.PromptTokensDetails; d != nil {
				usage.CachedInputTokens = d.CachedTokens
			}
			if d := chunk.Usage.CompletionTokensDetails; d != nil {
				usage.ReasoningTokens = d.ReasoningTokens
			}
		}
		choice := chunk.Choices[0]
		if choice.FinishReason != nil && *choice.FinishReason != "" {
			finishReason = *choice.FinishReason
		}
		if len(choice.Delta.ToolCalls) > 0 {
			toolCalls.AddDelta(choice.Delta.ToolCalls)
		}
		reasoning := choice.Delta.ReasoningText()
		if reasoning != "" {
			if err := sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeReasoningDelta,
				Delta:     reasoning,
				Timestamp: time.Now(),
			}); err != nil {
				return err
			}
		}

		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.