rig 0.37.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

//! AWS Bedrock streaming smoke tests inspired by the OpenAI and Anthropic provider tests.

use rig::client::CompletionClient;
use rig::completion::CompletionModel as _;
use rig::message::ToolChoice;
use rig::streaming::StreamingPrompt;
use rig::tool::Tool;

use super::{
    BEDROCK_COMPLETION_MODEL, client,
    support::{
        Adder, AlphaSignal, ORDERED_TOOL_STREAM_PREAMBLE, ORDERED_TOOL_STREAM_PROMPT,
        STREAMING_PREAMBLE, STREAMING_PROMPT, STREAMING_TOOLS_PREAMBLE, STREAMING_TOOLS_PROMPT,
        Subtract, assert_mentions_expected_number, assert_nonempty_response,
        assert_raw_stream_tool_call_precedes_text, collect_raw_stream_observation,
        collect_stream_final_response,
    },
};

#[tokio::test]
#[ignore = "requires AWS credentials and Bedrock model access"]
async fn streaming_smoke() {
    let agent = client()
        .agent(BEDROCK_COMPLETION_MODEL)
        .preamble(STREAMING_PREAMBLE)
        .build();

    let mut stream = agent.stream_prompt(STREAMING_PROMPT).await;
    let response = collect_stream_final_response(&mut stream)
        .await
        .expect("streaming prompt should succeed");

    assert_nonempty_response(&response);
}

#[tokio::test]
#[ignore = "requires AWS credentials and Bedrock model access"]
async fn streaming_tools_smoke() {
    let agent = client()
        .agent(BEDROCK_COMPLETION_MODEL)
        .preamble(STREAMING_TOOLS_PREAMBLE)
        .max_tokens(1024)
        .tool(Adder)
        .tool(Subtract)
        .build();

    let mut stream = agent.stream_prompt(STREAMING_TOOLS_PROMPT).await;
    let response = collect_stream_final_response(&mut stream)
        .await
        .expect("streaming tool prompt should succeed");

    assert_mentions_expected_number(&response, -3);
}

#[tokio::test]
#[ignore = "requires AWS credentials and Bedrock model access"]
async fn raw_streaming_tool_call_smoke() {
    let model = client().completion_model(BEDROCK_COMPLETION_MODEL);
    let request = model
        .completion_request(ORDERED_TOOL_STREAM_PROMPT)
        .preamble(ORDERED_TOOL_STREAM_PREAMBLE.to_string())
        .tool(AlphaSignal.definition(String::new()).await)
        .tool_choice(ToolChoice::Specific {
            function_names: vec![AlphaSignal::NAME.to_string()],
        })
        .build();

    let observation = collect_raw_stream_observation(
        model
            .stream(request)
            .await
            .expect("raw Bedrock stream should start"),
    )
    .await;

    assert_raw_stream_tool_call_precedes_text(&observation, "lookup_harbor_label");
}