overbrainer 0.10.0

//! The pipeline commands: `subtopics`, `questions`, `answers`, `split` and `run`.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{Context, anyhow, bail};
use serde_json::Value;
use tokio::sync::{Mutex, OnceCell};

use super::front::Frontend;
use super::record::Recorder;
use super::reload::Reloader;
use super::{LazyVault, StageArgs};
use crate::config::{RoleModel, Settings, Source};
use crate::dataset::{DataFiles, Subtopic, read};
use crate::dedup::{Embedding, Layered, Lexical};
use crate::events::{EventBus, Stage, StageStats};
use crate::history::Status;
use crate::llm::{ProtocolClient, connect};
use crate::pipeline::{self, Ctx, PipelineError, RoleClient, SplitReport};
use crate::pricing::{LISTING_TIMEOUT, Price, fetch_listing, listed_price};
use crate::prompts::Prompts;
use crate::secrets::Resolver;

/// Which pipeline command to run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
    /// `overbrainer subtopics`
    Subtopics,
    /// `overbrainer questions`
    Questions,
    /// `overbrainer answers`
    Answers,
    /// `overbrainer split`
    Split,
    /// `overbrainer run`: the four stages in order (training follows in `cli::run`).
    Run,
}

type Client = RoleClient<ProtocolClient>;

