ModelRouter/internal/provider/provider.go

// Package provider binds a configured source + Lua adapter and performs the
// HTTP call / stream / image generation against the upstream LLM, with
// per-source concurrency limiting and availability backoff.
package provider

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
	"llmsproxy/internal/types"
)

// health tracks availability with exponential backoff.
type health struct {
	failCount        int
	unavailableUntil time.Time
	permanent        bool
}

func (h *health) reset() { h.failCount = 0; h.unavailableUntil = time.Time{}; h.permanent = false }

func (h *health) available() bool {
	if h.permanent {
		return false
	}
	return time.Now().After(h.unavailableUntil)
}

func (h *health) backoff() {
	h.failCount++
	cooldown := 5 * time.Second * time.Duration(1<<(h.failCount-1))
	if cooldown > 30*time.Minute {
		cooldown = 30 * time.Minute
	}
	h.unavailableUntil = time.Now().Add(cooldown)
}

func (h *health) markPermanent() {
	h.permanent = true
	h.unavailableUntil = time.Time{}
}

// Provider is a single configured upstream LLM source.
type Provider struct {
	cfg     config.Source
	vm      *lua.VM
	adapter string
	client  *http.Client

	mu     sync.Mutex
	sem    chan struct{}
	health health
}

func New(cfg config.Source, vm *lua.VM) *Provider {
	p := &Provider{
		cfg:     cfg,
		vm:      vm,
		adapter: cfg.Adapter,
		client:  &http.Client{Timeout: cfg.Timeout},
		sem:     make(chan struct{}, cfg.MaxConcurrent),
	}
	if cfg.MaxConcurrent <= 0 {
		p.sem = nil
	}
	return p
}

func (p *Provider) Name() string            { return p.cfg.Name }
func (p *Provider) Adapter() string         { return p.cfg.Adapter }
func (p *Provider) MaxConcurrent() int      { return p.cfg.MaxConcurrent }
func (p *Provider) Config() *config.Source  { return &p.cfg }

// Models returns the model ids exposed by this source.
func (p *Provider) Models() []string {
	out := make([]string, 0, len(p.cfg.Models))
	for _, m := range p.cfg.Models {
		out = append(out, m.ID)
	}
	return out
}

// ModelByID returns the model definition if owned by this source.
func (p *Provider) ModelByID(id string) *config.Model {
	for i := range p.cfg.Models {
		if p.cfg.Models[i].ID == id {
			return &p.cfg.Models[i]
		}
	}
	return nil
}

// ModelFor resolves the model name this provider should send upstream.
// If the requested model is not owned by this provider (e.g. an AUTO chain
// fallback), it returns this provider's highest-priority chat model instead.
func (p *Provider) ModelFor(reqModel string) string {
	if reqModel == "" || isAutoID(reqModel) {
		return p.bestChatModel()
	}
	if p.ModelByID(reqModel) != nil {
		return reqModel
	}
	return p.bestChatModel()
}

// bestChatModel returns the highest-priority chat-kind model of this source.
func (p *Provider) bestChatModel() string {
	bestID, bestPrio := "", -1
	for _, m := range p.cfg.Models {
		if m.Kind != "" && m.Kind != "chat" {
			continue
		}
		if m.Priority > bestPrio {
			bestPrio = m.Priority
			bestID = m.ID
		}
	}
	if bestID == "" && len(p.cfg.Models) > 0 {
		bestID = p.cfg.Models[0].ID
	}
	return bestID
}

// IsAutoID reports whether s is an AUTO routing placeholder.
func isAutoID(s string) bool {
	s = strings.TrimSpace(s)
	return s == "" || strings.EqualFold(s, "AUTO")
}

// Endpoint resolves the upstream chat path.
func (p *Provider) Endpoint() string {
	if p.cfg.Endpoint != "" {
		return p.cfg.Endpoint
	}
	if ep := p.vm.Endpoint(p.adapter); ep != "" {
		return ep
	}
	return "/chat/completions"
}

// ImageEndpoint resolves the upstream image-generation path.
func (p *Provider) ImageEndpoint() string {
	if p.cfg.ImageEndpoint != "" {
		return p.cfg.ImageEndpoint
	}
	if ep := p.vm.Endpoint(p.adapter + "_image"); ep != "" {
		return ep
	}
	return "/v1/images/generations"
}

