//! LLM brain: async worker pool talking to Ollama, wake scheduling, action
//! application. The render loop never blocks on a model call.
use bevy::prelude::*;
use serde::Deserialize;
use serde_json::json;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use crate::agent::{cell_of, world_of, Brain, Memory, Names, Needs, Pathing, Speech, Villager};
use crate::persona::MODEL;
use crate::world::WorldMap;
const WORKERS: usize = 6; // in-flight cap
const WAKE_MIN: f32 = 45.0; // seconds between a villager's scheduled decisions
const WAKE_JITTER: f32 = 75.0;
const THINK_BUDGET: u32 = 64; // max new tokens per decision
pub struct DecisionJob {
pub entity: Entity,
pub name: String,
pub system: String,
pub user: String,
}
pub struct DecisionResult {
pub entity: Entity,
pub name: String,
pub action: Result<Action, String>,
pub elapsed: f32,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "action", rename_all = "snake_case")]
pub enum Action {
GoTo { place: String, reason: String },
Talk { to: String, say: String, reason: String },
Work { reason: String },
Rest { reason: String },
Idle { reason: String },
}
/// `Receiver` is `Send` but not `Sync`, so it rides behind a mutex to be a
/// valid Bevy `Resource`. Workers share the job receiver the same way.
#[derive(Resource)]
pub struct LlmPool {
pub jobs_tx: Sender<DecisionJob>,
pub results_rx: Mutex<Receiver<DecisionResult>>,
pub in_flight: usize,
pub completed: u64,
pub failed: u64,
pub peak_in_flight: usize,
}
pub fn start_pool(mut commands: Commands) {
let (jobs_tx, jobs_rx) = channel::<DecisionJob>();
let (results_tx, results_rx) = channel::<DecisionResult>();
let url = std::env::var("OLLAMA_URL").unwrap_or_else(|_| "http://localhost:11434".into());
let jobs_rx = Arc::new(Mutex::new(jobs_rx));
for _ in 0..WORKERS {
let jobs_rx = Arc::clone(&jobs_rx);
let results_tx = results_tx.clone();
let url = url.clone();
std::thread::spawn(move || worker_loop(jobs_rx, results_tx, &url));
}
commands.insert_resource(LlmPool {
jobs_tx,
results_rx: Mutex::new(results_rx),
in_flight: 0,
completed: 0,
failed: 0,
peak_in_flight: 0,
});
}
fn worker_loop(
jobs_rx: Arc<Mutex<Receiver<DecisionJob>>>,
results_tx: Sender<DecisionResult>,
url: &str,
) {
let client = reqwest::blocking::Client::builder()
.timeout(Duration::from_secs(120))
.build()
.expect("http client");
loop {
// Blocking recv while holding the lock is fine: workers take turns
// waiting, and the handoff is instant when the queue is non-empty.
let job = {
let rx = jobs_rx.lock().unwrap_or_else(|e| e.into_inner());
rx.recv()
};
let Ok(job) = job else { break };
let started = Instant::now();
let action = request_action(&client, url, &job.system, &job.user);
let _ = results_tx.send(DecisionResult {
entity: job.entity,
name: job.name,
action,
elapsed: started.elapsed().as_secs_f32(),
});
}
}
fn request_action(
client: &reqwest::blocking::Client,
url: &str,
system: &str,
user: &str,
) -> Result<Action, String> {
let schema = json!({
"type": "object",
"properties": {
"action": {"type": "string",
"enum": ["go_to", "talk", "work", "rest", "idle"]},
"place": {"type": "string"},
"to": {"type": "string"},
"say": {"type": "string"},
"reason": {"type": "string"}
},
"required": ["action", "reason"]
});
let body = json!({
"model": MODEL,
"stream": false,
"think": false,
"format": schema,
"messages": [
{"role": "system", "content": system},
{"role": "user", "content": user},
],
"options": {
"temperature": 0.9,
"num_predict": THINK_BUDGET,
"num_ctx": 4096,
"keep_alive": "30m",
}
});
let resp = client
.post(format!("{url}/api/chat"))
.json(&body)
.send()
.map_err(|e| format!("http: {e}"))?;
if !resp.status().is_success() {
return Err(format!("status {}", resp.status()));
}
let value: serde_json::Value = resp.json().map_err(|e| format!("json: {e}"))?;
let content = value["message"]["content"].as_str().unwrap_or_default();
let content = content
.trim()
.trim_start_matches("```json")
.trim_start_matches("```")
.trim_end_matches("```")
.trim();
serde_json::from_str::<Action>(content).map_err(|e| format!("{e}: {content:.120}"))
}
// ---------------------------------------------------------------------------
// Perception helpers
// ---------------------------------------------------------------------------
/// Rough named place for a position (nearest named POI within 6 units).
fn place_at(map: &WorldMap, pos: Vec3) -> String {
let cell = cell_of(pos);
let mut best = format!("({}, {})", cell.x, cell.y);
let mut best_d = 6.0f32;
for (name, poi) in &map.pois {
if name.starts_with("home_") {
continue;
}
let d = (world_of(*poi) - pos).length();
if d < best_d {
best_d = d;
best = name.clone();
}
}
best
}
fn build_nearby(
map: &WorldMap,
positions: &[(Entity, Vec3, String)],
me: Entity,
) -> Vec<(String, String)> {
let my_pos = match positions.iter().find(|(e, _, _)| *e == me) {
Some((_, p, _)) => *p,
None => return vec![],
};
let mut out = Vec::new();
for (e, pos, name) in positions {
if *e == me || pos.distance(my_pos) > 12.0 || out.len() >= 5 {
continue;
}
out.push((name.clone(), place_at(map, *pos)));
}
out
}
// ---------------------------------------------------------------------------
// Scheduling
// ---------------------------------------------------------------------------
pub fn schedule_wakes(
time: Res<Time>,
mut pool: ResMut<LlmPool>,
map: Res<WorldMap>,
mut query: Query<(Entity, &Villager, &Needs, &Memory, &Transform, &mut Brain, &Pathing)>,
all: Query<(Entity, &Transform, &Villager)>,
) {
if pool.in_flight >= WORKERS {
return;
}
let now = time.elapsed_secs();
let positions: Vec<(Entity, Vec3, String)> = all
.iter()
.map(|(e, t, v)| (e, t.translation, v.persona.name.clone()))
.collect();
for (entity, villager, needs, memory, transform, mut brain, pathing) in query.iter_mut() {
if brain.in_flight || now < brain.next_wake || pool.in_flight >= WORKERS {
continue;
}
let p = &villager.persona;
let home = map
.place(&format!("home_{}", p.home_index))
.map(|c| (c.x, c.y))
.unwrap_or((32, 32));
let at_place = if pathing.is_moving() {
format!("walking to {}", pathing.destination)
} else {
place_at(&map, transform.translation)
};
let nearby = build_nearby(&map, &positions, entity);
let nearby_str = if nearby.is_empty() {
"nobody is near you".to_string()
} else {
nearby
.iter()
.map(|(n, pl)| format!("{n} (at {pl})"))
.collect::<Vec<_>>()
.join(", ")
};
let memories: Vec<String> = memory.recent(6).cloned().collect();
let memories_str = if memories.is_empty() {
" (you remember nothing yet)".to_string()
} else {
memories
.iter()
.enumerate()
.map(|(i, m)| format!(" {}. {m}", i + 1))
.collect::<Vec<_>>()
.join("\n")
};
let reply_hint = if brain.expect_reply {
"\nSomeone recently spoke to you. You may respond to them with a talk action, \
or do whatever you actually want to do."
} else {
""
};
let user = format!(
"[minute {:.0}] You are at: {at_place}.\n\
Needs: hunger {hunger:.0}/100 (higher = hungrier), \
energy {energy:.0}/100, social {social:.0}/100 (lower = lonelier).\n\
Nearby: {nearby_str}.{reply_hint}\n\
Your memories:\n{memories_str}\n\
Choose your one action as JSON.",
now / 60.0,
hunger = needs.hunger,
energy = needs.energy,
social = needs.social,
);
let job = DecisionJob {
entity,
name: p.name.clone(),
system: p.system_prompt(home),
user,
};
if pool.jobs_tx.send(job).is_ok() {
brain.in_flight = true;
brain.expect_reply = false;
pool.in_flight += 1;
pool.peak_in_flight = pool.peak_in_flight.max(pool.in_flight);
}
}
}
// ---------------------------------------------------------------------------
// Results + action application
// ---------------------------------------------------------------------------
pub fn collect_results(
time: Res<Time>,
mut pool: ResMut<LlmPool>,
map: Res<WorldMap>,
mut events: MessageWriter<VillagerEvent>,
mut query: Query<(
Entity,
&Villager,
&Transform,
&mut Brain,
&mut Pathing,
&mut Speech,
&mut Memory,
&mut Needs,
)>,
) {
loop {
// Guard must drop before we mutate `pool` (edition-2021 while-let
// would keep it alive through the body).
let next = {
let rx = pool.results_rx.lock().unwrap_or_else(|e| e.into_inner());
rx.try_recv()
};
let Ok(result) = next else { break };
pool.in_flight = pool.in_flight.saturating_sub(1);
let now = time.elapsed_secs();
let Ok((entity, villager, transform, mut brain, mut pathing, mut speech, mut memory, mut needs)) =
query.get_mut(result.entity)
else {
continue;
};
brain.in_flight = false;
brain.next_wake = now + WAKE_MIN + jitter(&villager.persona.name, now) * WAKE_JITTER;
match result.action {
Ok(action) => {
pool.completed += 1;
apply_action(
action,
&map,
&villager.persona.name,
transform.translation,
&mut pathing,
&mut speech,
&mut memory,
&mut needs,
now,
&mut events,
entity,
);
}
Err(err) => {
pool.failed += 1;
brain.next_wake = now + 10.0; // don't stall the person
memory.push(format!("(could not decide: {err})"));
}
}
}
}
#[allow(clippy::too_many_arguments)]
fn apply_action(
action: Action,
map: &WorldMap,
my_name: &str,
my_pos: Vec3,
pathing: &mut Pathing,
speech: &mut Speech,
memory: &mut Memory,
needs: &mut Needs,
now: f32,
events: &mut MessageWriter<VillagerEvent>,
me: Entity,
) {
match action {
Action::GoTo { place, reason } => {
let target = map
.place(&place)
.or_else(|| map.place("plaza"))
.unwrap_or(IVec2::new(32, 32));
match map.path(cell_of(my_pos), target) {
Some(cells) => {
pathing.cells = cells;
pathing.step_progress = 0.0;
pathing.destination = place.clone();
memory.push(format!("Decided to go to the {place}: {reason}"));
events.write(VillagerEvent::Action { entity: me, text: format!("→ {place}") });
}
None => {
memory.push(format!("Wanted to go to the {place} but could not find a way."));
}
}
needs.social = (needs.social - 0.2).max(0.0);
}
Action::Talk { to, say, reason } => {
speech.text = say.clone();
speech.expires = now + 7.0;
memory.push(format!("Said to {to}: \"{say}\" ({reason})"));
events.write(VillagerEvent::Talk { from: my_name.to_string(), to, say });
}
Action::Work { reason } => {
memory.push(format!("Worked ({reason})"));
needs.energy = (needs.energy - 4.0).max(0.0);
needs.hunger = (needs.hunger + 2.0).min(100.0);
events.write(VillagerEvent::Action { entity: me, text: "work".into() });
}
Action::Rest { reason } => {
memory.push(format!("Rested ({reason})"));
needs.energy = (needs.energy + 12.0).min(100.0);
events.write(VillagerEvent::Action { entity: me, text: "rest".into() });
}
Action::Idle { reason } => {
memory.push(format!("Loitered ({reason})"));
events.write(VillagerEvent::Action { entity: me, text: "idle".into() });
}
}
}
/// Deterministic pseudo-jitter per (name, cycle) — no rand crate.
fn jitter(name: &str, t: f32) -> f32 {
let mut h: u32 = 2166136261;
for b in name.bytes() {
h ^= b as u32;
h = h.wrapping_mul(16777619);
}
h ^= (t / 1000.0) as u32;
(h % 1000) as f32 / 1000.0
}
#[derive(Message, Debug)]
pub enum VillagerEvent {
Talk { from: String, to: String, say: String },
Action { entity: Entity, text: String },
}
/// Record what was said into the listener's memory + wake them to respond.
pub fn deliver_talks(
time: Res<Time>,
mut events: MessageReader<VillagerEvent>,
names: Res<Names>,
mut query: Query<(&mut Memory, &mut Brain)>,
) {
let now = time.elapsed_secs();
for ev in events.read() {
if let VillagerEvent::Talk { from, to, say } = ev {
if let Some(&target) = names.0.get(&to.to_lowercase()) {
if let Ok((mut memory, mut brain)) = query.get_mut(target) {
memory.push(format!("{from} said: \"{say}\""));
if !brain.in_flight {
brain.next_wake = now + 1.0 + jitter(to, now) * 2.0;
brain.expect_reply = true;
}
}
}
}
}
}