types

type FallbackRouter interface {
	Router

	// SetFallback sets the fallback router
	SetFallback(RoutingAlgorithm, RouterProviderFunc)
}

type OutputPredictor interface {
	// AddTrace collects history input and output tokens data.
	AddTrace(inputTokens, outputTokens int, cnt int32)

	// Predict outputs the number of output tokens based on the number of input tokens.
	Predict(promptLen int) (outputLen int)
}

type OutputPredictorProvider interface {
	// GetOutputPredictor returns the output predictor
	GetOutputPredictor(modelName string) (OutputPredictor, error)
}

type PodList interface {
	// Len returns the number of pods in the list.
	Len() int

	// All returns a slice of all pods in the list.
	All() []*v1.Pod

	// Indexes returns a slice of indexes for querying pods by index.
	Indexes() []string

	// ListByIndex returns a slice of pods that match the given index.
	ListByIndex(index string) []*v1.Pod

	// ListPortsForPod returns a map of portList that bind with pod, key podname
	ListPortsForPod() map[string][]int
}

type QueueRouter interface {
	Router

	Len() int
}

type RequestFeatures []float64

type ResolvedConfigProfile struct {
	RoutingStrategy          string
	PromptLenBucketMinLength int
	PromptLenBucketMaxLength int
	Combined                 bool
}

type Router interface {
	// Route selects a target pod from the provided list of pods.
	// The input pods is guaranteed to be non-empty and contain only routable pods.
	Route(ctx *RoutingContext, readyPodList PodList) (string, error)
}

type RouterConstructor func() (Router, error)

type RouterProvider interface {
	// GetRouter returns the router
	GetRouter(ctx *RoutingContext) (Router, error)
}

type RouterProviderFunc func(*RoutingContext) (Router, error)

type RouterProviderRegistrationFunc func() RouterProviderFunc

type RouterQueue[V comparable] interface {
	Enqueue(V, time.Time) error
	Peek(time.Time, PodList) (V, error)
	Dequeue(time.Time) (V, error)
	Len() int
}

type RoutingAlgorithm string

type RoutingContext struct {
	context.Context
	Algorithm      RoutingAlgorithm
	Model          string
	Stream         bool
	Message        string
	RequestID      string
	User           *string
	RequestTime    time.Time // Time when the routing context is created.
	RequestEndTime time.Time // Time when the routing is done and sent to inference engine.
	PendingLoad    float64   // Normalized pending load of request, available after AddRequestCount call. See cache.PendingLoadProvider
	TraceTerm      int64     // Trace term identifier, available after AddRequestCount call.
	RoutedTime     time.Time // Time consumed during routing.

	ReqHeaders       map[string]string
	ReqBody          []byte
	ReqPath          string
	ReqConfigProfile string

	PrefillStartTime time.Time // Time when prefill request is started.
	PrefillEndTime   time.Time // Time consumed during prefill.

	// RespHeaders holds response headers that the router intends to set.
	// These are typically used to propagate control information back to the client,
	// such as session affinity id.
	// The router implementation (e.g., sessionAffinityRouter) may populate this field
	// during the Route() call.
	RespHeaders map[string]string

	// ConfigProfile holds the resolved model config profile for this request.
	// Set in HandleRequestBody from model.aibrix.ai/config (annotation)
	// based on config-profile header. Nil when no config is present.
	ConfigProfile *ResolvedConfigProfile
	// contains filtered or unexported fields
}