func (p *Provider) URL() string {
	return strings.TrimRight(p.cfg.BaseURL, "/") + p.Endpoint()
}

func (p *Provider) ImageURL() string {
	return strings.TrimRight(p.cfg.BaseURL, "/") + p.ImageEndpoint()
}

// ---- availability ----

func (p *Provider) Available() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.health.available()
}

// ReportStatus records an upstream HTTP status for backoff decisions.
func (p *Provider) ReportStatus(code int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if code == 401 || code == 403 {
		p.health.markPermanent()
		return
	}
	if code >= 500 || code == 429 {
		p.health.backoff()
	}
}

func (p *Provider) reportError() {
	p.mu.Lock()
	p.health.backoff()
	p.mu.Unlock()
}

func (p *Provider) reportOK() {
	p.mu.Lock()
	p.health.reset()
	p.mu.Unlock()
}

// ---- concurrency limiting ----

// Acquire waits for a free concurrency slot (bounded by cfg.QueueTimeout),
// or context cancel. The HTTP call itself is not truncated.
func (p *Provider) Acquire(ctx context.Context) error {
	if p.sem == nil {
		return nil
	}
	var qCtx context.Context
	var cancel context.CancelFunc
	if p.cfg.QueueTimeout > 0 {
		qCtx, cancel = context.WithTimeout(ctx, p.cfg.QueueTimeout)
	} else {
		qCtx, cancel = context.WithCancel(ctx)
	}
	defer cancel()
	select {
	case p.sem <- struct{}{}:
		return nil
	case <-qCtx.Done():
		return qCtx.Err()
	}
}

func (p *Provider) Release() {
	if p.sem == nil {
		return
	}
	<-p.sem
}

// ---- request construction ----

func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
	meta := map[string]interface{}{
		"url":       url,
		"method":    http.MethodPost,
		"body":      body,
		"api_key":   p.cfg.APIKey,
		"timestamp": types.Now(),
		"source": map[string]interface{}{
			"name": p.cfg.Name,
			"meta": p.cfg.Meta,
		},
	}
	hdrs, err := p.vm.BuildHeaders(p.adapter, meta)
	if err != nil {
		return nil, err
	}
	h := http.Header{}
	h.Set("Content-Type", "application/json")
	for k, v := range p.cfg.Headers {
		h.Set(k, v)
	}
	for k, v := range hdrs {
		if _, ok := p.cfg.Headers[k]; !ok {
			h.Set(k, v)
		}
	}
	if h.Get("Authorization") == "" && p.cfg.APIKey != "" {
		h.Set("Authorization", "Bearer "+p.cfg.APIKey)
	}
	return h, nil
}

// ---- chat ----

// Chat performs a non-streaming round trip and returns the unified response.
func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
	if err := p.Acquire(ctx); err != nil {
		return nil, err
	}
	defer p.Release()

	body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
	if err != nil {
		return nil, err
	}
	hdrs, err := p.buildHeaders(body, p.URL())
	if err != nil {
		return nil, err
	}
	raw, status, err := p.do(ctx, p.URL(), body, hdrs)
	if err != nil {
		p.reportError()
		return nil, err
	}
	if status != 200 {
		p.ReportStatus(status)
		return nil, fmt.Errorf("api error %d: %s", status, truncate(raw, 500))
	}
	unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
	if err != nil {
		return nil, err
	}
	var out types.UnifiedResponse
	if err := json.Unmarshal([]byte(unified), &out); err != nil {
		return nil, fmt.Errorf("unmarshal unified response: %w (body: %s)", err, unified)
	}
	p.reportOK()
	return &out, nil
}

