ai-gateway-go/internal/gateway/routing.go

package gateway

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
)

type ModelRouteQuery struct {
	ProviderCode string
	Model        string
	Endpoint     string
	APIKeyID     string
	TenantID     string
	Seed         string
}

type ModelRouteResult struct {
	ResolvedAdapter
	TargetModel string
	Matched     bool
	Known       bool
}

type ModelRouteResolver interface {
	ModelRoutingEnabled() bool
	ResolveModelRoute(ModelRouteQuery) (ModelRouteResult, error)
}

type modelRequest struct {
	body  map[string]json.RawMessage
	model string
}

func readModelRequest(request *http.Request, maxBody int64) (modelRequest, error) {
	if request.Body == nil || request.Method == http.MethodGet {
		return modelRequest{}, nil
	}
	body, err := io.ReadAll(io.LimitReader(request.Body, maxBody+1))
	if err != nil {
		return modelRequest{}, err
	}
	_ = request.Body.Close()
	if int64(len(body)) > maxBody {
		return modelRequest{}, errRequestBodyTooLarge
	}
	restoreRequestBody(request, body)
	var payload map[string]json.RawMessage
	if json.Unmarshal(body, &payload) != nil {
		return modelRequest{}, nil
	}
	var model string
	_ = json.Unmarshal(payload["model"], &model)
	return modelRequest{body: payload, model: model}, nil
}

func (m modelRequest) rewrite(request *http.Request, target string) error {
	if len(m.body) == 0 || target == "" || target == m.model {
		return nil
	}
	encodedModel, _ := json.Marshal(target)
	m.body["model"] = encodedModel
	body, err := json.Marshal(m.body)
	if err != nil {
		return err
	}
	restoreRequestBody(request, body)
	return nil
}

func restoreRequestBody(request *http.Request, body []byte) {
	request.Body = io.NopCloser(bytes.NewReader(body))
	request.ContentLength = int64(len(body))
	request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
}