dbx-tools-model-proxy 0.6.213

Multi-protocol Databricks model proxy
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

//! Shared request-completion logging for buffered and streaming routes.

use std::{net::SocketAddr, time::Instant};

use axum::http::StatusCode;

use crate::{
    protocol::{ClientWire, TargetWire},
    throttle::{ResponseTokenUsage, ThrottleAcquisition},
};

macro_rules! log_request {
    ($context:expr, $($fields:tt)*) => {
        tracing::info!(
            requested_model = %$context.requested_model,
            resolved_model = %$context.resolved_model,
            client_ip = %$context.peer.ip(),
            client_port = $context.peer.port(),
            request_bytes = $context.request_bytes,
            raw_estimated_input_tokens = $context.throttle.raw_estimated_input_tokens,
            estimate_factor = $context.throttle.estimate_factor,
            estimated_input_tokens = $context.throttle.estimated_input_tokens,
            reserved_output_tokens = $context.throttle.reserved_output_tokens,
            estimated_tokens = $context.throttle.estimated_tokens,
            upstream_attempt = $context.upstream_attempt,
            token_throttle_mode = ?$context.throttle.mode,
            token_throttle_active = $context.throttle.active,
            token_limit_input = $context.throttle.input_limit,
            token_reservation_input = $context.throttle.reserved_input_tokens,
            token_window_used_before = $context.throttle.input_window_used_before,
            token_window_wait_ms = $context.throttle.wait.as_millis(),
            oversized_request = false,
            $($fields)*
        );
    };
}

/// Request metadata shared by buffered and streamed completion events.
#[derive(Debug)]
pub(crate) struct RequestLogContext {
    requested_model: String,
    resolved_model: String,
    peer: SocketAddr,
    request_bytes: usize,
    started: Instant,
    throttle: ThrottleAcquisition,
    upstream_attempt: u32,
}

impl RequestLogContext {
    /// Build a context after the upstream attempt has acquired its token reservation.
    pub(crate) fn new(
        requested_model: String,
        resolved_model: String,
        peer: SocketAddr,
        request_bytes: usize,
        started: Instant,
        throttle: ThrottleAcquisition,
        upstream_attempt: u32,
    ) -> Self {
        Self {
            requested_model,
            resolved_model,
            peer,
            request_bytes,
            started,
            throttle,
            upstream_attempt,
        }
    }

    /// Reconcile the local token reservation with reported upstream usage.
    pub(crate) async fn reconcile(&self, usage: ResponseTokenUsage) {
        self.throttle.reconcile(usage).await;
    }

    /// Reconcile and log one buffered embeddings request.
    pub(crate) async fn complete_embedding(&self, status: StatusCode, usage: ResponseTokenUsage) {
        self.reconcile(usage).await;
        log_request!(
            self,
            route = "/v1/embeddings",
            status = status.as_u16(),
            input_tokens = usage.input,
            output_tokens = usage.output,
            total_tokens = usage.total,
            latency_ms = self.started.elapsed().as_millis(),
            "embedding request completed"
        );
    }

    /// Log the point at which a successful upstream stream is connected.
    pub(crate) fn stream_connected(
        &self,
        client_wire: ClientWire,
        target: TargetWire,
        status: StatusCode,
    ) {
        log_request!(
            self,
            ?client_wire,
            ?target,
            streaming = true,
            status = status.as_u16(),
            latency_ms = self.started.elapsed().as_millis(),
            "model stream connected"
        );
    }

    /// Reconcile and log one buffered model request.
    pub(crate) async fn complete_model(
        &self,
        client_wire: ClientWire,
        target: TargetWire,
        streaming: bool,
        status: StatusCode,
        usage: ResponseTokenUsage,
    ) {
        self.reconcile(usage).await;
        log_request!(
            self,
            ?client_wire,
            ?target,
            streaming,
            status = status.as_u16(),
            input_tokens = usage.input,
            output_tokens = usage.output,
            total_tokens = usage.total,
            latency_ms = self.started.elapsed().as_millis(),
            "model request completed"
        );
    }

    /// Log final body completion or cancellation for a streamed request.
    pub(crate) fn stream_completed(
        &self,
        client_wire: ClientWire,
        target: TargetWire,
        response_bytes: u64,
        usage: ResponseTokenUsage,
        finished: bool,
        failed: bool,
    ) {
        log_request!(
            self,
            ?client_wire,
            ?target,
            streaming = true,
            response_bytes,
            input_tokens = usage.input,
            output_tokens = usage.output,
            total_tokens = usage.total,
            duration_ms = self.started.elapsed().as_millis(),
            finished,
            failed,
            "model stream completed"
        );
    }
}