rswarm 0.1.8

A Rust implementation of the Swarm framework
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

#[cfg(test)]
mod tests {
    use crate::stream::Streamer;
    use crate::types::{Agent, ApiKey, ContextVariables, Instructions, Message};
    #[allow(unused)]
    use crate::SwarmError;
    use futures_util::{pin_mut, StreamExt};
    use reqwest::Client;
    use std::time::Duration;
    use tokio;
    use wiremock::matchers::{method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    // Helper function to create a simple test agent.
    fn test_agent() -> Agent {
        Agent::new(
            "test_agent",
            "gpt-4",
            Instructions::Text("You are a helpful assistant.".to_string()),
        )
        .expect("Failed to create test agent")
    }

    #[tokio::test]
    async fn test_stream_chat_returns_messages() {
        // Start a WireMock server.
        let mock_server = MockServer::start().await;

        // Create a response body that simulates streaming chunks.
        // The first chunk returns a partial response with a choice message,
        // followed by [DONE] to indicate the end of the stream.
        let body = "data: {\"id\":\"dummy\",\"object\":\"chat.completion\",\"created\":0,\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\"Hello from stream!\",\"name\":null,\"function_call\":null},\"finish_reason\":null}]}\n\
                    data: [DONE]\n";

        let response = ResponseTemplate::new(200).set_body_raw(body, "text/event-stream");

        // Set up the stub for POST /completions.
        Mock::given(method("POST"))
            .and(path("/completions"))
            .respond_with(response)
            .mount(&mock_server)
            .await;

        let api_url = format!("{}/completions", &mock_server.uri());

        // Create an HTTP client.
        let client = Client::builder()
            .timeout(Duration::from_secs(30))
            .build()
            .expect("Failed to build client");
        let api_key = ApiKey::new("sk-test123456789").expect("valid test key");

        let streamer = Streamer::new(client, api_key, api_url);
        let agent = test_agent();
        let history: Vec<Message> =
            vec![Message::user("Hello!").expect("Failed to create history message")];
        let context_variables = ContextVariables::new();

        let stream = streamer.stream_chat(&agent, &history, &context_variables, None, true);
        pin_mut!(stream);

        // Await one message from the stream.
        if let Some(result) = stream.next().await {
            match result {
                Ok(message) => {
                    // Verify that the message role and content are as expected.
                    assert_eq!(message.role().as_str(), "assistant");
                    assert_eq!(message.content(), Some("Hello from stream!"));
                }
                Err(e) => {
                    panic!("Stream returned an error: {:?}", e);
                }
            }
        } else {
            panic!("No messages returned from the stream");
        }
    }
}