/// How a pipeline command reads its settings.
pub(crate) enum Load<'a> {
    /// Once, from this source.
    Source(&'a Source),
    /// With the environment of this reloader, and, for `run`, again before
    /// each stage after the first when the files changed (on the command line).
    Reload(&'a mut Reloader),
}

/// Everything a pipeline command needs. Role clients are built on first use and at
/// most once per command, and each provider's model listing is read at most once.
struct Session {
    project_dir: PathBuf,
    settings: Settings,
    files: DataFiles,
    resolver: Resolver<LazyVault>,
    generator: OnceCell<Client>,
    parent: OnceCell<Arc<Client>>,
    embedder: OnceCell<Option<ProtocolClient>>,
    listings: Mutex<BTreeMap<String, Option<Value>>>,
}

impl Session {
    fn open(project_dir: &Path, source: &Source) -> anyhow::Result<Self> {
        let settings = source.load(project_dir)?;
        Ok(Self::new(project_dir, settings))
    }

    /// A session on `settings`: no client built yet, no listing read.
    fn new(project_dir: &Path, settings: Settings) -> Self {
        Self {
            project_dir: project_dir.to_path_buf(),
            settings,
            files: DataFiles::new(project_dir),
            resolver: super::resolver(),
            generator: OnceCell::new(),
            parent: OnceCell::new(),
            embedder: OnceCell::new(),
            listings: Mutex::new(BTreeMap::new()),
        }
    }

    /// The role `stage` uses; none for `split`.
    fn role_of(&self, stage: Stage) -> Option<&RoleModel> {
        let roles = &self.settings.roles;
        match stage {
            Stage::Subtopics | Stage::Questions => Some(&roles.generator),
            Stage::Answers => Some(&roles.parent),
            Stage::Split => None,
        }
    }

    fn ctx<'a>(&'a self, bus: &'a EventBus, args: &'a StageArgs, prompts: &'a Prompts) -> Ctx<'a> {
        Ctx {
            settings: &self.settings,
            files: &self.files,
            prompts,
            bus,
            topic: args.topic.as_deref(),
            force: args.force,
        }
    }

    async fn generator(&self) -> anyhow::Result<&Client> {
        self.generator
            .get_or_try_init(|| self.role(&self.settings.roles.generator))
            .await
    }

    async fn parent(&self) -> anyhow::Result<Arc<Client>> {
        let parent = self
            .parent
            .get_or_try_init(|| async {
                Ok::<_, anyhow::Error>(Arc::new(self.role(&self.settings.roles.parent).await?))
            })
            .await?;
        Ok(Arc::clone(parent))
    }

    async fn embedder(&self) -> anyhow::Result<Option<ProtocolClient>> {
        let embedder = self
            .embedder
            .get_or_try_init(|| async {
                match &self.settings.roles.embedder {
                    Some(role) => connect(&self.settings, role, &self.resolver)
                        .await
                        .map(Some)
                        .context("cannot set up the embedder"),
                    None => Ok(None),
                }
            })
            .await?;
        Ok(embedder.clone())
    }

    /// Connects to the provider of `role` and looks up the model price.
    async fn role(&self, role: &RoleModel) -> anyhow::Result<Client> {
        let client = connect(&self.settings, role, &self.resolver).await?;
        let price = self.price(role, &client).await;
        Ok(RoleClient {
            client,
            model: role.clone(),
            price,
        })
    }

    /// Price of `role`'s model, from its provider's listing read once per command.
    async fn price(&self, role: &RoleModel, client: &ProtocolClient) -> Option<Price> {
        let mut listings = self.listings.lock().await;
        if !listings.contains_key(&role.provider) {
            let listing = fetch_listing(client, LISTING_TIMEOUT).await;
            listings.insert(role.provider.clone(), listing);
        }
        let listing = listings.get(&role.provider)?.as_ref()?;
        listed_price(listing, &role.model)
    }
}

/// Runs `command` in `project_dir` with the settings `load` reads, publishing
/// progress on `front`'s bus, emitting a usage summary per stage through `front`
/// (stdout on the command line) and recording each stage it starts in the
/// project's history. The command stops at `front`'s interruption (Ctrl-C on the
/// command line).
///
/// # Errors
///
/// Returns an error if the configuration, a template or a data file cannot be loaded,
/// if `--topic` names no configured topic, if a provider cannot be reached, if items
/// failed (they are retried on the next run), or if the command was interrupted.
pub async fn run(
    project_dir: &Path,
    command: Command,
    args: &StageArgs,
    front: &Frontend,
    load: Load<'_>,
) -> anyhow::Result<()> {
    let (source, reloader) = match load {
        Load::Source(source) => (source.clone(), None),
        Load::Reload(reloader) => (Source::from(reloader.env().clone()), Some(reloader)),
    };
    let session = Session::open(project_dir, &source)?;
    let guard = front.open_bus();
    let mut recorder = Recorder::new(project_dir, &guard.bus);
    let steps = Steps {
        session: &session,
        front,
        recorder: &mut recorder,
    };
    let flow = execute(&guard.bus, command, args, steps, reloader);
    let result = if let Some(result) = front.interrupt().race(flow).await {
        result
    } else {
        recorder.interrupted().await;
        Err(anyhow!("interrupted: the stage resumes on its next run"))
    };
    drop(recorder);
    guard.close().await;
    result
}

async fn execute(
    bus: &EventBus,
    command: Command,
    args: &StageArgs,
    mut steps: Steps<'_>,
    reloader: Option<&mut Reloader>,
) -> anyhow::Result<()> {
    let session = steps.session;
    let empty = Prompts::empty();
    session.ctx(bus, args, &empty).topics()?;
    let prompts = if command == Command::Split {
        empty
    } else {
        Prompts::load(&session.project_dir)?
    };
    let ctx = session.ctx(bus, args, &prompts);
    match command {
        Command::Subtopics => {
            steps
                .stage(Stage::Subtopics, subtopics(session, &ctx))
                .await
        },
        Command::Questions => {
            if missing_subtopics(&ctx)? {
                let first = Ctx {
                    force: false,
                    ..ctx
                };
                steps
                    .stage(Stage::Subtopics, subtopics(session, &first))
                    .await?;
            }
            steps
                .stage(Stage::Questions, questions(session, &ctx))
                .await
        },
        Command::Answers => steps.stage(Stage::Answers, answers(session, &ctx)).await,
        Command::Split => steps.split(&ctx).await,
        // Boxed: its sessions and stages would weigh on every command's future.
        Command::Run => Box::pin(run_stages(bus, args, steps, reloader)).await,
    }
}

/// The stages of `run`, in order.
const RUN_STAGES: [Stage; 4] = [
    Stage::Subtopics,
    Stage::Questions,
    Stage::Answers,
    Stage::Split,
];

/// Runs the stages of `run` with the session of `steps`; before each stage
/// after the first, `reloader`, when given, reads the settings again when
/// their files changed, and the next stages use them.
async fn run_stages(
    bus: &EventBus,
    args: &StageArgs,
    steps: Steps<'_>,
    mut reloader: Option<&mut Reloader>,
) -> anyhow::Result<()> {
    let Steps {
        session: first,
        front,
        recorder,
    } = steps;
    let mut reread: Option<Session> = None;
    for (index, stage) in RUN_STAGES.into_iter().enumerate() {
        if index > 0
            && let Some(reloader) = reloader.as_deref_mut()
            && let Some(settings) = reloader.changed()?
        {
            reread = Some(Session::new(&first.project_dir, settings));
        }
        let session = reread.as_ref().unwrap_or(first);
        let prompts = if stage == Stage::Split {
            Prompts::empty()
        } else {
            Prompts::load(&session.project_dir)?
        };
        let ctx = session.ctx(bus, args, &prompts);
        let mut steps = Steps {
            session,
            front,
            recorder: &mut *recorder,
        };
        match stage {
            Stage::Subtopics => steps.stage(stage, subtopics(session, &ctx)).await?,
            Stage::Questions => steps.stage(stage, questions(session, &ctx)).await?,
            Stage::Answers => steps.stage(stage, answers(session, &ctx)).await?,
            Stage::Split => steps.split(&ctx).await?,
        }
    }
    Ok(())
}

/// Runs stages one at a time, reporting and recording each.
struct Steps<'a> {
    session: &'a Session,
    front: &'a Frontend,
    recorder: &'a mut Recorder,
}

