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

package vllm

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

	"go.uber.org/zap"

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

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

// Execute validates the request and emits the start event, sends the request
// with the auto-tool-choice and text-tool fallback retries, then consumes the
// SSE stream into RuntimeEvents.
func (v *Vllm) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
	exec, err := v.prepareExecution(ctx, spec, sink)
	if err != nil {
		return err
	}
	resp, textToolFallback, err := exec.requestWithFallback()
	if err != nil {
		return err
	}
	return v.consumeStream(ctx, spec, sink, resp, textToolFallback)
}

// vllmExecution holds the validated request context shared by the request and
// retry stages so Execute stays a linear orchestration without a goto.
type vllmExecution struct {
	adapter *Vllm
	ctx     context.Context
	sink    runtime.EventSink
	runID   string
	chatReq vllmChatRequest
	body    []byte
}

// prepareExecution validates the spec, emits the start event, and builds the
// marshaled chat request. Validation order and error text match the prior
// single-function Execute.
func (v *Vllm) prepareExecution(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) (*vllmExecution, error) {
	if v.endpoint == "" {
		return nil, fmt.Errorf("vllm adapter: endpoint is required")
	}
	model := strings.TrimSpace(spec.Target)
	if model == "" {
		model = stringInput(spec.Input, "model")
	}
	if model == "" {
		return nil, fmt.Errorf("vllm adapter: target/model is required")
	}
	messages := messagesFromInput(spec.Input)
	if len(messages) == 0 {
		return nil, 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 nil, err
	}

	chatReq := vllmChatRequest{
		Model:    model,
		Messages: messages,
		Stream:   true,
	}
	if val, ok := spec.Input["tools"]; ok {
		chatReq.Tools = val
	}
	if val, ok := spec.Input["tool_choice"]; ok {
		chatReq.ToolChoice = val
	}
	body, err := json.Marshal(chatReq)
	if err != nil {
		return nil, fmt.Errorf("vllm adapter: marshal request: %w", err)
	}

	v.logger.Info("vllm adapter executing",
		zap.String("run_id", spec.RunID),
		zap.String("target", model),
		zap.String("endpoint", v.endpoint),
	)

	return &vllmExecution{
		adapter: v,
		ctx:     ctx,
		sink:    sink,
		runID:   spec.RunID,
		chatReq: chatReq,
		body:    body,
	}, nil
}

// requestWithFallback sends the chat request and, on a non-2xx response, applies
// the forced-single-tool retry (only for the auto-tool-choice unsupported error)
// then the text-tool fallback retry. It returns the streaming response and
// whether the text-tool fallback was used, preserving the prior goto-based order.
func (e *vllmExecution) requestWithFallback() (*http.Response, bool, error) {
	v := e.adapter
	resp, err := v.doChatCompletion(e.ctx, e.body)
	if err != nil {
		_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm request failed: %v", err))
		return nil, false, fmt.Errorf("vllm adapter: request: %w", err)
	}
	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
		return resp, false, nil
	}

	msg := readLimited(resp.Body, 4096)
	_ = resp.Body.Close()

	if isAutoToolChoiceUnsupportedError(msg) {
		if forced, ok := forcedToolChoiceForSingleTool(e.chatReq.Tools); ok {
			e.chatReq.ToolChoice = forced
			body, err := json.Marshal(e.chatReq)
			if err != nil {
				_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm retry marshal failed: %v", err))
				return nil, false, fmt.Errorf("vllm adapter: retry marshal: %w", err)
			}
			resp, err = v.doChatCompletion(e.ctx, body)
			if err != nil {
				_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm retry request failed: %v", err))
				return nil, false, fmt.Errorf("vllm adapter: retry request: %w", err)
			}
			if resp.StatusCode >= 200 && resp.StatusCode < 300 {
				return resp, false, nil
			}
			msg = readLimited(resp.Body, 4096)
			_ = resp.Body.Close()
		}
	}

	if fallbackReq, ok := textToolFallbackChatRequest(e.chatReq, msg); ok {
		body, err := json.Marshal(fallbackReq)
		if err != nil {
			_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm text tool fallback marshal failed: %v", err))
			return nil, false, fmt.Errorf("vllm adapter: text tool fallback marshal: %w", err)
		}
		resp, err = v.doChatCompletion(e.ctx, body)
		if err != nil {
			_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm text tool fallback request failed: %v", err))
			return nil, false, fmt.Errorf("vllm adapter: text tool fallback request: %w", err)
		}
		if resp.StatusCode >= 200 && resp.StatusCode < 300 {
			return resp, true, nil
		}
		msg = readLimited(resp.Body, 4096)
		_ = resp.Body.Close()
	}

	if msg == "" {
		msg = resp.Status
	}
	_ = emitError(e.ctx, e.sink, e.runID, fmt.Sprintf("vllm returned %s: %s", resp.Status, msg))
	return nil, false, fmt.Errorf("vllm adapter: non-2xx response: %s", resp.Status)
}

