AMRS/src/main.rs at 90bdea3084776616a3f6a33d8d9da877b60ac672 · InftyAI/AMRS · GitHub

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
use tokio::runtime::Runtime;

use arms::client;
use arms::types::chat;

fn main() {
    // case 1: completion with DeepInfra provider.
    let config = client::Config::builder()
        .provider("deepinfra")
        .routing_mode(client::RoutingMode::WRR)
        .model(
            client::ModelConfig::builder()
                .name("nvidia/Nemotron-3-Nano-30B-A3B")
                .weight(1)
                .build()
                .unwrap(),
        )
        .model(
            client::ModelConfig::builder()
                .name("deepseek-ai/DeepSeek-V3.2")
                .weight(2)
                .build()
                .unwrap(),
        )
        .build()
        .unwrap();

    let mut client = client::Client::new(config);

    let request = chat::CreateChatCompletionRequestArgs::default()
        .messages([
            chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(),
            chat::ChatCompletionRequestUserMessage::from("Who won the world series in 2020?")
                .into(),
            chat::ChatCompletionRequestAssistantMessage::from(
                "The Los Angeles Dodgers won the World Series in 2020.",
            )
            .into(),
            chat::ChatCompletionRequestUserMessage::from("Where was it played?").into(),
        ])
        .build()
        .unwrap();

    let result = Runtime::new()
        .unwrap()
        .block_on(client.create_completion(request));

    match result {
        Ok(response) => {
            for choice in response.choices {
                println!("Response: {:?}", choice.message.content);
            }
        }
        Err(e) => {
            eprintln!("Error: {}", e);
        }
    }

    // case 2: response with DeepInfra provider.
    // let config = client::Config::builder()
    //     .provider("deepinfra")
    //     .routing_mode(client::RoutingMode::WRR)
    //     .model(
    //         client::ModelConfig::builder()
    //             .name("nvidia/Nemotron-3-Nano-30B-A3B")
    //             .weight(1)
    //             .build()
    //             .unwrap(),
    //     )
    //     .model(
    //         client::ModelConfig::builder()
    //             .name("deepseek-ai/DeepSeek-V3.2")
    //             .weight(2)
    //             .build()
    //             .unwrap(),
    //     )
    //     .build()
    //     .unwrap();

    // let mut client = client::Client::new(config);

    // let request = responses::CreateResponseArgs::default()
    //     .input(responses::InputParam::Items(vec![
    //         responses::InputItem::EasyMessage(responses::EasyInputMessage {
    //             r#type: responses::MessageType::Message,
    //             role: responses::Role::User,
    //             content: responses::EasyInputContent::Text("What is AGI?".to_string()),
    //         }),
    //     ]))
    //     .build()
    //     .unwrap();

    // let result = Runtime::new()
    //     .unwrap()
    //     .block_on(client.create_response(request));

    // match result {
    //     Ok(response) => {
    //         println!("Response ID: {}", response.id);
    //     }
    //     Err(e) => {
    //         eprintln!("Error: {}", e);
    //     }
    // }
}