aprender-orchestrate 0.31.2

Sovereign AI orchestration: autonomous agents, ML serving, code analysis, and transpilation pipelines
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

//! Batch inference tests.

use crate::serve::templates::ChatMessage;

use super::batch::{BatchItem, BatchItemResult, BatchStatus};

// ============================================================================
// BatchStore unit tests
// ============================================================================

#[test]
#[allow(non_snake_case)]
fn test_BATCH_001_run_batch() {
    let store = super::batch::BatchStore::new();
    let items = vec![
        BatchItem {
            id: "req-1".to_string(),
            messages: vec![ChatMessage::user("Hello")],
            max_tokens: 10,
        },
        BatchItem {
            id: "req-2".to_string(),
            messages: vec![ChatMessage::user("World")],
            max_tokens: 10,
        },
    ];
    let job = store.run(items, |item| BatchItemResult {
        id: item.id.clone(),
        content: format!("echo: {}", item.messages[0].content),
        finish_reason: "test".to_string(),
        tokens: 1,
    });
    assert!(job.batch_id.starts_with("batch-"));
    assert_eq!(job.status, BatchStatus::Complete);
    assert_eq!(job.total_items, 2);
    assert_eq!(job.results.len(), 2);
    assert_eq!(job.results[0].content, "echo: Hello");
}

#[test]
#[allow(non_snake_case)]
fn test_BATCH_002_get_job() {
    let store = super::batch::BatchStore::new();
    let job = store.run(vec![], |_| unreachable!());
    let retrieved = store.get(&job.batch_id).expect("should exist");
    assert_eq!(retrieved.batch_id, job.batch_id);
}

#[test]
#[allow(non_snake_case)]
fn test_BATCH_003_list_jobs() {
    let store = super::batch::BatchStore::new();
    store.run(vec![], |_| unreachable!());
    store.run(vec![], |_| unreachable!());
    assert_eq!(store.list().len(), 2);
}

// ============================================================================
// Batch endpoint tests
// ============================================================================

#[tokio::test]
#[allow(non_snake_case)]
async fn test_BATCH_HDL_001_submit() {
    use axum::{body::Body, http::Request};
    use tower::ServiceExt;

    let app = super::router::create_banco_router(super::state::BancoStateInner::with_defaults());
    let body = serde_json::json!({
        "items": [
            {"id": "q1", "messages": [{"role": "user", "content": "What is Rust?"}]},
            {"id": "q2", "messages": [{"role": "user", "content": "What is Python?"}]}
        ]
    });
    let response = app
        .oneshot(
            Request::post("/api/v1/batch")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&body).expect("json")))
                .expect("req"),
        )
        .await
        .expect("resp");
    assert_eq!(response.status(), axum::http::StatusCode::OK);
    let bytes = axum::body::to_bytes(response.into_body(), 1_048_576).await.expect("body");
    let json: serde_json::Value = serde_json::from_slice(&bytes).expect("parse");
    assert_eq!(json["status"], "complete");
    assert_eq!(json["total_items"], 2);
    assert_eq!(json["results"].as_array().expect("results").len(), 2);
}

#[tokio::test]
#[allow(non_snake_case)]
async fn test_BATCH_HDL_002_list() {
    use axum::{body::Body, http::Request};
    use tower::ServiceExt;

    let app = super::router::create_banco_router(super::state::BancoStateInner::with_defaults());
    let response = app
        .oneshot(Request::get("/api/v1/batch").body(Body::empty()).expect("req"))
        .await
        .expect("resp");
    assert_eq!(response.status(), axum::http::StatusCode::OK);
}