// ChatStream performs a streaming round trip, emitting unified chunks.
func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
	if err := p.Acquire(ctx); err != nil {
		return nil, err
	}
	req.Stream = true
	body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
	if err != nil {
		p.Release()
		return nil, err
	}
	hdrs, err := p.buildHeaders(body, p.URL())
	if err != nil {
		p.Release()
		return nil, err
	}

	type respOrErr struct {
		resp *http.Response
		err  error
	}
	rc := make(chan respOrErr, 1)
	go func() {
		resp, err := p.doRaw(ctx, p.URL(), body, hdrs)
		rc <- respOrErr{resp, err}
	}()

	ch := make(chan types.UnifiedChunk, 64)
	go func() {
		defer p.Release()
		defer close(ch)
		sel := <-rc
		if sel.err != nil {
			p.reportError()
			return
		}
		defer sel.resp.Body.Close()
		if sel.resp.StatusCode != 200 {
			raw, _ := io.ReadAll(sel.resp.Body)
			p.ReportStatus(sel.resp.StatusCode)
			_ = raw
			return
		}
		scanner := bufio.NewScanner(sel.resp.Body)
		scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" || !strings.HasPrefix(line, "data:") {
				continue
			}
			data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if data == "" {
				continue
			}
			if data == "[DONE]" {
				select {
				case ch <- types.UnifiedChunk{Done: true}:
				case <-ctx.Done():
				}
				continue
			}
			unified, err := p.vm.Transform(p.adapter, "transform_stream_chunk", data)
			if err != nil || unified == "" {
				continue
			}
			if unified == data {
				unified = standardSSEChunk(data)
				if unified == "" {
					continue
				}
			}
			var ck types.UnifiedChunk
			if err := json.Unmarshal([]byte(unified), &ck); err != nil {
				continue
			}
			select {
			case ch <- ck:
			case <-ctx.Done():
				return
			}
		}
	}()
	return ch, nil
}

// Image generates images via /v1/images/generations.
func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
	if err := p.Acquire(ctx); err != nil {
		return nil, err
	}
	defer p.Release()

	b, _ := json.Marshal(req)
	transformed, err := p.vm.Transform(p.adapter+"_image", "transform_request", string(b))
	if err != nil {
		// fall back to passthrough adapter (openai-style)
		transformed = string(b)
	}
	hdrs, err := p.buildHeaders(transformed, p.ImageURL())
	if err != nil {
		return nil, err
	}
	raw, status, err := p.do(ctx, p.ImageURL(), transformed, hdrs)
	if err != nil {
		p.reportError()
		return nil, err
	}
	if status != 200 {
		p.ReportStatus(status)
		return nil, fmt.Errorf("image api error %d: %s", status, truncate(raw, 500))
	}
	var out types.UnifiedResponse
	// try adapter transform_response; if missing, parse standard openai image format
	unified, terr := p.vm.Transform(p.adapter+"_image", "transform_response", raw)
	if terr == nil && unified != raw {
		if err := json.Unmarshal([]byte(unified), &out); err == nil {
			p.reportOK()
			return &out, nil
		}
	}
	var img types.ImageGenResponse
	if err := json.Unmarshal([]byte(raw), &img); err != nil {
		return nil, fmt.Errorf("unmarshal image response: %w", err)
	}
	out.ImageData = img.Data
	p.reportOK()
	return &out, nil
}

// ---- http helpers ----

func (p *Provider) do(ctx context.Context, url, body string, hdr http.Header) (string, int, error) {
	resp, err := p.doRaw(ctx, url, body, hdr)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	return string(raw), resp.StatusCode, nil
}

func (p *Provider) doRaw(ctx context.Context, url, body string, hdr http.Header) (*http.Response, error) {
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body)))
	if err != nil {
		return nil, err
	}
	httpReq.Header = hdr
	return p.client.Do(httpReq)
}

func marshalTransform(vm *lua.VM, adapter, fn string, v interface{}) (string, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	out, err := vm.Transform(adapter, fn, string(b))
	if err != nil {
		return "", err
	}
	return out, nil
}

func standardSSEChunk(data string) string {
	var raw struct {
		Choices []struct {
			Delta struct {
				Content string `json:"content"`
			} `json:"delta"`
			FinishReason *string `json:"finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal([]byte(data), &raw); err != nil {
		return ""
	}
	if len(raw.Choices) == 0 {
		return ""
	}
	out, _ := json.Marshal(types.UnifiedChunk{
		Content: raw.Choices[0].Delta.Content,
		Done:    raw.Choices[0].FinishReason != nil,
	})
	return string(out)
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}