iop/apps/edge/internal/openai/run_result.go

package openai

import (
	"context"
	"fmt"
	"strings"
	"time"

	edgeservice "iop/apps/edge/internal/service"
)

func collectRunResult(ctx context.Context, stream edgeservice.RunStream, timeout time.Duration) (string, string, *openAIUsage, error) {
	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	var usage *openAIUsage
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", "", nil, ctx.Err()
		case <-timer.C:
			return "", "", nil, fmt.Errorf("run timed out")
		case nodeEvent := <-stream.NodeEvents:
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				return "", "", nil, fmt.Errorf("node disconnected")
			}
		case event := <-stream.Events:
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				contentBuilder.WriteString(event.GetDelta())
			case "reasoning_delta":
				reasoningBuilder.WriteString(event.GetDelta())
			case "complete":
				if u := event.GetUsage(); u != nil {
					usage = &openAIUsage{
						PromptTokens:     int(u.GetInputTokens()),
						CompletionTokens: int(u.GetOutputTokens()),
						TotalTokens:      int(u.GetInputTokens() + u.GetOutputTokens()),
					}
				}
				return contentBuilder.String(), reasoningBuilder.String(), usage, nil
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				return "", "", nil, fmt.Errorf("%s", msg)
			}
		}
	}
}