kiromi-ai-cli 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! `kiromi-ai-memory context` (Plan 12 phase J).

use kiromi_ai_memory::{ContextOpts, MemoryRef, NodeRef, PartitionPath};

use crate::cli::{ContextArgs, GlobalArgs};
use crate::error::{CliError, ExitCode};
use crate::runtime::Runtime;

pub(crate) async fn run(args: ContextArgs, globals: &GlobalArgs) -> Result<(), CliError> {
    let rt = Runtime::open(globals).await?;
    let focus = parse_focus(&args.focus)?;
    let opts = ContextOpts::default()
        .with_budget(args.budget)
        .with_top_k(args.top_k)
        .with_include_tenant_memo(args.memo);
    let blocks = rt.mem.build_context(focus, opts).await?;
    if rt.json {
        let arr: Vec<serde_json::Value> = blocks
            .iter()
            .map(|b| {
                serde_json::json!({
                    "kind": format!("{:?}", b.kind),
                    "anchor": &b.anchor,
                    "text": b.text,
                    "tokens_estimated": b.tokens_estimated,
                })
            })
            .collect();
        println!("{}", serde_json::to_string_pretty(&arr).unwrap_or_default());
    } else {
        for b in &blocks {
            println!(
                "[{:?}] tokens={}\t{}",
                b.kind,
                b.tokens_estimated,
                b.text.lines().next().unwrap_or("")
            );
        }
    }
    rt.mem.close().await?;
    Ok(())
}

fn parse_focus(s: &str) -> Result<NodeRef, CliError> {
    if let Some(rest) = s.strip_prefix("memory:") {
        let mut parts = rest.splitn(2, ':');
        let id_s = parts.next().ok_or_else(|| CliError {
            kind: ExitCode::Config,
            source: anyhow::anyhow!("focus memory missing id"),
        })?;
        let path_s = parts.next().ok_or_else(|| CliError {
            kind: ExitCode::Config,
            source: anyhow::anyhow!("focus memory missing partition"),
        })?;
        let id = id_s
            .parse::<kiromi_ai_memory::MemoryId>()
            .map_err(|e| CliError {
                kind: ExitCode::Config,
                source: anyhow::anyhow!("focus memory id: {e}"),
            })?;
        let p: PartitionPath = path_s.parse().map_err(|e| CliError {
            kind: ExitCode::Config,
            source: anyhow::anyhow!("focus partition: {e}"),
        })?;
        return Ok(NodeRef::Memory(MemoryRef { id, partition: p }));
    }
    if let Some(rest) = s.strip_prefix("partition:") {
        let p: PartitionPath = rest.parse().map_err(|e| CliError {
            kind: ExitCode::Config,
            source: anyhow::anyhow!("focus partition: {e}"),
        })?;
        return Ok(NodeRef::Partition(p));
    }
    Err(CliError {
        kind: ExitCode::Config,
        source: anyhow::anyhow!("--focus must be `memory:<id>:<partition>` or `partition:<path>`"),
    })
}