use crate::math::{clamp, softmax};
use crate::types::gate::{GateResult, GateType};
use crate::types::world::GainMode;
pub const BUDGET_SIMPLE: f64 = 3000.0;
pub const BUDGET_MODERATE: f64 = 5000.0;
pub const BUDGET_COMPLEX: f64 = 8000.0;
pub const BUDGET_CEILING: f64 = 12000.0;
const YERKES_SIMPLE_BOOST: f64 = 0.15;
const YERKES_COMPLEX_PENALTY: f64 = 0.20;
const YERKES_AROUSAL_THRESHOLD: f64 = 0.6;
pub const PINNED_CHARS_PER_ATOM: f64 = 400.0;
pub const PROSPECTIVE_CHARS: f64 = 200.0;
const TEMP_PHASIC: f64 = 0.15;
const TEMP_NEUTRAL: f64 = 0.30;
const TEMP_TONIC: f64 = 0.50;
const UNCERTAINTY_THRESHOLD: f64 = 0.45;
const EXPLORATION_FLOOR: f64 = 0.3;
const EPISTEMIC_BONUS: f64 = 0.15;
pub const TOPIC_PRECISION_DEFAULT: f64 = 0.5;
const THREAT_MEMORY_BOOST: f64 = 0.25;
pub const LAYER_MEMORY: &str = "layer:memory";
pub const LAYER_GRAPH: &str = "layer:graph";
pub const LAYER_ARTIFACT: &str = "layer:artifact";
pub const LAYER_RECAP: &str = "layer:recap";
const MODEL_TIER_NANO: f64 = 0.6;
const MODEL_TIER_MEDIUM: f64 = 1.0;
const MODEL_TIER_LARGE: f64 = 1.5;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskLoad {
Simple,
Moderate,
Complex,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ModelTier {
Nano,
#[default]
Medium,
Large,
}
#[derive(Debug, Clone)]
pub struct AllocatorContext<'a> {
pub query: &'a str,
pub gate_result: Option<&'a GateResult>,
pub gain_mode: GainMode,
pub arousal: f64,
pub fok_average: Option<f64>,
pub model_tier: ModelTier,
pub has_graph_data: bool,
pub active_file_count: usize,
pub pinned_count: usize,
pub has_prospective: bool,
pub message_count: usize,
pub has_threat_topics: bool,
}
#[derive(Debug, Clone)]
pub struct LayerBudget {
pub layer_id: String,
pub budget: f64,
pub query_relevance: f64,
pub precision: f64,
pub effective_precision: f64,
pub expected_utility: f64,
}
#[derive(Debug, Clone)]
pub struct AllocationResult {
pub total_budget: f64,
pub guaranteed_used: f64,
pub competitive_pool: f64,
pub task_load: TaskLoad,
pub layers: Vec<LayerBudget>,
}
pub fn estimate_task_load(query: &str, gate_result: Option<&GateResult>) -> TaskLoad {
let question_count = query.matches('?').count();
if question_count >= 2 || query.len() > 200 {
return TaskLoad::Complex;
}
let compare_patterns = ["compare", "difference", "versus", "vs", "pros and cons"];
if compare_patterns
.iter()
.any(|p| query.to_lowercase().contains(p))
{
return TaskLoad::Complex;
}
if query.len() < 50 && question_count == 0 {
return TaskLoad::Simple;
}
if let Some(gr) = gate_result {
if gr.gate == GateType::Routine {
return TaskLoad::Simple;
}
}
TaskLoad::Moderate
}
pub fn compute_total_budget(task_load: TaskLoad, arousal: f64, model_tier: ModelTier) -> f64 {
let mut budget = match task_load {
TaskLoad::Simple => BUDGET_SIMPLE,
TaskLoad::Moderate => BUDGET_MODERATE,
TaskLoad::Complex => BUDGET_COMPLEX,
};
if arousal >= YERKES_AROUSAL_THRESHOLD {
match task_load {
TaskLoad::Simple => budget *= 1.0 + YERKES_SIMPLE_BOOST,
TaskLoad::Complex => budget *= 1.0 - YERKES_COMPLEX_PENALTY,
TaskLoad::Moderate => {} }
}
let tier_scale = match model_tier {
ModelTier::Nano => MODEL_TIER_NANO,
ModelTier::Medium => MODEL_TIER_MEDIUM,
ModelTier::Large => MODEL_TIER_LARGE,
};
budget *= tier_scale;
budget.min(BUDGET_CEILING)
}
fn get_temperature(gain_mode: GainMode) -> f64 {
match gain_mode {
GainMode::Phasic => TEMP_PHASIC,
GainMode::Tonic => TEMP_TONIC,
GainMode::Neutral => TEMP_NEUTRAL,
}
}
pub fn compute_effective_precision(
raw_precision: f64,
fok_average: Option<f64>,
layer_has_content: bool,
) -> f64 {
if !layer_has_content {
return 0.0;
}
match fok_average {
None => raw_precision,
Some(fok) if fok >= UNCERTAINTY_THRESHOLD => raw_precision, Some(_) => {
let bonus = if raw_precision < TOPIC_PRECISION_DEFAULT {
EPISTEMIC_BONUS
} else {
0.0
};
(raw_precision + bonus).max(EXPLORATION_FLOOR)
}
}
}
fn estimate_query_relevance(layer_id: &str, ctx: &AllocatorContext) -> f64 {
match layer_id {
LAYER_MEMORY => {
let mut relevance = 0.3; if ctx.has_threat_topics {
relevance += THREAT_MEMORY_BOOST;
}
let q = ctx.query.to_lowercase();
if q.contains("before") || q.contains("earlier") || q.contains("previous") {
relevance += 0.3;
}
relevance.min(1.0)
}
LAYER_GRAPH => {
if !ctx.has_graph_data {
0.0
} else {
0.4
}
}
LAYER_ARTIFACT => {
if ctx.active_file_count == 0 {
0.0
} else {
0.4
}
}
LAYER_RECAP => {
if ctx.message_count >= 20 {
0.4
} else {
0.0
}
}
_ => 0.3,
}
}
pub fn allocate_context_budget(
ctx: &AllocatorContext,
layer_precisions: &[(String, f64)],
) -> Option<AllocationResult> {
let task_load = estimate_task_load(ctx.query, ctx.gate_result);
let total_budget = compute_total_budget(task_load, ctx.arousal, ctx.model_tier);
let pinned_chars = ctx.pinned_count as f64 * PINNED_CHARS_PER_ATOM;
let prospective_chars = if ctx.has_prospective {
PROSPECTIVE_CHARS
} else {
0.0
};
let guaranteed_used = pinned_chars + prospective_chars;
let competitive_pool = (total_budget - guaranteed_used).max(0.0);
let layer_ids = [LAYER_MEMORY, LAYER_GRAPH, LAYER_ARTIFACT, LAYER_RECAP];
let mut layers = Vec::new();
for layer_id in &layer_ids {
let query_relevance = estimate_query_relevance(layer_id, ctx);
let raw_precision = layer_precisions
.iter()
.find(|(id, _)| id == *layer_id)
.map(|(_, p)| *p)
.unwrap_or(TOPIC_PRECISION_DEFAULT);
let has_content = match *layer_id {
LAYER_GRAPH => ctx.has_graph_data,
LAYER_ARTIFACT => ctx.active_file_count > 0,
LAYER_RECAP => ctx.message_count >= 20,
_ => true, };
let effective_precision =
compute_effective_precision(raw_precision, ctx.fok_average, has_content);
let expected_utility = query_relevance * effective_precision;
layers.push(LayerBudget {
layer_id: layer_id.to_string(),
budget: 0.0, query_relevance,
precision: raw_precision,
effective_precision,
expected_utility,
});
}
let utilities: Vec<f64> = layers.iter().map(|l| l.expected_utility).collect();
let temp = get_temperature(ctx.gain_mode);
let weights = softmax(&utilities, temp);
for (layer, weight) in layers.iter_mut().zip(weights.iter()) {
layer.budget = competitive_pool * weight;
}
Some(AllocationResult {
total_budget,
guaranteed_used,
competitive_pool,
task_load,
layers,
})
}
pub fn compute_resource_pressure(allocation: Option<&AllocationResult>) -> f64 {
match allocation {
None => 0.5, Some(alloc) => {
if alloc.competitive_pool <= 0.0 {
return 1.0;
}
let summed_utility: f64 = alloc
.layers
.iter()
.map(|l| l.query_relevance * l.effective_precision)
.sum();
clamp(summed_utility, 0.0, 1.0)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn task_load_simple() {
assert_eq!(estimate_task_load("ok", None), TaskLoad::Simple);
}
#[test]
fn task_load_complex_multiple_questions() {
assert_eq!(
estimate_task_load("What is X? And how does Y work?", None),
TaskLoad::Complex
);
}
#[test]
fn task_load_complex_compare() {
assert_eq!(
estimate_task_load("Compare React vs Vue", None),
TaskLoad::Complex
);
}
#[test]
fn total_budget_simple() {
let budget = compute_total_budget(TaskLoad::Simple, 0.0, ModelTier::Medium);
assert!((budget - BUDGET_SIMPLE).abs() < f64::EPSILON);
}
#[test]
fn total_budget_yerkes_dodson_simple_boost() {
let budget = compute_total_budget(TaskLoad::Simple, 0.8, ModelTier::Medium);
assert!(budget > BUDGET_SIMPLE); }
#[test]
fn total_budget_yerkes_dodson_complex_penalty() {
let budget = compute_total_budget(TaskLoad::Complex, 0.8, ModelTier::Medium);
assert!(budget < BUDGET_COMPLEX); }
#[test]
fn total_budget_model_tier_scaling() {
let nano = compute_total_budget(TaskLoad::Moderate, 0.0, ModelTier::Nano);
let large = compute_total_budget(TaskLoad::Moderate, 0.0, ModelTier::Large);
assert!(nano < large);
}
#[test]
fn total_budget_ceiling() {
let budget = compute_total_budget(TaskLoad::Complex, 0.0, ModelTier::Large);
assert!(budget <= BUDGET_CEILING);
}
#[test]
fn effective_precision_exploit_mode() {
let eff = compute_effective_precision(0.8, Some(0.7), true);
assert!((eff - 0.8).abs() < f64::EPSILON);
}
#[test]
fn effective_precision_explore_mode() {
let eff = compute_effective_precision(0.2, Some(0.3), true);
assert!(eff >= EXPLORATION_FLOOR);
assert!(eff > 0.2); }
#[test]
fn effective_precision_no_content() {
let eff = compute_effective_precision(0.8, Some(0.7), false);
assert_eq!(eff, 0.0);
}
#[test]
fn allocate_budget_returns_result() {
let ctx = AllocatorContext {
query: "Tell me about Rust",
gate_result: None,
gain_mode: GainMode::Neutral,
arousal: 0.3,
fok_average: Some(0.5),
model_tier: ModelTier::Medium,
has_graph_data: false,
active_file_count: 0,
pinned_count: 0,
has_prospective: false,
message_count: 5,
has_threat_topics: false,
};
let result = allocate_context_budget(&ctx, &[]);
assert!(result.is_some());
let alloc = result.unwrap();
assert!(alloc.total_budget > 0.0);
assert_eq!(alloc.layers.len(), 4);
}
#[test]
fn allocate_budget_softmax_sums_to_pool() {
let ctx = AllocatorContext {
query: "Complex question about multiple topics here?",
gate_result: None,
gain_mode: GainMode::Neutral,
arousal: 0.0,
fok_average: None,
model_tier: ModelTier::Medium,
has_graph_data: true,
active_file_count: 3,
pinned_count: 0,
has_prospective: false,
message_count: 25,
has_threat_topics: false,
};
let result = allocate_context_budget(&ctx, &[]).unwrap();
let sum: f64 = result.layers.iter().map(|l| l.budget).sum();
assert!((sum - result.competitive_pool).abs() < 1.0); }
#[test]
fn phasic_concentrates_budget() {
let make_ctx = |gain| AllocatorContext {
query: "Tell me about Rust",
gate_result: None,
gain_mode: gain,
arousal: 0.0,
fok_average: None,
model_tier: ModelTier::Medium,
has_graph_data: true,
active_file_count: 3,
pinned_count: 0,
has_prospective: false,
message_count: 25,
has_threat_topics: false,
};
let phasic = allocate_context_budget(&make_ctx(GainMode::Phasic), &[]).unwrap();
let tonic = allocate_context_budget(&make_ctx(GainMode::Tonic), &[]).unwrap();
let phasic_max = phasic
.layers
.iter()
.map(|l| l.budget)
.fold(0.0_f64, f64::max);
let tonic_max = tonic
.layers
.iter()
.map(|l| l.budget)
.fold(0.0_f64, f64::max);
let phasic_ratio = phasic_max / phasic.competitive_pool;
let tonic_ratio = tonic_max / tonic.competitive_pool;
assert!(phasic_ratio > tonic_ratio);
}
#[test]
fn resource_pressure_within_bounds() {
let ctx = AllocatorContext {
query: "test query",
gate_result: None,
gain_mode: GainMode::Neutral,
arousal: 0.5,
fok_average: None,
model_tier: ModelTier::Medium,
has_graph_data: true,
active_file_count: 3,
pinned_count: 2,
has_prospective: true,
message_count: 30,
has_threat_topics: true,
};
let alloc = allocate_context_budget(&ctx, &[]).unwrap();
let pressure = compute_resource_pressure(Some(&alloc));
assert!((0.0..=1.0).contains(&pressure));
}
}