// vllmStreamSession owns the accumulated stream state (usage, finish reason,
// tool-call accumulator, and counted output tokens) while SSE payloads are
// decoded into RuntimeEvents.
type vllmStreamSession struct {
	spec             runtime.ExecutionSpec
	textToolFallback bool
	outputTokens     int
	finishReason     string
	usage            *runtime.UsageStats
	toolCalls        openAIToolCallAccumulator
}

// consumeStream decodes the vLLM SSE stream, emitting reasoning and content
// deltas and, on [DONE], the completion event. The scan/emit ordering and
// terminal error handling match the prior single-function Execute.
func (v *Vllm) consumeStream(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink, resp *http.Response, textToolFallback bool) error {
	defer resp.Body.Close()
	session := &vllmStreamSession{spec: spec, textToolFallback: textToolFallback}
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		payload := strings.TrimPrefix(line, "data: ")
		if payload == "[DONE]" {
			return session.emitComplete(ctx, sink)
		}
		if err := session.handlePayload(ctx, sink, payload); 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]")
}

// emitComplete emits the terminal completion event with the accumulated finish
// reason, usage, counted output tokens, and tool calls.
func (s *vllmStreamSession) emitComplete(ctx context.Context, sink runtime.EventSink) error {
	return sink.Emit(ctx, completeEvent(s.spec.RunID, s.finishReason, s.usage, s.outputTokens, s.toolCalls.ToolCalls(), s.textToolFallback))
}

// handlePayload decodes one SSE data payload, updating usage/finish/tool-call
// state and emitting reasoning and content deltas. A malformed or choice-less
// chunk is skipped, matching the prior continue semantics.
func (s *vllmStreamSession) handlePayload(ctx context.Context, sink runtime.EventSink, payload string) error {
	var chunk vllmChatChunk
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		return nil
	}
	if len(chunk.Choices) == 0 {
		return nil
	}
	if chunk.Usage != nil {
		s.usage = &runtime.UsageStats{
			InputTokens:  chunk.Usage.PromptTokens,
			OutputTokens: chunk.Usage.CompletionTokens,
		}
		if d := chunk.Usage.PromptTokensDetails; d != nil {
			s.usage.CachedInputTokens = d.CachedTokens
		}
		if d := chunk.Usage.CompletionTokensDetails; d != nil {
			s.usage.ReasoningTokens = d.ReasoningTokens
		}
	}
	choice := chunk.Choices[0]
	if choice.FinishReason != nil && *choice.FinishReason != "" {
		s.finishReason = *choice.FinishReason
	}
	if len(choice.Delta.ToolCalls) > 0 {
		s.toolCalls.AddDelta(choice.Delta.ToolCalls)
	}
	reasoning := choice.Delta.ReasoningText()
	if reasoning != "" {
		if err := sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     s.spec.RunID,
			Type:      runtime.EventTypeReasoningDelta,
			Delta:     reasoning,
			Timestamp: time.Now(),
		}); err != nil {
			return err
		}
	}

	content := choice.Delta.Content
	if content != "" {
		s.outputTokens += len(strings.Fields(content))
		if err := sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     s.spec.RunID,
			Type:      runtime.EventTypeDelta,
			Delta:     content,
			Timestamp: time.Now(),
		}); err != nil {
			return err
		}
	}
	return nil
}