unigateway 1.3.0

Lightweight, local-first LLM gateway for developers. A stable, single-binary unified entry point for all your AI tools and models.
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

use std::sync::Arc;
use std::time::Instant;

use axum::{
    Json,
    body::Body,
    http::header,
    response::{IntoResponse, Response},
};
use unigateway_host::flow::HostResponseResult;
use unigateway_protocol::{RuntimeHttpResponse, RuntimeResponseBody};

use crate::middleware::{GatewayAuth, record_stat};
use crate::types::GatewayRequestState;

use super::request_flow::PreparedGatewayRequest;

pub(super) async fn respond_prepared_host_result(
    state: &Arc<GatewayRequestState>,
    prepared: &PreparedGatewayRequest<'_>,
    endpoint: &str,
    result: HostResponseResult,
) -> Response {
    match prepared.auth.as_ref() {
        Some(auth) => {
            respond_authenticated_host_result(state, auth, endpoint, &prepared.start, result).await
        }
        None => respond_env_host_result(state, endpoint, &prepared.start, result).await,
    }
}

pub(super) async fn respond_authenticated_host_result(
    state: &Arc<GatewayRequestState>,
    auth: &GatewayAuth,
    endpoint: &str,
    start: &Instant,
    result: HostResponseResult,
) -> Response {
    match result {
        Ok(response) => gateway_success_response(state, auth, endpoint, start, response).await,
        Err(response) => gateway_error_response(state, auth, endpoint, start, response).await,
    }
}

pub(super) async fn respond_env_host_result(
    state: &Arc<GatewayRequestState>,
    endpoint: &str,
    start: &Instant,
    result: HostResponseResult,
) -> Response {
    match result {
        Ok(response) => success_response(state, endpoint, start, response).await,
        Err(response) => error_response(state, endpoint, start, response).await,
    }
}

async fn gateway_success_response(
    state: &Arc<GatewayRequestState>,
    auth: &GatewayAuth,
    endpoint: &str,
    start: &Instant,
    response: RuntimeHttpResponse,
) -> Response {
    auth.finalize(state).await;
    success_response(state, endpoint, start, response).await
}

async fn gateway_error_response(
    state: &Arc<GatewayRequestState>,
    auth: &GatewayAuth,
    endpoint: &str,
    start: &Instant,
    response: RuntimeHttpResponse,
) -> Response {
    auth.release(state).await;
    error_response(state, endpoint, start, response).await
}

async fn success_response(
    state: &Arc<GatewayRequestState>,
    endpoint: &str,
    start: &Instant,
    response: RuntimeHttpResponse,
) -> Response {
    record_stat(state, endpoint, 200, start).await;
    into_axum_response(response)
}

async fn error_response(
    state: &Arc<GatewayRequestState>,
    endpoint: &str,
    start: &Instant,
    response: RuntimeHttpResponse,
) -> Response {
    record_stat(state, endpoint, 500, start).await;
    into_axum_response(response)
}

fn into_axum_response(response: RuntimeHttpResponse) -> Response {
    let (status, body) = response.into_parts();

    match body {
        RuntimeResponseBody::Json(value) => (status, Json(value)).into_response(),
        RuntimeResponseBody::ServerSentEvents(stream) => (
            status,
            [(header::CONTENT_TYPE, "text/event-stream")],
            Body::from_stream(stream),
        )
            .into_response(),
    }
}