homelab-frontend/internal/proxy/router.go

// Package proxy provides request routing and forwarding.
package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"

	"github.com/Riotpiaole/homelab-frontend/internal/config"
)

// RouteRequest determines which upstream should handle the request.
// For /v1/chat/completions, it uses body-based dispatch (reads JSON to find "model" field).
// For other routes, it looks up by path prefix.
func (h *Handler) RouteRequest(r *http.Request) (*Route, error) {
	// For /v1/chat/completions, use body-based dispatch
	if r.URL.Path == "/v1/chat/completions" && r.Method == "POST" {
		return h.routeByModel(r)
	}

	// For other paths, try to find a matching route by path
	// Look for exact path match or path prefix match
	for _, route := range h.routes {
		// Check if route's pathRewrite matches request path
		if route.Upstream.PathRewrite != "" && route.Upstream.PathRewrite == r.URL.Path {
			return route, nil
		}
	}

	// If no routes configured, return error
	if len(h.routes) == 0 {
		return nil, fmt.Errorf("no routes configured")
	}

	// Return the first configured route as fallback
	for _, route := range h.routes {
		return route, nil
	}

	return nil, fmt.Errorf("no route available")
}

// routeByModel reads the request body to find the "model" field and routes accordingly.
// The body is preserved for forwarding to the upstream.
func (h *Handler) routeByModel(r *http.Request) (*Route, error) {
	// If there's no body, we can't determine the model
	if r.Body == nil {
		return nil, fmt.Errorf("request body required")
	}

	// Read the body to extract the model name
	// We need to be careful to preserve the body for the upstream
	bodyBytes, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read request body: %w", err)
	}

	// Restore the body so it can be read again by the upstream
	r.Body = io.NopCloser(bytes.NewReader(bodyBytes))

	// Parse the JSON to find the model field
	var payload map[string]interface{}
	if err := json.Unmarshal(bodyBytes, &payload); err != nil {
		return nil, fmt.Errorf("invalid JSON in request body: %w", err)
	}

	// Extract the model name
	modelName, ok := payload["model"].(string)
	if !ok {
		return nil, fmt.Errorf("model field missing or not a string")
	}

	// Look up the model in the registry
	modelUpstream := h.config.LookupModel(modelName)
	if modelUpstream == nil {
		return nil, fmt.Errorf("unknown model: %q", modelName)
	}

	// Create a route for this model with appropriate timeouts
	// These are sensible defaults for LLM models
	upstreamCfg := &config.Upstream{
		Address:        modelUpstream.Address,
		PathRewrite:    "/v1/chat/completions",
		ConnectTimeout: h.defaultConnectTimeout,
		ReadTimeout:    h.defaultReadTimeout,
		WriteTimeout:   h.defaultWriteTimeout,
		MaxBodySize:    h.defaultMaxBodySize,
		AuthRequired:   false,
	}

	targetURL, _ := url.Parse("http://" + modelUpstream.Address)

	route := &Route{
		Name:      "v1-chat-" + modelName,
		Upstream:  upstreamCfg,
		Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg),
		Director: func(req *http.Request) {
			directorFunc(req, targetURL, upstreamCfg)
		},
	}

	return route, nil
}