use std::{collections::HashMap, path::PathBuf, time::SystemTime};
use miette::{Result, miette};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
use tokio::fs;
use crate::{
daat_locus_paths::daat_locus_paths,
persistence::{PersistenceFileMode, write_bytes_atomic},
};
use super::{
examples::{ExampleField, ProgramExample},
optimizer::PromptTuningConfig,
program::Program,
};
pub const COMPILED_DIR_NAME: &str = "reasoning_compiled";
pub const RUNTIME_SYSTEM_PROMPT_COMPILE_KEY: &str = "runtime_agent_system";
pub const RUNTIME_SYSTEM_PROMPT_PREVIOUS_COMPILE_KEY: &str = "runtime_agent_system_previous";
fn model_scoped_runtime_compile_key(base: &str, model_name: &str) -> String {
let normalized = model_name
.trim()
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() {
ch.to_ascii_lowercase()
} else {
'-'
}
})
.collect::<String>()
.trim_matches('-')
.to_string();
if normalized.is_empty() {
base.to_string()
} else {
format!("{base}--{normalized}")
}
}
fn model_scope_suffix(model_name: &str) -> String {
model_scoped_runtime_compile_key("", model_name)
}
fn stem_has_model_scope(stem: &str, model_name: &str) -> bool {
let suffix = model_scope_suffix(model_name);
!suffix.is_empty() && stem.ends_with(&suffix)
}
#[derive(Clone, Serialize, Deserialize)]
pub struct StoredProgramExample {
pub title: String,
pub inputs: Vec<ExampleField>,
pub output: Value,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct StoredPromptTuningConfig {
pub extra_instructions: Vec<String>,
pub examples: Vec<StoredProgramExample>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledProgram {
pub suite: String,
pub compile_key: String,
pub best_candidate: String,
pub score: usize,
pub total_cases: usize,
pub tuning: StoredPromptTuningConfig,
#[serde(default)]
pub report: Option<CompiledProgramReport>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledProgramReport {
pub train_score: usize,
pub train_total_cases: usize,
pub train_attempts_used: usize,
#[serde(default)]
pub acceptance_score: Option<usize>,
#[serde(default)]
pub acceptance_total_cases: Option<usize>,
#[serde(default)]
pub acceptance_attempts_used: Option<usize>,
pub dev_score: usize,
pub dev_total_cases: usize,
pub dev_attempts_used: usize,
#[serde(default)]
pub ranking_label: Option<String>,
pub selected_extra_instructions: Vec<String>,
pub selected_example_titles: Vec<String>,
pub candidates: Vec<CompiledCandidateReport>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledCandidateReport {
pub name: String,
#[serde(default)]
pub acceptance_score: Option<usize>,
#[serde(default)]
pub acceptance_total_cases: Option<usize>,
#[serde(default)]
pub acceptance_attempts_used: Option<usize>,
pub score: usize,
pub total_cases: usize,
pub attempts_used: usize,
#[serde(default)]
pub judge_wins: usize,
#[serde(default)]
pub judge_losses: usize,
#[serde(default)]
pub judge_ties: usize,
pub extra_instructions: Vec<String>,
pub example_titles: Vec<String>,
pub failed_cases: Vec<CompiledFailureCaseReport>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledFailureCaseReport {
pub case_name: String,
pub detail: String,
}
#[derive(Clone, Default)]
pub struct CompiledPromptStore {
entries: HashMap<String, CompiledProgram>,
runtime_system_prompt: Option<CompiledRuntimeSystemPrompt>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledRuntimeSystemPrompt {
pub compile_key: String,
pub best_candidate: String,
#[serde(default)]
pub system_additions: Vec<String>,
#[serde(default)]
pub selected_demo_titles: Vec<String>,
#[serde(default)]
pub report: Option<CompiledRuntimeSystemPromptReport>,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct CompiledRuntimeSystemPromptReport {
pub score: usize,
pub total_cases: usize,
#[serde(default)]
pub judge_summary: Option<String>,
}
impl CompiledRuntimeSystemPrompt {
pub fn with_compile_key(mut self, compile_key: impl Into<String>) -> Self {
self.compile_key = compile_key.into();
self
}
}
impl CompiledPromptStore {
pub fn from_entries(entries: Vec<CompiledProgram>) -> Self {
let entries = entries
.into_iter()
.map(|entry| (entry.suite.clone(), entry))
.collect();
Self {
entries,
runtime_system_prompt: None,
}
}
pub fn with_runtime_system_prompt(
mut self,
runtime_system_prompt: Option<CompiledRuntimeSystemPrompt>,
) -> Self {
self.runtime_system_prompt = runtime_system_prompt;
self
}
pub fn get_tuning<P: Program>(&self, program: &P) -> Option<PromptTuningConfig<P::Output>> {
self.entries
.get(&program.tuning_key())
.and_then(|entry| entry.tuning.to_typed::<P::Output>().ok())
}
pub fn runtime_system_additions(&self) -> &[String] {
self.runtime_system_prompt
.as_ref()
.map(|prompt| prompt.system_additions.as_slice())
.unwrap_or(&[])
}
}
impl StoredPromptTuningConfig {
pub fn from_typed<O: Serialize + Clone + DeserializeOwned>(
tuning: &PromptTuningConfig<O>,
) -> Self {
Self {
extra_instructions: tuning.extra_instructions.clone(),
examples: tuning
.examples
.iter()
.map(|example| StoredProgramExample {
title: example.title.clone(),
inputs: example.inputs.clone(),
output: serde_json::to_value(&example.output).unwrap(),
})
.collect(),
}
}
pub fn to_typed<O: Serialize + Clone + DeserializeOwned>(
&self,
) -> Result<PromptTuningConfig<O>> {
let mut examples = Vec::with_capacity(self.examples.len());
for example in &self.examples {
let output = serde_json::from_value::<O>(example.output.clone())
.map_err(|err| miette!("failed to deserialize stored example output: {err}"))?;
examples.push(ProgramExample {
title: example.title.clone(),
inputs: example.inputs.clone(),
output,
});
}
Ok(PromptTuningConfig {
extra_instructions: self.extra_instructions.clone(),
examples,
})
}
}
pub async fn load_all_compiled_programs_for_model(
model_name: &str,
) -> Result<Vec<CompiledProgram>> {
load_all_compiled_programs_from_dir(COMPILED_DIR_NAME, model_name).await
}
pub async fn load_all_compiled_programs_from_dir(
dir_name: &str,
model_name: &str,
) -> Result<Vec<CompiledProgram>> {
let dir = compiled_dir_named(dir_name).await;
let Ok(mut entries) = fs::read_dir(&dir).await else {
return Ok(Vec::new());
};
let mut newest_by_suite: HashMap<String, (SystemTime, CompiledProgram)> = HashMap::new();
while let Some(entry) = entries
.next_entry()
.await
.map_err(|err| miette!("failed to iterate compiled dir {}: {err}", dir.display()))?
{
let path = entry.path();
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
let stem = path
.file_stem()
.and_then(|stem| stem.to_str())
.unwrap_or("");
if stem == RUNTIME_SYSTEM_PROMPT_COMPILE_KEY
|| stem == RUNTIME_SYSTEM_PROMPT_PREVIOUS_COMPILE_KEY
|| stem.starts_with(&format!("{RUNTIME_SYSTEM_PROMPT_COMPILE_KEY}--"))
|| stem.starts_with(&format!("{RUNTIME_SYSTEM_PROMPT_PREVIOUS_COMPILE_KEY}--"))
{
continue;
}
if !stem_has_model_scope(stem, model_name) {
continue;
}
let bytes = fs::read(&path).await.map_err(|err| {
miette!(
"failed to read compiled prompt config {}: {err}",
path.display()
)
})?;
let program = serde_json::from_slice::<CompiledProgram>(&bytes).map_err(|err| {
miette!(
"failed to decode compiled prompt config {}: {err}",
path.display()
)
})?;
let modified = entry
.metadata()
.await
.ok()
.and_then(|metadata| metadata.modified().ok())
.unwrap_or(SystemTime::UNIX_EPOCH);
match newest_by_suite.get(&program.suite) {
Some((existing_modified, _)) if *existing_modified >= modified => {}
_ => {
newest_by_suite.insert(program.suite.clone(), (modified, program));
}
}
}
Ok(newest_by_suite
.into_values()
.map(|(_, program)| program)
.collect())
}
pub async fn save_compiled_program_for_model(
model_name: &str,
compiled: &CompiledProgram,
) -> Result<()> {
let dir = compiled_dir_named(COMPILED_DIR_NAME).await;
if !dir.exists() {
fs::create_dir_all(&dir)
.await
.map_err(|err| miette!("failed to create compiled prompt dir: {err}"))?;
}
let compile_key = model_scoped_runtime_compile_key(&compiled.compile_key, model_name);
let path = dir.join(format!("{compile_key}.json"));
let bytes = serde_json::to_vec_pretty(compiled)
.map_err(|err| miette!("failed to serialize compiled prompt config: {err}"))?;
write_bytes_atomic(path, bytes, PersistenceFileMode::Default)
.await
.map_err(|err| miette!("failed to write compiled prompt config: {err}"))?;
Ok(())
}
pub async fn seed_compiled_program_from_tuning_for_model<P: Program>(
model_name: &str,
program: &P,
tuning: &PromptTuningConfig<P::Output>,
) -> Result<()> {
let compiled = CompiledProgram {
suite: program.tuning_key(),
compile_key: program.tuning_key(),
best_candidate: "default_seed".to_string(),
score: 0,
total_cases: 0,
tuning: StoredPromptTuningConfig::from_typed(tuning),
report: None,
};
save_compiled_program_for_model(model_name, &compiled).await
}
pub async fn load_compiled_runtime_system_prompt_for_model(
model_name: &str,
) -> Result<Option<CompiledRuntimeSystemPrompt>> {
let scoped_key =
model_scoped_runtime_compile_key(RUNTIME_SYSTEM_PROMPT_COMPILE_KEY, model_name);
load_compiled_runtime_system_prompt_by_key(&scoped_key).await
}
pub async fn save_compiled_runtime_system_prompt_for_model(
model_name: &str,
compiled: &CompiledRuntimeSystemPrompt,
) -> Result<()> {
let scoped_key =
model_scoped_runtime_compile_key(RUNTIME_SYSTEM_PROMPT_COMPILE_KEY, model_name);
save_compiled_runtime_system_prompt_by_key(&scoped_key, compiled).await
}
async fn load_compiled_runtime_system_prompt_by_key(
compile_key: &str,
) -> Result<Option<CompiledRuntimeSystemPrompt>> {
let path = compiled_dir_named(COMPILED_DIR_NAME)
.await
.join(format!("{compile_key}.json"));
let Ok(bytes) = fs::read(path).await else {
return Ok(None);
};
let compiled = serde_json::from_slice::<CompiledRuntimeSystemPrompt>(&bytes)
.map_err(|err| miette!("failed to decode runtime system prompt config: {err}"))?;
Ok(Some(compiled))
}
async fn save_compiled_runtime_system_prompt_by_key(
compile_key: &str,
compiled: &CompiledRuntimeSystemPrompt,
) -> Result<()> {
let dir = compiled_dir_named(COMPILED_DIR_NAME).await;
if !dir.exists() {
fs::create_dir_all(&dir)
.await
.map_err(|err| miette!("failed to create compiled prompt dir: {err}"))?;
}
let path = dir.join(format!("{compile_key}.json"));
let bytes = serde_json::to_vec_pretty(&compiled.clone().with_compile_key(compile_key))
.map_err(|err| miette!("failed to serialize runtime system prompt config: {err}"))?;
write_bytes_atomic(path, bytes, PersistenceFileMode::Default)
.await
.map_err(|err| miette!("failed to write runtime system prompt config: {err}"))?;
Ok(())
}
async fn compiled_dir_named(dir_name: &str) -> PathBuf {
daat_locus_paths().await.artifact_dir(dir_name)
}