latencypredictorclient

type BucketCounts struct {
	TTFTBuckets map[int]int `json:"ttft_buckets"`
	TPOTBuckets map[int]int `json:"tpot_buckets"`
}

type BulkPredictionError struct {
	Index   int               `json:"index"`
	Error   string            `json:"error"`
	Request PredictionRequest `json:"request"`
}

type BulkPredictionRequest struct {
	Requests []PredictionRequest `json:"requests"`
}

type BulkPredictionResponse struct {
	Predictions           []PredictionResponse `json:"predictions"`
	TotalRequests         int                  `json:"total_requests"`
	SuccessfulPredictions int                  `json:"successful_predictions"`
	FailedPredictions     int                  `json:"failed_predictions"`
	ProcessingTimeMs      float64              `json:"processing_time_ms"`
}

type BulkPredictionResponseWithErrors struct {
	Predictions           []*PredictionResponse `json:"predictions"`
	Errors                []BulkPredictionError `json:"errors"`
	TotalRequests         int                   `json:"total_requests"`
	SuccessfulPredictions int                   `json:"successful_predictions"`
	FailedPredictions     int                   `json:"failed_predictions"`
	ProcessingTimeMs      float64               `json:"processing_time_ms"`
}

type BulkTrainingRequest struct {
	Entries []TrainingEntry `json:"entries"`
}

type Config struct {
	// TrainingURL is the base URL of the Python training server.
	TrainingURL string
	// PredictionURLs is a list of prediction server URLs for load balancing.
	PredictionURLs []string
	// MaxSampleSize is the maximum number of training entries to send in each flush.
	// If the buffer contains more entries, they will be randomly sampled.
	MaxSampleSize int
	// FlushInterval determines how often to flush training & refresh metrics.
	FlushInterval time.Duration
	// UseNativeXGBoost when true, attempts to use local XGBoost models for prediction.
	// When false, falls back to HTTP calls to the Python server for XGBoost predictions.
	UseNativeXGBoost bool
	// HTTPTimeout is the timeout for HTTP requests to the Python server.
	HTTPTimeout time.Duration
	// MetricsRefreshInterval determines how often to refresh cached metrics.
	MetricsRefreshInterval time.Duration
	// MaxBulkSize is the maximum number of predictions to send in a single bulk request.
	MaxBulkSize int
	// CoalesceWindow is how long the coalescer waits to accumulate concurrent
	// PredictBulkStrict callers before firing one mega-batch HTTP call.
	// Set to 0 to disable coalescing (each caller gets its own HTTP call).
	CoalesceWindow time.Duration
	// MaxCoalescedRows caps the total number of rows in one coalesced mega-batch,
	// causing an early dispatch before the window expires.
	// This is separate from MaxBulkSize (the per-caller row limit).
	// Default 0 means no row cap (window-only dispatch).
	MaxCoalescedRows int
	// MaxConcurrentDispatches limits how many coalesced HTTP calls can be
	// in-flight to the prediction server simultaneously. Defaults to 36.
	MaxConcurrentDispatches int
}

type MetricsResponse struct {
	ModelType    string             `json:"model_type"`
	Coefficients *ModelCoefficients `json:"coefficients"`
	XGBoostTrees *XGBoostTrees      `json:"xgboost_trees"`
	BucketCounts *BucketCounts      `json:"bucket_counts"`
	RawMetrics   string             `json:"raw_metrics"`
}

type ModelCoefficients struct {
	TTFTIntercept float64            `json:"ttft_intercept"`
	TTFTCoeffs    map[string]float64 `json:"ttft_coefficients"`
	TPOTIntercept float64            `json:"tpot_intercept"`
	TPOTCoeffs    map[string]float64 `json:"tpot_coefficients"`
}

type ModelInfo struct {
	ModelType     string          `json:"model_type"`
	ObjectiveType string          `json:"objective_type,omitempty"`
	ModelStatus   map[string]bool `json:"model_status"`
	Quantile      float64         `json:"quantile"`
}

type PredictionRequest struct {
	KVCachePercentage     float64 `json:"kv_cache_percentage"`
	InputTokenLength      int     `json:"input_token_length"`
	NumRequestWaiting     int     `json:"num_request_waiting"`
	NumRequestRunning     int     `json:"num_request_running"`
	NumTokensGenerated    int     `json:"num_tokens_generated"`
	PrefixCacheScore      float64 `json:"prefix_cache_score"`
	PodType               string  `json:"pod_type,omitempty"` // "prefill", "decode", or "" for monolithic
	PrefillTokensInFlight int64   `json:"prefill_tokens_in_flight"`
	DecodeTokensInFlight  int64   `json:"decode_tokens_in_flight"`
}

type PredictionResponse struct {
	TTFT                 float64    `json:"ttft_ms"`
	TPOT                 float64    `json:"tpot_ms"`
	TTFTUncertainty      float64    `json:"ttft_uncertainty,omitempty"`
	TPOTUncertainty      float64    `json:"tpot_uncertainty,omitempty"`
	TTFTPredictionBounds [2]float64 `json:"ttft_prediction_bounds,omitempty"`
	TPOTPredictionBounds [2]float64 `json:"tpot_prediction_bounds,omitempty"`
	PredictedAt          time.Time  `json:"predicted_at"`
	ModelType            string     `json:"model_type"`
	ObjectiveType        string     `json:"objective_type,omitempty"`
	Quantile             float64    `json:"quantile"`
	LastModelLoad        *time.Time `json:"last_model_load"`
}

type Predictor struct {
	// contains filtered or unexported fields
}

type PredictorInterface interface {
	Predict(ctx context.Context, req PredictionRequest) (*PredictionResponse, error)
	PredictBulk(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error)
	PredictBulkStrict(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error)
	AddTrainingDataBulk(entry []TrainingEntry) error
}

type ServerStatusResponse struct {
	IsReady           bool            `json:"is_ready"`
	ModelType         string          `json:"model_type"`
	ObjectiveType     string          `json:"objective_type,omitempty"`
	Quantile          float64         `json:"quantile"`
	LastModelLoad     *time.Time      `json:"last_model_load"`
	TrainingServerURL string          `json:"training_server_url"`
	ModelsExist       map[string]bool `json:"models_exist"`
}

type TrainingEntry struct {
	KVCachePercentage     float64   `json:"kv_cache_percentage"`
	InputTokenLength      int       `json:"input_token_length"`
	NumRequestWaiting     int       `json:"num_request_waiting"`
	NumRequestRunning     int       `json:"num_request_running"`
	NumTokensGenerated    int       `json:"num_tokens_generated"`
	ActualTTFT            float64   `json:"actual_ttft_ms"`
	ActualTPOT            float64   `json:"actual_tpot_ms"`
	PrefixCacheScore      float64   `json:"prefix_cache_score"`
	PodType               string    `json:"pod_type,omitempty"` // "prefill", "decode", or "" for monolithic
	PrefillTokensInFlight int64     `json:"prefill_tokens_in_flight"`
	DecodeTokensInFlight  int64     `json:"decode_tokens_in_flight"`
	Timestamp             time.Time `json:"timestamp"`
}

type XGBoostTrees struct {
	TTFTTrees []any `json:"ttft_trees"`
	TPOTTrees []any `json:"tpot_trees"`
}