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
78
79
80
81
82
83
84
85
86

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

use rig::agent::AgentBuilder;
use rig::client::CompletionClient;
use rig::completion::Prompt;

use super::{
    BEDROCK_COMPLETION_MODEL, client,
    support::{
        Adder, BASIC_PREAMBLE, BASIC_PROMPT, CONTEXT_DOCS, CONTEXT_PROMPT,
        STREAMING_TOOLS_PREAMBLE, STREAMING_TOOLS_PROMPT, Subtract,
        assert_contains_any_case_insensitive, assert_mentions_expected_number,
        assert_nonempty_response,
    },
};

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

    let response = agent
        .prompt(BASIC_PROMPT)
        .await
        .expect("completion should succeed");

    assert_nonempty_response(&response);
}

#[tokio::test]
#[ignore = "requires AWS credentials and Bedrock model access"]
async fn completion_with_context_smoke() {
    let agent = client()
        .agent(BEDROCK_COMPLETION_MODEL)
        .preamble("Answer the user using only the supplied context.")
        .context(CONTEXT_DOCS[0])
        .context(CONTEXT_DOCS[1])
        .context(CONTEXT_DOCS[2])
        .build();

    let response = agent
        .prompt(CONTEXT_PROMPT)
        .await
        .expect("context completion should succeed");

    assert_contains_any_case_insensitive(&response, &["ancient tool", "farm"]);
}

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

    let response = agent
        .prompt(STREAMING_TOOLS_PROMPT)
        .await
        .expect("tool prompt should succeed");

    assert_mentions_expected_number(&response, -3);
}

#[tokio::test]
#[ignore = "requires AWS credentials and Bedrock model access"]
async fn prompt_caching_completion_smoke() {
    let bedrock_client = client();
    let model = bedrock_client
        .completion_model(BEDROCK_COMPLETION_MODEL)
        .with_prompt_caching();
    let agent = AgentBuilder::new(model).preamble(BASIC_PREAMBLE).build();

    let response = agent
        .prompt(BASIC_PROMPT)
        .await
        .expect("prompt-caching completion should succeed");

    assert_nonempty_response(&response);
}