394 lines
11 KiB
Go
394 lines
11 KiB
Go
package openai_compat
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"iop/apps/node/internal/runtime"
|
|
)
|
|
|
|
const (
|
|
runtimeMetadataOpenAIToolCalls = "openai_tool_calls"
|
|
runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
|
|
streamTraceMetadataKey = "iop_trace_stream"
|
|
streamTraceEnvKey = "IOP_OPENAI_COMPAT_TRACE_STREAM"
|
|
)
|
|
|
|
// chatStreamSession owns the accumulated stream state (usage, finish reason,
|
|
// tool-call accumulator, counted output tokens, and trace sequence) while the
|
|
// SSE payloads are decoded into RuntimeEvents.
|
|
type chatStreamSession struct {
|
|
adapter *Adapter
|
|
spec runtime.ExecutionSpec
|
|
textToolFallback bool
|
|
traceStream bool
|
|
traceSeq int
|
|
outputTokens int
|
|
finishReason string
|
|
usage *runtime.UsageStats
|
|
toolCalls openAIToolCallAccumulator
|
|
}
|
|
|
|
// consumeStream decodes the chat completions 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 (a *Adapter) consumeStream(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink, resp *http.Response, textToolFallback bool) error {
|
|
defer resp.Body.Close()
|
|
session := &chatStreamSession{
|
|
adapter: a,
|
|
spec: spec,
|
|
textToolFallback: textToolFallback,
|
|
traceStream: openAICompatTraceStreamEnabled(spec.Metadata),
|
|
}
|
|
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("openai_compat stream scan failed: %v", err))
|
|
return fmt.Errorf("openai_compat adapter: scan stream: %w", err)
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, "openai_compat stream ended without [DONE]")
|
|
return fmt.Errorf("openai_compat 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 *chatStreamSession) emitComplete(ctx context.Context, sink runtime.EventSink) error {
|
|
if s.traceStream {
|
|
s.adapter.logger.Info("openai_compat provider stream done",
|
|
zap.String("run_id", s.spec.RunID),
|
|
zap.Int("seq", s.traceSeq+1),
|
|
)
|
|
}
|
|
return sink.Emit(ctx, completeEvent(s.spec.RunID, s.finishReason, s.usage, s.outputTokens, s.toolCalls.ToolCalls(), s.textToolFallback))
|
|
}
|
|
|
|
// applyUsage records the usage stats from a chunk when present.
|
|
func (s *chatStreamSession) applyUsage(chunk chatChunk) {
|
|
if chunk.Usage == nil {
|
|
return
|
|
}
|
|
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
|
|
}
|
|
}
|
|
|
|
// handlePayload decodes one SSE data payload, updating usage/finish/tool-call
|
|
// state and emitting reasoning and content deltas. A malformed chunk or a
|
|
// choice-less chunk is skipped, matching the prior continue semantics.
|
|
func (s *chatStreamSession) handlePayload(ctx context.Context, sink runtime.EventSink, payload string) error {
|
|
a := s.adapter
|
|
s.traceSeq++
|
|
if s.traceStream {
|
|
a.logger.Info("openai_compat provider stream raw chunk",
|
|
zap.String("run_id", s.spec.RunID),
|
|
zap.Int("seq", s.traceSeq),
|
|
zap.Int("raw_len", len(payload)),
|
|
zap.Bool("raw_has_open_think", hasOpenThinkTag(payload)),
|
|
zap.Bool("raw_has_close_think", hasCloseThinkTag(payload)),
|
|
zap.String("raw_preview", tracePreview(payload, 4000)),
|
|
)
|
|
}
|
|
var chunk chatChunk
|
|
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
|
if s.traceStream {
|
|
a.logger.Warn("openai_compat provider stream raw chunk ignored",
|
|
zap.String("run_id", s.spec.RunID),
|
|
zap.Int("seq", s.traceSeq),
|
|
zap.Error(err),
|
|
)
|
|
}
|
|
return nil
|
|
}
|
|
s.applyUsage(chunk)
|
|
if len(chunk.Choices) == 0 {
|
|
return nil
|
|
}
|
|
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 s.traceStream {
|
|
a.logger.Info("openai_compat provider stream parsed chunk",
|
|
zap.String("run_id", s.spec.RunID),
|
|
zap.Int("seq", s.traceSeq),
|
|
zap.String("finish_reason", s.finishReason),
|
|
zap.Int("tool_call_delta_count", len(choice.Delta.ToolCalls)),
|
|
zap.Int("content_len", len(choice.Delta.Content)),
|
|
zap.Bool("content_has_open_think", hasOpenThinkTag(choice.Delta.Content)),
|
|
zap.Bool("content_has_close_think", hasCloseThinkTag(choice.Delta.Content)),
|
|
zap.String("content_preview", tracePreview(choice.Delta.Content, 2000)),
|
|
zap.Int("reasoning_len", len(reasoning)),
|
|
zap.Bool("reasoning_has_open_think", hasOpenThinkTag(reasoning)),
|
|
zap.Bool("reasoning_has_close_think", hasCloseThinkTag(reasoning)),
|
|
zap.String("reasoning_preview", tracePreview(reasoning, 2000)),
|
|
)
|
|
}
|
|
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
|
|
}
|
|
}
|
|
if content := choice.Delta.Content; 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
|
|
}
|
|
|
|
func openAICompatTraceStreamEnabled(metadata map[string]string) bool {
|
|
if traceBool(os.Getenv(streamTraceEnvKey)) {
|
|
return true
|
|
}
|
|
if metadata == nil {
|
|
return false
|
|
}
|
|
return traceBool(metadata[streamTraceMetadataKey])
|
|
}
|
|
|
|
func traceBool(v string) bool {
|
|
switch strings.ToLower(strings.TrimSpace(v)) {
|
|
case "1", "true", "yes", "on":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func tracePreview(s string, max int) string {
|
|
if max <= 0 || len(s) <= max {
|
|
return s
|
|
}
|
|
return s[:max] + "...[truncated]"
|
|
}
|
|
|
|
func hasOpenThinkTag(s string) bool {
|
|
return strings.Contains(strings.ToLower(s), "<think")
|
|
}
|
|
|
|
func hasCloseThinkTag(s string) bool {
|
|
return strings.Contains(strings.ToLower(s), "</think")
|
|
}
|
|
|
|
// joinOpenAIPath appends an OpenAI-compatible path to the endpoint. When the
|
|
// endpoint already targets the /v1 root it strips the leading /v1 from the path
|
|
// so a configured ".../v1" endpoint does not produce ".../v1/v1/...".
|
|
func joinOpenAIPath(baseURL, path string) string {
|
|
u, err := url.Parse(baseURL)
|
|
if err != nil {
|
|
return strings.TrimRight(baseURL, "/") + path
|
|
}
|
|
basePath := strings.TrimRight(u.Path, "/")
|
|
if strings.HasSuffix(basePath, "/v1") && strings.HasPrefix(path, "/v1/") {
|
|
path = strings.TrimPrefix(path, "/v1")
|
|
}
|
|
u.Path = basePath + 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 chatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
ToolCalls []any `json:"tool_calls,omitempty"`
|
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
|
}
|
|
|
|
type chatChunk struct {
|
|
Choices []struct {
|
|
Delta chatChunkDelta `json:"delta"`
|
|
FinishReason *string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
Usage *struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
PromptTokensDetails *struct {
|
|
CachedTokens int `json:"cached_tokens"`
|
|
} `json:"prompt_tokens_details"`
|
|
CompletionTokensDetails *struct {
|
|
ReasoningTokens int `json:"reasoning_tokens"`
|
|
} `json:"completion_tokens_details"`
|
|
} `json:"usage"`
|
|
}
|
|
|
|
type chatChunkDelta struct {
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
Reasoning string `json:"reasoning"`
|
|
ToolCalls []any `json:"tool_calls"`
|
|
}
|
|
|
|
func (d chatChunkDelta) ReasoningText() string {
|
|
if d.ReasoningContent != "" {
|
|
return d.ReasoningContent
|
|
}
|
|
return d.Reasoning
|
|
}
|
|
|
|
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 modelsResponse struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
}
|