rig-core 0.35.0

An opinionated library for building LLM powered applications.
Documentation 
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

use futures::stream;
use rig::OneOrMany;
use rig::agent::MultiTurnStreamItem;
use rig::completion::Usage;
use rig::message::{ToolCall, ToolFunction, ToolResult, ToolResultContent};
use rig::streaming::{StreamedAssistantContent, StreamedUserContent};

use crate::reasoning::collect_stream_stats;

#[tokio::test]
async fn collect_stream_stats_tracks_only_final_turn_text() {
    let internal_call_id = "call_internal_1".to_string();
    let tool_call = ToolCall::new(
        "tool_1".to_string(),
        ToolFunction::new(
            "get_weather".to_string(),
            serde_json::json!({ "city": "Tokyo" }),
        ),
    );
    let tool_result = ToolResult {
        id: "tool_1".to_string(),
        call_id: None,
        content: OneOrMany::one(ToolResultContent::text("72F and sunny")),
    };

    let items = vec![
        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::<()>::text("Sure! Let me check the weather right away!"),
        )),
        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::ToolCall {
                tool_call,
                internal_call_id: internal_call_id.clone(),
            },
        )),
        Ok(MultiTurnStreamItem::StreamUserItem(
            StreamedUserContent::tool_result(tool_result, internal_call_id),
        )),
        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::<()>::text("It's 72F and sunny in Tokyo."),
        )),
        Ok(MultiTurnStreamItem::final_response(
            "It's 72F and sunny in Tokyo.",
            Usage::new(),
        )),
    ];

    let stats = collect_stream_stats(stream::iter(items), "test").await;

    assert_eq!(stats.tool_calls_in_stream, vec!["get_weather".to_string()]);
    assert_eq!(stats.tool_results_in_stream, 1);
    assert!(stats.got_final_response, "expected final response event");
    assert_eq!(
        stats.final_turn_text, "It's 72F and sunny in Tokyo.",
        "pre-tool assistant text should not be counted as final-turn text"
    );
    assert_eq!(
        stats.final_response_text.as_deref(),
        Some(stats.final_turn_text.as_str()),
        "final response text should match the final turn's streamed text"
    );
}