Skip to main content

pai_engine/
main.rs

1mod bootstrap;
2
3use anyhow::Result;
4use bootstrap::{load_config, log_domain_stack_init, wait_for_shutdown_signal};
5use clap::{ArgAction, Parser};
6use pai_core::adapters::HardcodedFlowRunner;
7use pai_core::domain::{EventBus, SessionManager};
8use pai_core::ports::{InferenceError, InferencePort, InferenceRequest};
9use pai_core::types::Token;
10use std::path::PathBuf;
11use std::sync::Arc;
12use tokio::sync::mpsc;
13use tracing::{debug, info};
14
15/// Command line arguments for the paiOS engine.
16#[derive(Parser, Debug)]
17#[command(author, version, about, long_about = None)]
18struct Args {
19    /// Path to the configuration file
20    #[arg(short, long)]
21    config: Option<PathBuf>,
22
23    /// Increase verbosity (-v, -vv, -vvv)
24    #[arg(short, long, action = ArgAction::Count)]
25    verbose: u8,
26}
27
28#[derive(Debug)]
29struct StubInference;
30
31impl InferencePort for StubInference {
32    async fn infer(
33        &self,
34        req: InferenceRequest,
35    ) -> Result<mpsc::Receiver<Result<Token, InferenceError>>, InferenceError> {
36        let (tx, rx) = mpsc::channel(1);
37        let content = format!("[stub] {}", req.prompt);
38        // Spawn a task to honour the cancellation token and send the stub token.
39        tokio::spawn(async move {
40            tokio::select! {
41                _ = req.cancellation.cancelled() => {}
42                res = tx.send(Ok(Token { content })) => {
43                    let _ = res;
44                }
45            }
46        });
47        Ok(rx)
48    }
49}
50
51#[tokio::main]
52async fn main() -> Result<()> {
53    let args = Args::parse();
54
55    // `RUST_LOG` overrides the verbosity flag; -v / -vv / -vvv set the fallback.
56    let default_level = match args.verbose {
57        0 => "info",
58        1 => "debug",
59        _ => "trace",
60    };
61
62    common::logging::try_init(default_level)
63        .map_err(|e| anyhow::anyhow!("tracing subscriber init failed: {e}"))?;
64
65    load_config(args.config.as_deref())?;
66
67    info!(target: "pai_engine::bootstrap", "pai-engine composition root starting");
68
69    log_domain_stack_init();
70
71    let (event_bus, event_rx) = EventBus::channel(64);
72    let flow_runner = Arc::new(HardcodedFlowRunner::new(
73        Arc::new(StubInference),
74        event_bus.clone(),
75    ));
76    let session = SessionManager::new(flow_runner, event_bus);
77
78    // Keep the sole consumer alive so the mpsc channel stays open; drain events so publishes never
79    // fail with Closed/Full during normal operation.
80    tokio::spawn(async move {
81        let mut event_rx = event_rx;
82        while let Some(ev) = event_rx.recv().await {
83            debug!(target: "pai_engine::event_bus", ?ev, "domain event");
84        }
85    });
86    info!(
87        target: "pai_engine::bootstrap",
88        "session orchestration ready (state: {:?})",
89        session.state_machine().state()
90    );
91
92    info!(
93        target: "pai_engine::bootstrap",
94        "engine main loop running (waiting for shutdown signal)"
95    );
96
97    wait_for_shutdown_signal().await;
98
99    info!(
100        target: "pai_engine::bootstrap",
101        "shutdown complete; exiting pai-engine"
102    );
103
104    Ok(())
105}