rig 0.41.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

//! Focused DeepSeek cassette coverage for request document ordering.

use rig::OneOrMany;
use rig::completion::{AssistantContent, CompletionModel, Document, Message};
use rig::prelude::*;
use rig::providers::deepseek;
use serde::Deserialize;
use serde_json::Value;

use super::support::with_deepseek_cassette;
use crate::support::assert_contains_any_case_insensitive;

const SYSTEM_INSTRUCTION: &str = "Answer with the exact token from the document only.";
const DOCUMENT_ANSWER: &str = "violet-needle";
const PROMPT: &str = "According to the document, what is the ordering token?";

#[derive(Deserialize)]
struct RecordedInteraction {
    when: RecordedRequest,
}

#[derive(Deserialize)]
struct RecordedRequest {
    body: Option<String>,
}

fn ordering_document() -> Document {
    Document {
        id: "ordering-note".to_string(),
        text: format!("The ordering token is {DOCUMENT_ANSWER}."),
        additional_props: Default::default(),
    }
}

fn assistant_text(choice: &OneOrMany<AssistantContent>) -> String {
    choice
        .iter()
        .filter_map(|content| match content {
            AssistantContent::Text(text) => Some(text.text.clone()),
            _ => None,
        })
        .collect::<Vec<_>>()
        .join("")
}

#[tokio::test]
async fn chat_completions_keeps_documents_after_system_before_history() {
    with_deepseek_cassette(
        "document_ordering/chat_completions_keeps_documents_after_system_before_history",
        |client| async move {
            let response = client
                .completion_model(deepseek::DEEPSEEK_V4_FLASH)
                .completion_request(PROMPT)
                .message(Message::system(SYSTEM_INSTRUCTION))
                .message(Message::assistant("Acknowledged."))
                .document(ordering_document())
                .temperature(0.0)
                // Needs headroom for deepseek-v4-flash's thinking tokens now
                // that max_tokens is actually forwarded to the API.
                .max_tokens(512)
                .send()
                .await
                .expect("DeepSeek document ordering request should succeed");

            assert_contains_any_case_insensitive(
                &assistant_text(&response.choice),
                &[DOCUMENT_ANSWER],
            );
        },
    )
    .await;

    assert_chat_request_order(
        "document_ordering/chat_completions_keeps_documents_after_system_before_history",
    );
}

fn recorded_request_body(scenario: &str) -> Value {
    let cassette_path = crate::cassettes::cassette_path("deepseek", scenario);
    let contents = std::fs::read_to_string(&cassette_path).unwrap_or_else(|error| {
        panic!(
            "provider cassette {} should be readable after recording: {error}",
            cassette_path.display()
        )
    });

    serde_yaml::Deserializer::from_str(&contents)
        .find_map(|document| {
            let interaction = RecordedInteraction::deserialize(document)
                .expect("cassette interaction should deserialize");
            interaction
                .when
                .body
                .and_then(|body| serde_json::from_str::<Value>(&body).ok())
        })
        .unwrap_or_else(|| panic!("expected cassette {scenario} to contain a JSON request body"))
}

fn assert_chat_request_order(scenario: &str) {
    let body = recorded_request_body(scenario);
    let messages = body["messages"]
        .as_array()
        .expect("DeepSeek request should contain messages[]");
    assert_eq!(
        messages.len(),
        4,
        "expected system, document, assistant history, and prompt messages: {body:#}"
    );
    assert_eq!(messages[0]["role"], "system");
    assert!(messages[0].to_string().contains(SYSTEM_INSTRUCTION));
    assert_eq!(messages[1]["role"], "user");
    assert!(
        messages[1].to_string().contains("<file id: ordering-note>"),
        "expected second message to contain normalized document: {body:#}"
    );
    assert_eq!(messages[2]["role"], "assistant");
    assert!(messages[2].to_string().contains("Acknowledged."));
    assert_eq!(messages[3]["role"], "user");
    assert!(messages[3].to_string().contains(PROMPT));
}