Luminary/internal/provider/forward.go

package provider

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

	"github.com/rose_cat707/luminary/internal/model"
)

type ForwardRequest struct {
	Endpoint EndpointType
	Method   string
	Body     []byte
	Headers  map[string]string
}

type ForwardResponse struct {
	StatusCode int
	Body       []byte
	TokensIn   int
	TokensOut  int
}

type ProviderConfig struct {
	Type           model.ProviderType
	Category       model.ProviderCategory
	BaseURL        string
	EndpointPaths  map[string]string // endpoint key -> path override
}

func ParseEndpointPaths(jsonStr string) map[string]string {
	out := map[string]string{}
	if jsonStr == "" || jsonStr == "{}" {
		return out
	}
	_ = json.Unmarshal([]byte(jsonStr), &out)
	return out
}

func ParseAllowedEndpoints(jsonStr string, category model.ProviderCategory) []string {
	var out []string
	if jsonStr == "" {
		return DefaultEndpointsForCategory(category)
	}
	_ = json.Unmarshal([]byte(jsonStr), &out)
	if len(out) == 0 {
		return DefaultEndpointsForCategory(category)
	}
	return out
}

func (cfg *ProviderConfig) ResolveEndpointURL(endpoint EndpointType) string {
	override := cfg.EndpointPaths[string(endpoint)]
	defaultPath := DefaultPath(cfg.Type, endpoint)
	return ResolveURL(cfg.BaseURL, override, defaultPath)
}

func ForwardHTTP(ctx context.Context, client *http.Client, method, url string, apiKey string, providerType model.ProviderType, body []byte, extraHeaders map[string]string) (*ForwardResponse, error) {
	var reader io.Reader
	if len(body) > 0 {
		reader = bytes.NewReader(body)
	}
	req, err := http.NewRequestWithContext(ctx, method, url, reader)
	if err != nil {
		return nil, err
	}
	setAuthHeaders(req, apiKey, providerType)
	req.Header.Set("Content-Type", "application/json")
	for k, v := range extraHeaders {
		req.Header.Set(k, v)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	tokensIn, tokensOut := parseUsageJSON(respBody)
	return &ForwardResponse{
		StatusCode: resp.StatusCode,
		Body:       respBody,
		TokensIn:   tokensIn,
		TokensOut:  tokensOut,
	}, nil
}

func setAuthHeaders(req *http.Request, apiKey string, providerType model.ProviderType) {
	switch providerType {
	case model.ProviderAnthropic:
		req.Header.Set("x-api-key", apiKey)
		req.Header.Set("anthropic-version", "2023-06-01")
	case model.ProviderAzureOpenAI:
		if apiKey != "" {
			req.Header.Set("api-key", apiKey)
		}
	case model.ProviderOllama:
		if apiKey != "" {
			req.Header.Set("Authorization", "Bearer "+apiKey)
		}
	default:
		if apiKey != "" {
			req.Header.Set("Authorization", "Bearer "+apiKey)
		}
	}
}

func parseUsageJSON(body []byte) (int, int) {
	var u struct {
		Usage struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
			InputTokens      int `json:"input_tokens"`
			OutputTokens     int `json:"output_tokens"`
		} `json:"usage"`
	}
	_ = json.Unmarshal(body, &u)
	in := u.Usage.PromptTokens
	out := u.Usage.CompletionTokens
	if in == 0 {
		in = u.Usage.InputTokens
	}
	if out == 0 {
		out = u.Usage.OutputTokens
	}
	return in, out
}

func DefaultBaseURL(providerType model.ProviderType) string {
	switch providerType {
	case model.ProviderAnthropic:
		return "https://api.anthropic.com/v1"
	case model.ProviderAzureOpenAI:
		return "" // must be configured per deployment
	case model.ProviderOpenRouter:
		return "https://openrouter.ai/api/v1"
	case model.ProviderOllama:
		return "http://127.0.0.1:11434/v1"
	default:
		return "https://api.openai.com/v1"
	}
}

func IsOpenAICompatible(t model.ProviderType) bool {
	switch t {
	case model.ProviderOpenAI, model.ProviderCustom, model.ProviderAzureOpenAI, model.ProviderOpenRouter, model.ProviderOllama:
		return true
	default:
		return false
	}
}

func NormalizeBaseURL(baseURL string, providerType model.ProviderType) string {
	if baseURL == "" {
		return DefaultBaseURL(providerType)
	}
	return baseURL
}

func IsEndpointAllowed(allowed []string, endpoint EndpointType) bool {
	for _, a := range allowed {
		if a == string(endpoint) || a == "*" {
			return true
		}
	}
	return false
}

func NewHTTPClient() *http.Client {
	return &http.Client{Timeout: 120 * time.Second}
}

func ValidateCategoryEndpoints(category model.ProviderCategory, allowed []string) error {
	if !model.IsCategorySupported(category) {
		return fmt.Errorf("category %s is not supported yet", category)
	}
	valid := EndpointsForCategory(category)
	validSet := map[string]bool{}
	for _, e := range valid {
		validSet[string(e.Key)] = true
	}
	for _, a := range allowed {
		if a == "*" {
			continue
		}
		if !validSet[a] {
			return fmt.Errorf("endpoint %s is not available for category %s", a, category)
		}
	}
	return nil
}