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

package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
)

type providerModelRewriter struct {
	streaming bool
	model     string
	pending   []byte
}

func newProviderModelRewriter(streaming bool, model string) *providerModelRewriter {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	return &providerModelRewriter{streaming: streaming, model: model}
}

func (r *providerModelRewriter) AppendStream(chunk []byte) []byte {
	if r == nil || !r.streaming || len(chunk) == 0 {
		return chunk
	}
	r.pending = append(r.pending, chunk...)
	var out bytes.Buffer
	for {
		idx := bytes.IndexByte(r.pending, '\n')
		if idx < 0 {
			break
		}
		line := r.pending[:idx+1]
		out.Write(rewriteProviderSSEModelLine(line, r.model))
		r.pending = r.pending[idx+1:]
	}
	return out.Bytes()
}

func (r *providerModelRewriter) FlushStream() []byte {
	if r == nil || len(r.pending) == 0 {
		return nil
	}
	pending := r.pending
	r.pending = nil
	return rewriteProviderSSEModelLine(pending, r.model)
}

func (r *providerModelRewriter) RewriteComplete(body []byte) []byte {
	if r == nil || len(body) == 0 {
		return body
	}
	return rewriteProviderJSONModel(body, r.model)
}

func rewriteProviderSSEModelLine(line []byte, model string) []byte {
	body, ending := splitLineEnding(line)
	prefix, payload, ok := bytes.Cut(body, []byte(":"))
	if !ok || strings.TrimSpace(string(prefix)) != "data" {
		return line
	}
	payload = bytes.TrimSpace(payload)
	if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
		return line
	}
	rewritten := rewriteProviderJSONModel(payload, model)
	if bytes.Equal(rewritten, payload) {
		return line
	}
	out := make([]byte, 0, len("data: ")+len(rewritten)+len(ending))
	out = append(out, "data: "...)
	out = append(out, rewritten...)
	out = append(out, ending...)
	return out
}

func splitLineEnding(line []byte) ([]byte, []byte) {
	if len(line) == 0 || line[len(line)-1] != '\n' {
		return line, nil
	}
	if len(line) >= 2 && line[len(line)-2] == '\r' {
		return line[:len(line)-2], []byte("\r\n")
	}
	return line[:len(line)-1], []byte("\n")
}

func rewriteProviderJSONModel(body []byte, model string) []byte {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(body, &raw); err != nil {
		return body
	}
	if _, ok := raw["model"]; !ok {
		return body
	}
	modelJSON, err := json.Marshal(model)
	if err != nil {
		return body
	}
	raw["model"] = modelJSON
	rewritten, err := json.Marshal(raw)
	if err != nil {
		return body
	}
	return rewritten
}

// rewriteChatCompletionModel replaces only the model field of the caller's
// original request JSON so the provider receives its served model name. The
// rest of the caller payload is forwarded without IOP rewriting.
func rewriteChatCompletionModel(rawBody []byte, target string, req chatCompletionRequest) ([]byte, error) {
	if strings.TrimSpace(target) == "" && req.MaxTokens == nil && req.ThinkingTokenBudget == nil && req.Think == nil {
		return rawBody, nil
	}
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(rawBody, &raw); err != nil {
		return nil, fmt.Errorf("invalid JSON request")
	}
	if strings.TrimSpace(target) != "" {
		modelJSON, err := json.Marshal(target)
		if err != nil {
			return nil, err
		}
		raw["model"] = modelJSON
	}
	if req.MaxTokens != nil {
		maxTokensJSON, err := json.Marshal(*req.MaxTokens)
		if err != nil {
			return nil, err
		}
		raw["max_tokens"] = maxTokensJSON
		delete(raw, "max_completion_tokens")
	}
	if req.ThinkingTokenBudget != nil {
		budgetJSON, err := json.Marshal(*req.ThinkingTokenBudget)
		if err != nil {
			return nil, err
		}
		raw["thinking_token_budget"] = budgetJSON
	}
	if req.Think != nil {
		thinkJSON, err := json.Marshal(*req.Think)
		if err != nil {
			return nil, err
		}
		raw["think"] = thinkJSON
	}
	return json.Marshal(raw)
}

// rewriteResponsesModel replaces only the model field of the caller's original
// /v1/responses request JSON so the provider receives its served model name.
// Every other field (input, instructions, tools, max_output_tokens, and any
// Codex/Responses-specific field) is forwarded without IOP rewriting. An empty
// target leaves the body untouched; invalid JSON is rejected.
func rewriteResponsesModel(rawBody []byte, target string) ([]byte, error) {
	if strings.TrimSpace(target) == "" {
		return rawBody, nil
	}
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(rawBody, &raw); err != nil {
		return nil, fmt.Errorf("invalid JSON request")
	}
	modelJSON, err := json.Marshal(target)
	if err != nil {
		return nil, err
	}
	raw["model"] = modelJSON
	return json.Marshal(raw)
}