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, /v1/embeddings, and /v1/rerank, it uses body-based dispatch.
// For other routes, it looks up by path in the configured routes.
func (h *Handler) RouteRequest(r *http.Request) (*Route, error) {
	// For /v1/chat/completions, /v1/embeddings, /v1/rerank use body-based dispatch
	if r.Method == "POST" && (r.URL.Path == "/v1/chat/completions" || r.URL.Path == "/v1/embeddings" || r.URL.Path == "/v1/rerank") {
		return h.routeByModel(r, r.URL.Path)
	}

	// 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.
// Returns a modelValidationError for client errors (invalid JSON, missing/unknown model).
// The path parameter indicates which endpoint is being called (/v1/chat/completions, /v1/embeddings, /v1/rerank)
func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
	// If there's no body, we can't determine the model
	if r.Body == nil {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "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, &modelValidationError{
			Kind:    "invalid_request",
			Message: fmt.Sprintf("failed to read request body: %v", 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, &modelValidationError{
			Kind:    "invalid_json",
			Message: "request body is not valid JSON",
		}
	}

	// Extract the model name
	modelVal, hasModel := payload["model"]
	if !hasModel {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "'model' field is missing",
		}
	}

	// Handle null model
	if modelVal == nil {
		return nil, &modelValidationError{
			Kind:    "null_model",
			Message: "'model' field is null",
		}
	}

	// Extract as string
	modelName, ok := modelVal.(string)
	if !ok {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "'model' field must be a string",
		}
	}

	// Handle empty string
	if modelName == "" {
		return nil, &modelValidationError{
			Kind:    "empty_model",
			Message: "'model' field cannot be empty",
		}
	}

	// Look up the model in the registry
	modelUpstream := h.config.LookupModel(modelName)
	if modelUpstream == nil {
		return nil, &modelValidationError{
			Kind:    "unknown_model",
			Message: fmt.Sprintf("unknown model: %q", modelName),
			Model:   modelName,
		}
	}

	// Determine the upstream path based on the request path
	upstreamPath := path
	if path == "/v1/rerank" {
		// Rerank endpoint uses /rerank path on upstream
		upstreamPath = "/rerank"
	}

	// Create a route for this model with appropriate timeouts
	// These are sensible defaults for LLM models
	upstreamCfg := &config.Upstream{
		Address:        modelUpstream.Address,
		PathRewrite:    upstreamPath,
		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
}