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

140 lines
4.1 KiB
Go

package openai
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
edgeservice "iop/apps/edge/internal/service"
)
const (
runtimeMetadataOpenAIToolCalls = "openai_tool_calls"
runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
)
// errRunTimedOut is returned by collectRunResult and the direct streaming
// loop when handle.WaitTimeout() elapses without a terminal run event.
var errRunTimedOut = errors.New("run timed out")
// isCancelWorthyRunError reports whether err means the HTTP caller gave up
// (context cancellation/deadline or a WaitTimeout expiry) before the run
// reached a terminal state, so Edge should propagate CancelRun to Node.
// Terminal run outcomes (complete/error/cancelled events, closed stream,
// node disconnect) are not cancel-worthy: Node already reached a terminal
// state or is unreachable, so sending CancelRun would be redundant or futile.
func isCancelWorthyRunError(err error) bool {
return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, errRunTimedOut)
}
func collectRunResult(ctx context.Context, stream edgeservice.RunStream, timeout time.Duration) (string, string, string, []any, *openAIUsage, bool, error) {
if stream.Events == nil {
return "", "", "", nil, nil, false, fmt.Errorf("run stream unavailable")
}
var contentBuilder strings.Builder
var reasoningBuilder strings.Builder
finishReason := "stop"
var usage *openAIUsage
timer := time.NewTimer(timeout)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return "", "", "", nil, nil, false, ctx.Err()
case <-timer.C:
return "", "", "", nil, nil, false, errRunTimedOut
case nodeEvent, ok := <-stream.NodeEvents:
if !ok {
stream.NodeEvents = nil
continue
}
if edgeservice.IsNodeDisconnected(nodeEvent) {
return "", "", "", nil, nil, false, fmt.Errorf("node disconnected")
}
case event, ok := <-stream.Events:
if !ok {
return "", "", "", nil, nil, false, fmt.Errorf("run stream closed")
}
if event == nil {
continue
}
switch event.GetType() {
case "delta":
contentBuilder.WriteString(event.GetDelta())
case "reasoning_delta":
reasoningBuilder.WriteString(event.GetDelta())
case "complete":
if reason := event.GetMetadata()["finish_reason"]; reason != "" {
finishReason = reason
}
toolCalls := toolCallsFromRunMetadata(event.GetMetadata())
if len(toolCalls) > 0 {
finishReason = "tool_calls"
}
if u := event.GetUsage(); u != nil {
usage = &openAIUsage{
PromptTokens: int(u.GetInputTokens()),
CompletionTokens: int(u.GetOutputTokens()),
TotalTokens: int(u.GetInputTokens() + u.GetOutputTokens()),
ReasoningTokens: int(u.GetReasoningTokens()),
CachedInputTokens: int(u.GetCachedInputTokens()),
}
}
return contentBuilder.String(), reasoningBuilder.String(), finishReason, toolCalls, usage, isTextToolFallback(event.GetMetadata()), nil
case "error", "cancelled":
msg := event.GetError()
if msg == "" {
msg = event.GetMessage()
}
if msg == "" {
msg = "run failed"
}
return "", "", "", nil, nil, false, fmt.Errorf("%s", msg)
}
}
}
}
func toolCallsFromRunMetadata(metadata map[string]string) []any {
raw := strings.TrimSpace(metadata[runtimeMetadataOpenAIToolCalls])
if raw == "" {
return nil
}
var toolCalls []any
if err := json.Unmarshal([]byte(raw), &toolCalls); err != nil {
return nil
}
if len(toolCalls) == 0 {
return nil
}
return toolCalls
}
func isTextToolFallback(metadata map[string]string) bool {
return strings.EqualFold(strings.TrimSpace(metadata[runtimeMetadataOpenAITextToolFallback]), "true")
}
func extractToolCallNames(toolCalls []any) []string {
var names []string
for _, tc := range toolCalls {
if tc == nil {
continue
}
b, err := json.Marshal(tc)
if err != nil {
continue
}
var parsed struct {
Function struct {
Name string `json:"name"`
} `json:"function"`
}
if err := json.Unmarshal(b, &parsed); err == nil && parsed.Function.Name != "" {
names = append(names, parsed.Function.Name)
}
}
return names
}