malvin 0.2.5

Non-interactive research and coding agent
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

use pi::sdk::{ContentBlock, Message};
use serde_json::{Value, json};

use crate::bridge_protocol::{BridgeEvent, RunDoneStatus};

use super::usage_cost::aggregate_cost_usd;

#[must_use]
pub(crate) fn map_agent_end(messages: &[Message], error: Option<&str>) -> BridgeEvent {
    let status = if error.is_some() {
        RunDoneStatus::Error
    } else {
        RunDoneStatus::Finished
    };
    BridgeEvent::RunDone {
        status,
        result: last_assistant_text(messages),
        usage: aggregate_usage(messages),
        error: error.map(str::to_string),
        duration_ms: None,
    }
}

fn last_assistant_text(messages: &[Message]) -> Option<String> {
    messages.iter().rev().find_map(|msg| match msg {
        Message::Assistant(assistant) => text_from_blocks(&assistant.content),
        _ => None,
    })
}

fn text_from_blocks(blocks: &[ContentBlock]) -> Option<String> {
    let mut text = String::new();
    for block in blocks {
        if let ContentBlock::Text(part) = block {
            text.push_str(&part.text);
        }
    }
    (!text.is_empty()).then_some(text)
}

fn aggregate_usage(messages: &[Message]) -> Option<Value> {
    let mut input = 0_u64;
    let mut output = 0_u64;
    let mut cache_read = 0_u64;
    let mut cache_write = 0_u64;
    let mut seen = false;
    for msg in messages {
        let Message::Assistant(assistant) = msg else {
            continue;
        };
        let assistant = assistant.as_ref();
        seen = true;
        input = input.saturating_add(assistant.usage.input);
        output = output.saturating_add(assistant.usage.output);
        cache_read = cache_read.saturating_add(assistant.usage.cache_read);
        cache_write = cache_write.saturating_add(assistant.usage.cache_write);
    }
    if !seen {
        return None;
    }
    let mut usage = serde_json::Map::from_iter([
        ("inputTokens".into(), json!(input)),
        ("outputTokens".into(), json!(output)),
        ("cacheReadTokens".into(), json!(cache_read)),
        ("cacheWriteTokens".into(), json!(cache_write)),
    ]);
    let cost = aggregate_cost_usd(messages);
    if cost.is_present() {
        usage.insert(
            "costUsd".into(),
            json!({
                "input": cost.input,
                "output": cost.output,
                "cacheRead": cost.cache_read,
                "cacheWrite": cost.cache_write,
                "total": cost.total,
            }),
        );
    }
    Some(Value::Object(usage))
}