poimen-memory/crates/mem-core/src/gated_loop.rs

use crate::domain::{Chunk, Level};
use crate::gate_parser::parse_gate_response;
use crate::prompt::PromptBuilder;
use crate::query::Query;
use anyhow::Result;

/// LLM client trait for dependency injection.
pub trait LlmClient: Send + Sync {
    fn complete_blocking(&self, system: &str, user: &str, max_tokens: usize) -> Result<String>;
}

/// Loop configuration.
#[derive(Debug, Clone)]
pub struct LoopConfig {
    pub level: Level,
    pub query: Query,
    pub memory_budget: u32,
    pub use_exit_gate: bool,
}

/// Events emitted by the loop.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LoopEvent {
    Evidence { turn: u32 },
    Memory { turn: u32, update: bool },
    Gate { turn: u32, update: bool, exit: bool },
    ParseFailed { turn: u32, attempts: u32 },
    BudgetExceeded { turn: u32 },
    RunEnd { chunks_seen: u32, chunks_used: u32 },
}

/// Outcome of a loop run.
#[derive(Debug, Clone)]
pub struct RunOutcome {
    pub chunks_seen: u32,
    pub chunks_used: u32,
    pub final_memory: String,
    pub events: Vec<LoopEvent>,
}

/// Run the gated loop over chunks.
pub fn run_loop(
    config: LoopConfig,
    chunks: Vec<Chunk>,
    llm: &dyn LlmClient,
) -> Result<RunOutcome> {
    let mut memory = String::new();
    let mut chunks_seen = 0u32;
    let mut chunks_used = 0u32;
    let mut events = Vec::new();

    for chunk in chunks {
        chunks_seen += 1;
        let turn = chunks_seen;

        // Build prompt
        let memory_ref = if memory.is_empty() { None } else { Some(memory.as_str()) };
        let (system_prompt, user_prompt) = PromptBuilder::build(&config.query, memory_ref, &chunk)?;

        // Try parse up to 3 times
        let mut should_exit = false;
        let mut parse_ok = false;
        
        for attempt in 1..=3 {
            match llm.complete_blocking(&system_prompt, &user_prompt, 2048) {
                Ok(response) => match parse_gate_response(&response) {
                    Ok(gated) => {
                        // Check memory budget
                        if gated.candidate.len() as u32 > config.memory_budget {
                            events.push(LoopEvent::BudgetExceeded { turn });
                            events.push(LoopEvent::Gate {
                                turn,
                                update: false,
                                exit: gated.exit_gate,
                            });
                            parse_ok = true;
                            should_exit = gated.exit_gate && config.use_exit_gate;
                            break;
                        }

                        // Apply update rule
                        if gated.update_gate {
                            memory = gated.candidate.clone();
                            chunks_used += 1;
                            events.push(LoopEvent::Evidence { turn });
                        }

                        events.push(LoopEvent::Memory {
                            turn,
                            update: gated.update_gate,
                        });
                        events.push(LoopEvent::Gate {
                            turn,
                            update: gated.update_gate,
                            exit: gated.exit_gate,
                        });

                        parse_ok = true;
                        should_exit = gated.exit_gate && config.use_exit_gate;
                        break;
                    }
                    Err(_) if attempt < 3 => continue,
                    Err(_) => {
                        events.push(LoopEvent::ParseFailed { turn, attempts: attempt });
                        parse_ok = true;
                        break;
                    }
                },
                Err(_) if attempt < 3 => continue,
                Err(e) => return Err(e),
            }
        }

        if !parse_ok {
            return Err(anyhow::anyhow!("Failed to parse after all retries"));
        }

        if should_exit {
            break;
        }
    }

    events.push(LoopEvent::RunEnd {
        chunks_seen,
        chunks_used,
    });

    Ok(RunOutcome {
        chunks_seen,
        chunks_used,
        final_memory: memory,
        events,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_loop_basic() {
        // Placeholder test to verify it compiles
        assert!(true);
    }
}