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

package openai

import (
	"bytes"
	"encoding/json"
	iop "iop/proto/gen/iop"
	"strings"
)

// providerAssembledObservation is the Edge-internal human-readable view of a
// provider Chat Completions/Responses body assembled by providerChatAssembler.
// It feeds Edge-local logging and usage metrics only and is never written to
// the caller response; pure passthrough relays provider bytes unmodified.
type providerAssembledObservation struct {
	Content       string   `json:"content,omitempty"`
	Reasoning     string   `json:"reasoning,omitempty"`
	ToolCallNames []string `json:"tool_call_names,omitempty"`
	BodyBytes     int      `json:"body_bytes"`
}

// providerChatAssembler accumulates a human-readable view of the provider
// Chat Completions response for the assembled observation. It
// tolerates non-JSON and partial payloads: whatever cannot be parsed simply
// yields an empty summary.
type providerChatAssembler struct {
	streaming          bool
	nonStreamingParsed bool
	bodyBytes          int
	pending            []byte
	content            strings.Builder
	reasoning          strings.Builder
	toolCallNames      []string
	// usage holds provider-reported token counts observed in the passthrough
	// body. It feeds Edge-internal metrics only.
	usage usageObservation
}

// providerChatDeltaEnvelope matches the provider fields the assembler
// observes; streaming uses delta, non-streaming uses message.
type providerChatDeltaEnvelope struct {
	Content          string `json:"content"`
	ReasoningContent string `json:"reasoning_content"`
	Reasoning        string `json:"reasoning"`
	ToolCalls        []struct {
		Function struct {
			Name string `json:"name"`
		} `json:"function"`
	} `json:"tool_calls"`
}

// providerUsageEnvelope matches the OpenAI-compatible usage object from
// both Chat Completions and Responses. Chat Completions uses prompt_tokens/
// completion_tokens; Responses uses input_tokens/output_tokens. The optional
// detail objects carry reasoning and cached-input tokens for both formats.
type providerUsageEnvelope struct {
	PromptTokens        int `json:"prompt_tokens"`
	CompletionTokens    int `json:"completion_tokens"`
	InputTokens         int `json:"input_tokens"`
	OutputTokens        int `json:"output_tokens"`
	PromptTokensDetails *struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
	CompletionTokensDetails *struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"completion_tokens_details"`
	InputTokensDetails *struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"input_tokens_details"`
	OutputTokensDetails *struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"output_tokens_details"`
}

// recordUsage stores the latest provider-reported usage. The provider tunnel
// body is never mutated; this observation feeds Edge-internal metrics only
// (SDD S05/D04). Chat Completions keys (prompt_tokens/completion_tokens) take
// precedence; if absent, Responses keys (input_tokens/output_tokens) are used.
func (a *providerChatAssembler) recordUsage(u *providerUsageEnvelope) {
	if u == nil {
		return
	}
	// Prefer Chat Completions keys; fall back to Responses keys.
	if u.PromptTokens != 0 {
		a.usage.inputTokens = u.PromptTokens
		a.usage.outputTokens = u.CompletionTokens
		if d := u.PromptTokensDetails; d != nil {
			a.usage.cachedInputTokens = d.CachedTokens
		}
		if d := u.CompletionTokensDetails; d != nil {
			a.usage.reasoningTokens = d.ReasoningTokens
		}
		return
	}
	if u.InputTokens != 0 {
		a.usage.inputTokens = u.InputTokens
		a.usage.outputTokens = u.OutputTokens
		if d := u.InputTokensDetails; d != nil {
			a.usage.cachedInputTokens = d.CachedTokens
		}
		if d := u.OutputTokensDetails; d != nil {
			a.usage.reasoningTokens = d.ReasoningTokens
		}
	}
}

// recordProtoUsage stores usage carried on a provider tunnel USAGE frame. Like
// recordUsage it observes for metrics only and never alters the response body.
func (a *providerChatAssembler) recordProtoUsage(u *iop.Usage) {
	if u == nil {
		return
	}
	a.usage.inputTokens = int(u.GetInputTokens())
	a.usage.outputTokens = int(u.GetOutputTokens())
	a.usage.reasoningTokens = int(u.GetReasoningTokens())
	a.usage.cachedInputTokens = int(u.GetCachedInputTokens())
}

// usageObservation returns the observed token usage plus the observed reasoning
// character count used by auxiliary and estimated-token reasoning metrics.
func (a *providerChatAssembler) usageObservation() usageObservation {
	obs := a.usage
	obs.reasoningChars = a.reasoning.Len()
	return obs
}

func (a *providerChatAssembler) Write(chunk []byte) {
	a.bodyBytes += len(chunk)
	a.pending = append(a.pending, chunk...)
	if !a.streaming {
		return
	}
	for {
		idx := bytes.IndexByte(a.pending, '\n')
		if idx < 0 {
			return
		}
		line := strings.TrimRight(string(a.pending[:idx]), "\r")
		a.pending = a.pending[idx+1:]
		a.consumeSSELine(line)
	}
}

func (a *providerChatAssembler) consumeSSELine(line string) {
	payload, ok := strings.CutPrefix(line, "data:")
	if !ok {
		return
	}
	payload = strings.TrimSpace(payload)
	if payload == "" || payload == "[DONE]" {
		return
	}
	var chunk struct {
		Choices []struct {
			Delta providerChatDeltaEnvelope `json:"delta"`
		} `json:"choices"`
		Usage *providerUsageEnvelope `json:"usage"`
	}
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		return
	}
	for _, choice := range chunk.Choices {
		a.consumeDelta(choice.Delta)
	}
	a.recordUsage(chunk.Usage)

	// Responses streaming: nested response.usage from events like
	// response.completed. This captures provider-reported token usage from
	// the Responses API SSE payload (SDD S05/D04).
	if a.usage.inputTokens == 0 && a.usage.outputTokens == 0 {
		var respEvent struct {
			Response struct {
				Usage *providerUsageEnvelope `json:"usage"`
			} `json:"response"`
		}
		if err := json.Unmarshal([]byte(payload), &respEvent); err == nil {
			if respEvent.Response.Usage != nil {
				a.recordUsage(respEvent.Response.Usage)
			}
		}
	}
}

func (a *providerChatAssembler) consumeDelta(delta providerChatDeltaEnvelope) {
	a.content.WriteString(delta.Content)
	a.reasoning.WriteString(delta.ReasoningContent)
	a.reasoning.WriteString(delta.Reasoning)
	for _, call := range delta.ToolCalls {
		if name := strings.TrimSpace(call.Function.Name); name != "" {
			a.toolCallNames = append(a.toolCallNames, name)
		}
	}
}

func (a *providerChatAssembler) observation() *providerAssembledObservation {
	if !a.streaming && !a.nonStreamingParsed {
		a.nonStreamingParsed = true
		var resp struct {
			Choices []struct {
				Message providerChatDeltaEnvelope `json:"message"`
			} `json:"choices"`
			Usage *providerUsageEnvelope `json:"usage"`
		}
		if err := json.Unmarshal(a.pending, &resp); err == nil {
			for _, choice := range resp.Choices {
				a.consumeDelta(choice.Message)
			}
			a.recordUsage(resp.Usage)
		}
	}
	return &providerAssembledObservation{
		Content:       a.content.String(),
		Reasoning:     a.reasoning.String(),
		ToolCallNames: a.toolCallNames,
		BodyBytes:     a.bodyBytes,
	}
}