146 lines
4.3 KiB
Rust
146 lines
4.3 KiB
Rust
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);
|
|
}
|
|
}
|