impl Steps<'_> {
    /// Runs `stage` through `work`, then reports and records how it ended.
    async fn stage(
        &mut self,
        stage: Stage,
        work: impl Future<Output = anyhow::Result<StageStats>>,
    ) -> anyhow::Result<()> {
        self.recorder.begin(stage, self.session.role_of(stage));
        report(stage, work.await, self.front, self.recorder).await
    }

    async fn split(&mut self, ctx: &Ctx<'_>) -> anyhow::Result<()> {
        self.recorder.begin(Stage::Split, None);
        match pipeline::split(ctx) {
            Ok(report) => {
                self.recorder.split(&report);
                print_split(&report, self.front);
                Ok(())
            },
            Err(error) => {
                self.recorder.fail().await;
                Err(error.into())
            },
        }
    }
}

/// Whether a selected topic has no subtopic yet (`questions` then generates them).
fn missing_subtopics(ctx: &Ctx<'_>) -> anyhow::Result<bool> {
    let existing: Vec<Subtopic> = read(&ctx.files.subtopics)?;
    Ok(ctx
        .topics()?
        .iter()
        .any(|topic| !existing.iter().any(|subtopic| subtopic.topic == topic.name)))
}

async fn subtopics(session: &Session, ctx: &Ctx<'_>) -> anyhow::Result<StageStats> {
    let generator = session.generator().await?;
    Ok(pipeline::subtopics(ctx, generator).await?)
}

async fn questions(session: &Session, ctx: &Ctx<'_>) -> anyhow::Result<StageStats> {
    let generator = session.generator().await?;
    let embedder = session.embedder().await?;
    let pipeline = &session.settings.pipeline;
    let policy = ctx.policy();
    let new_dedup = || {
        Layered::new(
            Lexical::new(pipeline.dedup_threshold),
            embedder
                .clone()
                .map(|client| Embedding::new(client, pipeline.embedding_threshold, policy)),
        )
    };
    Ok(pipeline::questions(ctx, generator, new_dedup).await?)
}

