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

//! Health check endpoints

use axum::{http::StatusCode, response::IntoResponse, Json};
use serde::{Deserialize, Serialize};

/// Health check response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthResponse {
    /// Service status
    pub status: String,

    /// Service version
    pub version: String,

    /// Uptime in seconds
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uptime_secs: Option<u64>,
}

/// Basic health check endpoint
///
/// Returns 200 OK if service is running
///
/// ## TDD Design
/// - Tests written first in tests module
/// - Simple, fast endpoint (< 1ms)
/// - No dependencies on external services
pub async fn health() -> impl IntoResponse {
    let response = HealthResponse {
        status: "ok".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        uptime_secs: None,
    };

    (StatusCode::OK, Json(response))
}

/// Readiness check endpoint
///
/// Returns 200 OK if service is ready to accept requests
/// Checks:
/// - Models loaded (if required)
/// - Cache accessible
///
/// ## Kubernetes Integration
/// Used by readiness probe to determine if pod should receive traffic
pub async fn ready() -> impl IntoResponse {
    // TODO: Check if models are loaded
    // TODO: Check if cache is accessible

    let response = HealthResponse {
        status: "ready".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        uptime_secs: None,
    };

    (StatusCode::OK, Json(response))
}

/// Liveness check endpoint
///
/// Returns 200 OK if service is alive and responding
///
/// ## Kubernetes Integration
/// Used by liveness probe to determine if pod should be restarted
pub async fn live() -> impl IntoResponse {
    let response = HealthResponse {
        status: "live".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        uptime_secs: None,
    };

    (StatusCode::OK, Json(response))
}

/// Version information endpoint
pub async fn version() -> impl IntoResponse {
    #[derive(Serialize)]
    struct VersionInfo {
        version: String,
        build_time: String,
        git_commit: String,
    }

    let info = VersionInfo {
        version: env!("CARGO_PKG_VERSION").to_string(),
        build_time: "unknown".to_string(), // TODO: Add build time
        git_commit: "unknown".to_string(), // TODO: Add git commit
    };

    (StatusCode::OK, Json(info))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_health_endpoint() {
        let response = health().await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_ready_endpoint() {
        let response = ready().await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_live_endpoint() {
        let response = live().await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_version_endpoint() {
        let response = version().await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
    }
}