llm-shield-api 0.1.0

Production-grade REST API for LLM Shield
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
121
122
123
124

//! Route configuration

use axum::{routing::{get, post}, Router};

use crate::handlers;
use crate::state::AppState;

/// Create the application router
///
/// ## Routes
/// - GET /health - Basic health check
/// - GET /health/ready - Readiness probe
/// - GET /health/live - Liveness probe
/// - GET /version - Version information
/// - POST /v1/scan/prompt - Scan user prompt
pub fn create_router() -> Router {
    Router::new()
        .route("/health", get(handlers::health))
        .route("/health/ready", get(handlers::ready))
        .route("/health/live", get(handlers::live))
        .route("/version", get(handlers::version))
}

/// Create the application router with state
pub fn create_router_with_state(state: AppState) -> Router {
    Router::new()
        .route("/health", get(handlers::health))
        .route("/health/ready", get(handlers::ready))
        .route("/health/live", get(handlers::live))
        .route("/version", get(handlers::version))
        .route("/v1/scan/prompt", post(handlers::scan_prompt))
        .route("/v1/scan/output", post(handlers::scan_output))
        .route("/v1/scan/batch", post(handlers::scan_batch))
        .route("/v1/scanners", get(handlers::list_scanners))
        .with_state(state)
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use tower::ServiceExt; // For `oneshot`

    #[tokio::test]
    async fn test_health_route() {
        let app = create_router();

        let response = app
            .oneshot(Request::builder().uri("/health").body(Body::empty()).unwrap())
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_ready_route() {
        let app = create_router();

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/health/ready")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_live_route() {
        let app = create_router();

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/health/live")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_version_route() {
        let app = create_router();

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/version")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_not_found() {
        let app = create_router();

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/notfound")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::NOT_FOUND);
    }
}