async fn answers(session: &Session, ctx: &Ctx<'_>) -> anyhow::Result<StageStats> {
    Ok(pipeline::answers(ctx, session.parent().await?).await?)
}

/// Emits the stage summary through `front`, records how the stage ended, and fails
/// when items failed. A stage stopped by a fatal provider error still emits and
/// records what it produced and spent before stopping; a stage stopped by another
/// error records what its events counted.
async fn report(
    stage: Stage,
    result: anyhow::Result<StageStats>,
    front: &Frontend,
    recorder: &mut Recorder,
) -> anyhow::Result<()> {
    let stats = match result {
        Ok(stats) => stats,
        Err(error) => {
            if let Some(PipelineError::Llm { stage, spent, .. }) = error.downcast_ref() {
                front.line(&summary(*stage, spent));
                recorder.end(Status::Failed, spent);
            } else {
                recorder.fail().await;
            }
            return Err(error);
        },
    };
    front.line(&summary(stage, &stats));
    let outcome = if stats.failed == 0 {
        Status::Ok
    } else {
        Status::Failed
    };
    recorder.end(outcome, &stats);
    if stats.failed > 0 {
        bail!(
            "{} {stage} item(s) failed; run the command again to retry them",
            stats.failed
        );
    }
    Ok(())
}

fn print_split(report: &SplitReport, front: &Frontend) {
    front.line(&split_summary(report));
}

/// One line: counts, tokens, and cost when the price is known.
#[must_use]
pub fn summary(stage: Stage, stats: &StageStats) -> String {
    let cost = stats.cost.map_or_else(
        || "cost unknown".to_string(),
        |cost| format!("cost ${cost:.4}"),
    );
    format!(
        "{stage}: {} done, {} skipped, {} failed, {} excluded; tokens {} in, {} out; {cost}",
        stats.done,
        stats.skipped,
        stats.failed,
        stats.excluded,
        stats.usage.input_tokens,
        stats.usage.output_tokens,
    )
}

/// One line: train and eval sizes, exclusions by reason, then orphaned examples.
#[must_use]
pub fn split_summary(report: &SplitReport) -> String {
    let excluded: usize = report.excluded.values().sum();
    let reasons: Vec<String> = report
        .excluded
        .iter()
        .map(|(reason, count)| {
            let name = serde_json::to_value(reason)
                .ok()
                .and_then(|value| value.as_str().map(str::to_string))
                .unwrap_or_default();
            format!("{name} {count}")
        })
        .collect();
    let detail = if reasons.is_empty() {
        String::new()
    } else {
        format!(" ({})", reasons.join(", "))
    };
    format!(
        "split: {} train, {} eval, {excluded} excluded{detail}, {} orphaned",
        report.train, report.eval, report.orphaned
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dataset::Exclusion;
    use crate::llm::Usage;

    #[test]
    fn summary_shows_cost_only_when_known() {
        let mut stats = StageStats {
            done: 3,
            usage: Usage {
                input_tokens: 10,
                output_tokens: 20,
            },
            ..StageStats::default()
        };
        assert_eq!(
            summary(Stage::Answers, &stats),
            "answers: 3 done, 0 skipped, 0 failed, 0 excluded; tokens 10 in, 20 out; cost unknown"
        );
        stats.cost = Some(0.123_45);
        assert!(summary(Stage::Answers, &stats).ends_with("cost $0.1235"));
    }

    #[test]
    fn split_summary_lists_reasons() {
        let report = SplitReport {
            train: 9,
            eval: 1,
            excluded: BTreeMap::from([(Exclusion::Truncated, 2), (Exclusion::NoRawReasoning, 1)]),
            orphaned: 4,
        };
        assert_eq!(
            split_summary(&report),
            "split: 9 train, 1 eval, 3 excluded (truncated 2, no_raw_reasoning 1), 4 orphaned"
        );
    }
}