feat(lprovider): implement LLM Provider Gateway crate (admin/ai proto+grpc, affinity, audit, billing, quota, rotation, selector, upstream, wallet, http, main, tests)
zhenyi committed
Added agent/lprovider/tests/affinity.rs agent/lprovider/tests/affinity.rs +78 −0 View agent/lprovider/tests/affinity.rs
@@ -0,0 +1,78 @@
1use std::sync::Arc;
2
3use config::lprovider::LproviderRedisSettings;
4use llm::provider::LlmProvider;
5use llm::provider::openai::{OpenAIConfig, OpenAIProvider};
6use llm::types::ChatMessage;
7use lprovider::affinity::AffinityManager;
8
9fn estimator() -> Arc<dyn LlmProvider> {
10 let cfg = OpenAIConfig::new("test-key").with_model("test-model");
11 Arc::new(OpenAIProvider::new(cfg).expect("provider builds"))
12}
13
14#[test]
15fn session_key_override_validation() {
16 let mgr = AffinityManager::new(None, &LproviderRedisSettings::default());
17 assert!(mgr.validate_session_key_override("abc-123.X_y").is_ok());
18 assert!(mgr.validate_session_key_override("").is_err());
19 assert!(mgr.validate_session_key_override(&"a".repeat(129)).is_err());
20 assert!(
21 mgr.validate_session_key_override("bad key with space")
22 .is_err()
23 );
24 assert!(mgr.validate_session_key_override("semi;colon").is_err());
25 assert!(mgr.validate_session_key_override("quote\"x").is_err());
26}
27
28#[test]
29fn short_prefix_does_not_bind() {
30 let mgr = AffinityManager::new(None, &LproviderRedisSettings::default());
31 let messages = vec![ChatMessage::user("hi there")];
32 let key = mgr.compute_session_key("auto", &messages, &estimator());
33 assert!(key.is_none(), "short prefix must not produce a session key");
34}
35
36#[test]
37fn long_prefix_produces_key() {
38 let settings = LproviderRedisSettings {
39 min_prefix_tokens: 8,
40 prefix_token_budget: 512,
41 ..LproviderRedisSettings::default()
42 };
43 let mgr = AffinityManager::new(None, &settings);
44 let long_text = "lorem ipsum dolor sit amet ".repeat(10);
45 let messages = vec![
46 ChatMessage::system("you are helpful"),
47 ChatMessage::user(long_text),
48 ];
49 let (key, tokens) = mgr
50 .compute_session_key("auto", &messages, &estimator())
51 .expect("long prefix must produce a key");
52 assert!(!key.is_empty());
53 assert!(tokens >= 8);
54 assert_eq!(key.len(), 64, "session key must be sha256 hex");
55}
56
57#[test]
58fn distinct_prefixes_distinct_keys() {
59 let settings = LproviderRedisSettings {
60 min_prefix_tokens: 8,
61 ..LproviderRedisSettings::default()
62 };
63 let mgr = AffinityManager::new(None, &settings);
64 let mk = |text: &str| {
65 mgr.compute_session_key("auto", &[ChatMessage::user(text.to_string())], &estimator())
66 .map(|(k, _)| k)
67 };
68 let a = mk("the quick brown fox jumps over the lazy dog");
69 let b = mk("a completely different opening statement here");
70 assert_ne!(a, b);
71}
72
73#[test]
74fn bind_key_format() {
75 let mgr = AffinityManager::new(None, &LproviderRedisSettings::default());
76 let key = mgr.bind_key(42, "auto", "sesskey");
77 assert_eq!(key, "lprovider:bind:42:auto:sesskey");
78}
Added agent/lprovider/tests/billing.rs agent/lprovider/tests/billing.rs +53 −0 View agent/lprovider/tests/billing.rs
@@ -0,0 +1,53 @@
1use config::lprovider::PointsPer1k;
2use lprovider::billing::{ceil_div, compute_bill, points_for_tokens, worst_case_estimate};
3
4#[test]
5fn ceil_div_rounds_up() {
6 assert_eq!(ceil_div(0, 1000), 0);
7 assert_eq!(ceil_div(-5, 1000), 0);
8 assert_eq!(ceil_div(1, 1000), 1);
9 assert_eq!(ceil_div(999, 1000), 1);
10 assert_eq!(ceil_div(1000, 1000), 1);
11 assert_eq!(ceil_div(1001, 1000), 2);
12 assert_eq!(ceil_div(2000, 1000), 2);
13}
14
15#[test]
16fn points_for_tokens_uses_rate() {
17 let rate = PointsPer1k {
18 points_per_1k_input_tokens: 10,
19 points_per_1k_output_tokens: 30,
20 points_per_1k_cache_hit_tokens: 1,
21 };
22 assert_eq!(points_for_tokens(1500, rate.points_per_1k_input_tokens), 15);
23 assert_eq!(points_for_tokens(1, rate.points_per_1k_output_tokens), 1);
24 assert_eq!(points_for_tokens(0, rate.points_per_1k_cache_hit_tokens), 0);
25}
26
27#[test]
28fn compute_bill_input_excludes_cache_hit() {
29 let rate = PointsPer1k {
30 points_per_1k_input_tokens: 10,
31 points_per_1k_output_tokens: 30,
32 points_per_1k_cache_hit_tokens: 1,
33 };
34 let bill = compute_bill(2000, 500, 1500, &rate);
35 assert_eq!(bill.billable_input_tokens, 500);
36 assert_eq!(bill.billable_output_tokens, 500);
37 assert_eq!(bill.billable_cache_hit_tokens, 1500);
38 assert_eq!(bill.points_input, 5);
39 assert_eq!(bill.points_output, 15);
40 assert_eq!(bill.points_cache_hit, 2);
41 assert_eq!(bill.points_total, 22);
42}
43
44#[test]
45fn worst_case_estimate_uses_max_rate() {
46 let rate = PointsPer1k {
47 points_per_1k_input_tokens: 10,
48 points_per_1k_output_tokens: 30,
49 points_per_1k_cache_hit_tokens: 1,
50 };
51 assert_eq!(worst_case_estimate(1000, 2000, &rate), 90);
52 assert_eq!(worst_case_estimate(0, 0, &rate), 0);
53}
Added agent/lprovider/tests/error.rs agent/lprovider/tests/error.rs +65 −0 View agent/lprovider/tests/error.rs
@@ -0,0 +1,65 @@
1use lprovider::error::{AuthError, LproviderError, UpstreamError, WalletError};
2
3#[test]
4fn status_code_mapping() {
5 assert_eq!(
6 LproviderError::Auth(AuthError::MissingToken).status_code(),
7 401
8 );
9 assert_eq!(
10 LproviderError::Auth(AuthError::InvalidToken).status_code(),
11 401
12 );
13 assert_eq!(
14 LproviderError::Auth(AuthError::PermissionDenied("x".into())).status_code(),
15 403
16 );
17 assert_eq!(
18 LproviderError::Wallet(WalletError::WalletRequired).status_code(),
19 402
20 );
21 assert_eq!(
22 LproviderError::Wallet(WalletError::InsufficientBalance).status_code(),
23 402
24 );
25 assert_eq!(
26 LproviderError::Wallet(WalletError::WalletNotFound(7)).status_code(),
27 404
28 );
29 assert_eq!(
30 LproviderError::Wallet(WalletError::OwnerMismatch).status_code(),
31 403
32 );
33 assert_eq!(
34 LproviderError::Upstream(UpstreamError::NoUpstreamAvailable).status_code(),
35 503
36 );
37 assert_eq!(
38 LproviderError::Upstream(UpstreamError::AllFailed("x".into())).status_code(),
39 502
40 );
41 assert_eq!(
42 LproviderError::Upstream(UpstreamError::ContextLengthExceeded).status_code(),
43 400
44 );
45 assert_eq!(LproviderError::QuotaExceeded("q".into()).status_code(), 429);
46 assert_eq!(
47 LproviderError::InvalidRequest("x".into()).status_code(),
48 400
49 );
50 assert_eq!(
51 LproviderError::ResourceLocked("x".into()).status_code(),
52 409
53 );
54 assert_eq!(LproviderError::NotFound("x".into()).status_code(), 404);
55 assert_eq!(LproviderError::Internal("x".into()).status_code(), 500);
56 assert_eq!(LproviderError::Config("x".into()).status_code(), 500);
57 assert_eq!(LproviderError::Redis("x".into()).status_code(), 503);
58}
59
60#[test]
61fn error_types_in_message() {
62 let err = LproviderError::Wallet(WalletError::InsufficientBalance);
63 let msg = err.to_string();
64 assert!(msg.contains("insufficient"), "unexpected message: {msg}");
65}
Added agent/lprovider/tests/selector.rs agent/lprovider/tests/selector.rs +177 −0 View agent/lprovider/tests/selector.rs
@@ -0,0 +1,177 @@
1use std::sync::Arc;
2use std::sync::atomic::{AtomicI64, AtomicU64};
3
4use config::lprovider::{LproviderConfig, PriorityStrategyOverride};
5use llm::provider::LlmProvider;
6use llm::provider::openai::{OpenAIConfig, OpenAIProvider};
7use lprovider::selector::{CandidateFilter, rank, resolve_strategy};
8use lprovider::upstream::{KeyEntry, UpstreamEntry, UpstreamSnapshot};
9use tokio::sync::RwLock;
10
11fn dummy_provider() -> Arc<dyn LlmProvider> {
12 let cfg = OpenAIConfig::new("test-key").with_model("test-model");
13 Arc::new(OpenAIProvider::new(cfg).expect("provider builds"))
14}
15
16fn key(label: &str) -> KeyEntry {
17 KeyEntry {
18 label: label.to_string(),
19 fingerprint: "fp".to_string(),
20 qps_limit: 0,
21 rpm_limit: 0,
22 provider: dummy_provider(),
23 in_flight: Arc::new(AtomicI64::new(0)),
24 cooldown_until: RwLock::new(None),
25 rr: AtomicU64::new(0),
26 }
27}
28
29fn upstream(code: &str, cost: i32, quality: i32, latency: i32, ctx: usize) -> UpstreamEntry {
30 UpstreamEntry {
31 code: code.to_string(),
32 default_model: format!("{code}-model"),
33 capabilities: vec!["streaming".to_string(), "tools".to_string()],
34 context_window: ctx,
35 max_output: 8192,
36 cost_tier: cost,
37 quality_tier: quality,
38 latency_tier: latency,
39 tags: vec![],
40 keys: vec![key("k1")],
41 }
42}
43
44fn snapshot(ups: Vec<UpstreamEntry>) -> UpstreamSnapshot {
45 UpstreamSnapshot {
46 upstreams: ups,
47 estimator: dummy_provider(),
48 }
49}
50
51#[test]
52fn resolve_strategy_priority() {
53 let overrides = vec![
54 PriorityStrategyOverride {
55 min_priority: 10,
56 max_priority: None,
57 strategy: "QUALITY".to_string(),
58 },
59 PriorityStrategyOverride {
60 min_priority: -100,
61 max_priority: Some(-10),
62 strategy: "COST".to_string(),
63 },
64 ];
65 assert_eq!(
66 resolve_strategy("BALANCED", &overrides, 5, 10_000, None, 1000),
67 "BALANCED"
68 );
69 assert_eq!(
70 resolve_strategy("BALANCED", &overrides, 12, 10_000, None, 1000),
71 "QUALITY"
72 );
73 assert_eq!(
74 resolve_strategy("BALANCED", &overrides, -50, 10_000, None, 1000),
75 "COST"
76 );
77 assert_eq!(
78 resolve_strategy("BALANCED", &overrides, -10, 10_000, None, 1000),
79 "BALANCED"
80 );
81}
82
83#[test]
84fn resolve_strategy_explicit_and_low_balance() {
85 assert_eq!(
86 resolve_strategy("BALANCED", &[], 0, 10_000, Some("LATENCY"), 1000),
87 "LATENCY"
88 );
89 assert_eq!(
90 resolve_strategy("QUALITY", &[], 0, 999, Some("LATENCY"), 1000),
91 "LATENCY"
92 );
93 assert_eq!(resolve_strategy("QUALITY", &[], 0, 500, None, 1000), "COST");
94 assert_eq!(
95 resolve_strategy("QUALITY", &[], 0, 1000, None, 1000),
96 "QUALITY"
97 );
98}
99
100#[test]
101fn rank_filters_alias_and_context() {
102 let snap = snapshot(vec![
103 upstream("cheap", 1, 3, 5, 64_000),
104 upstream("best", 5, 5, 1, 200_000),
105 ]);
106 let filter = CandidateFilter {
107 required_capabilities: vec!["streaming".to_string()],
108 input_tokens: 1000,
109 max_output: 8000,
110 };
111 let ranked =
112 rank(&snap, "auto", &filter, "COST", &[], &[], 0.0, 200_000).expect("has candidates");
113 assert_eq!(ranked.len(), 2);
114 assert_eq!(ranked[0].upstream_idx, 0);
115 assert_eq!(ranked[0].model, "cheap-model");
116
117 let q =
118 rank(&snap, "auto", &filter, "QUALITY", &[], &[], 0.0, 200_000).expect("has candidates");
119 assert_eq!(q[0].upstream_idx, 1);
120
121 let miss = rank(
122 &snap,
123 "nonexistent",
124 &filter,
125 "COST",
126 &[],
127 &[],
128 0.0,
129 200_000,
130 );
131 assert!(miss.is_err());
132}
133
134#[test]
135fn rank_context_exceeded() {
136 let snap = snapshot(vec![upstream("small", 1, 1, 1, 4_000)]);
137 let filter = CandidateFilter {
138 required_capabilities: vec![],
139 input_tokens: 10_000,
140 max_output: 1000,
141 };
142 let err = rank(&snap, "auto", &filter, "COST", &[], &[], 0.0, 200_000)
143 .unwrap_err()
144 .to_string();
145 assert!(err.contains("context"), "unexpected error: {err}");
146}
147
148#[test]
149fn rank_explicit_allowed_order() {
150 let snap = snapshot(vec![
151 upstream("b", 1, 1, 1, 64_000),
152 upstream("a", 1, 1, 1, 64_000),
153 ]);
154 let filter = CandidateFilter {
155 required_capabilities: vec![],
156 input_tokens: 10,
157 max_output: 10,
158 };
159 let ranked = rank(
160 &snap,
161 "auto",
162 &filter,
163 "EXPLICIT",
164 &["a".to_string(), "b".to_string()],
165 &[],
166 0.0,
167 200_000,
168 )
169 .expect("has candidates");
170 assert_eq!(ranked[0].model, "a-model");
171}
172
173#[test]
174fn rank_respects_default_config_low_balance_strategy() {
175 let cfg = LproviderConfig::default();
176 assert_eq!(cfg.low_balance_force_cost_points, 1000);
177}
Added agent/lprovider/admin.proto agent/lprovider/admin.proto +479 −0 View agent/lprovider/admin.proto
@@ -0,0 +1,479 @@
1syntax = "proto3";
2package lprovider.admin.v1;
3
4import "google/protobuf/struct.proto";
5
6service LproviderAdmin {
7 rpc GetConfig(GetConfigRequest) returns (GetConfigResponse);
8 rpc UpdateConfig(UpdateConfigRequest) returns (UpdateConfigResponse);
9 rpc ReloadConfig(ReloadConfigRequest) returns (ReloadConfigResponse);
10 rpc RestartService(RestartServiceRequest) returns (RestartServiceResponse);
11
12 rpc ListUpstreams(ListUpstreamsRequest) returns (ListUpstreamsResponse);
13 rpc GetUpstream(GetUpstreamRequest) returns (Upstream);
14
15 rpc ListUpstreamKeys(ListUpstreamKeysRequest) returns (ListUpstreamKeysResponse);
16 rpc GetUpstreamKey(GetUpstreamKeyRequest) returns (UpstreamKey);
17
18 rpc ListBindings(ListBindingsRequest) returns (ListBindingsResponse);
19 rpc InvalidateBinding(InvalidateBindingRequest) returns (InvalidateBindingResponse);
20
21 rpc QueryAudit(QueryAuditRequest) returns (QueryAuditResponse);
22 rpc AggregateBilling(AggregateBillingRequest) returns (AggregateBillingResponse);
23 rpc ListTenantQuotas(ListTenantQuotasRequest) returns (ListTenantQuotasResponse);
24 rpc UpsertTenantQuota(UpsertTenantQuotaRequest) returns (TenantQuota);
25
26 rpc GetWallet(GetWalletRequest) returns (Wallet);
27 rpc ListWallets(ListWalletsRequest) returns (ListWalletsResponse);
28 rpc TopupWallet(TopupWalletRequest) returns (Wallet);
29 rpc RefundWallet(RefundWalletRequest) returns (Wallet);
30 rpc AdjustWallet(AdjustWalletRequest) returns (Wallet);
31 rpc FreezeWallet(FreezeWalletRequest) returns (Wallet);
32 rpc UnfreezeWallet(UnfreezeWalletRequest) returns (Wallet);
33 rpc DisableWallet(DisableWalletRequest) returns (Wallet);
34 rpc ListWalletLedger(ListWalletLedgerRequest) returns (ListWalletLedgerResponse);
35 rpc ListWalletFreezeLog(ListWalletFreezeLogRequest) returns (ListWalletFreezeLogResponse);
36
37 rpc GetRuntimeStatus(GetRuntimeStatusRequest) returns (RuntimeStatus);
38 rpc ListInstances(ListInstancesRequest) returns (ListInstancesResponse);
39}
40
41message Upstream {
42 string code = 1;
43 string provider_type = 2;
44 string base_url = 3;
45 string default_model = 4;
46 repeated string capabilities = 5;
47 int32 context_window = 6;
48 int32 max_output = 7;
49 bool enabled = 8;
50 int32 cost_tier = 9;
51 int32 quality_tier = 10;
52 int32 latency_tier = 11;
53 repeated string tags = 12;
54}
55
56message UpstreamKey {
57 string upstream_code = 1;
58 string label = 2;
59 string credential_ref = 3;
60 string key_fingerprint = 4;
61 string status = 5;
62 int32 qps_limit = 6;
63 int32 rpm_limit = 7;
64 int64 cooldown_until_unix = 8;
65 KeyStats stats = 9;
66}
67
68message KeyStats {
69 int64 success_count = 1;
70 int64 error_count = 2;
71 int64 rate_limited_count = 3;
72 int32 in_flight = 4;
73}
74
75message SessionBinding {
76 string model_alias = 1;
77 string session_key = 2;
78 string upstream_code = 3;
79 string key_label = 4;
80 string prefix_hash = 5;
81 int64 bound_at_unix = 6;
82 int64 expires_at_unix = 7;
83 int32 last_hit_tokens = 8;
84 int64 hit_count = 9;
85}
86
87message CallRecord {
88 int64 id = 1;
89 string request_id = 2;
90 string trace_id = 3;
91 int64 tenant_id = 4;
92 int64 api_key_id = 5;
93 string subject_type = 6;
94 int64 subject_id = 7;
95 int64 wallet_id = 8;
96 bool anonymous = 9;
97 string subject_path = 10;
98 string session_key = 11;
99 string model_alias = 12;
100 string upstream_code = 13;
101 string requested_model = 14;
102 string chosen_strategy = 15;
103 int32 tier_cost = 16;
104 int32 tier_quality = 17;
105 int32 tier_latency = 18;
106 string key_label = 19;
107 string caller = 20;
108 string caller_service = 21;
109 int32 priority = 22;
110 bool streamed = 23;
111 string status = 24;
112 string error_kind = 25;
113 string error_message = 26;
114 int32 prompt_tokens = 27;
115 int32 completion_tokens = 28;
116 int32 cache_hit_tokens = 29;
117 int32 cache_miss_tokens = 30;
118 int32 billable_input_tokens = 31;
119 int32 billable_output_tokens = 32;
120 int64 points_input = 33;
121 int64 points_output = 34;
122 int64 points_cache_hit = 35;
123 int64 points_total = 36;
124 int64 points_rate_input = 37;
125 int64 points_rate_output = 38;
126 int64 points_rate_cache_hit = 39;
127 int64 request_started_at_ms = 40;
128 int64 response_completed_at_ms = 41;
129 int32 latency_ms = 42;
130 string client_ip = 43;
131 string user_agent = 44;
132 int64 created_at_ms = 45;
133}
134
135message BillingAggregate {
136 int64 id = 1;
137 int64 tenant_id = 2;
138 string subject_type = 3;
139 int64 subject_id = 4;
140 int64 period_start_ms = 5;
141 int64 period_end_ms = 6;
142 string period_kind = 7;
143 string model_alias = 8;
144 string upstream_code = 9;
145 int64 call_count = 10;
146 int64 success_count = 11;
147 int64 error_count = 12;
148 int64 total_input_tokens = 13;
149 int64 total_output_tokens = 14;
150 int64 total_cache_hit_tokens = 15;
151 int64 total_points = 16;
152 int64 updated_at_ms = 17;
153}
154
155message Wallet {
156 int64 wallet_id = 1;
157 string owner_type = 2;
158 int64 owner_id = 3;
159 int64 tenant_id = 4;
160 string status = 5;
161 int64 balance_points = 6;
162 int64 total_topped_up = 7;
163 int64 total_consumed = 8;
164 int64 total_refunded = 9;
165 int64 created_at_unix = 10;
166 int64 updated_at_unix = 11;
167}
168
169message GetWalletRequest {
170 int64 wallet_id = 1;
171}
172
173message ListWalletsRequest {
174 int64 tenant_id = 1;
175 string owner_type = 2;
176 int64 owner_id = 3;
177}
178
179message ListWalletsResponse {
180 repeated Wallet wallets = 1;
181}
182
183message TopupWalletRequest {
184 int64 wallet_id = 1;
185 int64 delta_points = 2;
186 string reason = 3;
187 string operator = 4;
188 string idempotency_key = 5;
189}
190
191message RefundWalletRequest {
192 int64 wallet_id = 1;
193 int64 delta_points = 2;
194 string request_id = 3;
195 string reason = 4;
196 string operator = 5;
197}
198
199message AdjustWalletRequest {
200 int64 wallet_id = 1;
201 int64 delta_points = 2;
202 string reason = 3;
203 string operator = 4;
204 string idempotency_key = 5;
205}
206
207message FreezeWalletRequest {
208 int64 wallet_id = 1;
209 string reason = 2;
210 string operator = 3;
211}
212
213message UnfreezeWalletRequest {
214 int64 wallet_id = 1;
215 string reason = 2;
216 string operator = 3;
217}
218
219message DisableWalletRequest {
220 int64 wallet_id = 1;
221 string reason = 2;
222 string operator = 3;
223}
224
225message WalletLedgerEntry {
226 int64 id = 1;
227 int64 wallet_id = 2;
228 string request_id = 3;
229 int64 delta_points = 4;
230 int64 balance_after = 5;
231 string kind = 6;
232 string reason = 7;
233 string operator = 8;
234 int64 created_at_unix = 9;
235}
236
237message ListWalletLedgerRequest {
238 int64 wallet_id = 1;
239 int64 from_unix = 2;
240 int64 to_unix = 3;
241}
242
243message ListWalletLedgerResponse {
244 repeated WalletLedgerEntry entries = 1;
245}
246
247message ListWalletFreezeLogRequest {
248 int64 wallet_id = 1;
249}
250
251message WalletFreezeLogEntry {
252 int64 id = 1;
253 int64 wallet_id = 2;
254 string action = 3;
255 string reason = 4;
256 string operator = 5;
257 int64 created_at_unix = 6;
258}
259
260message ListWalletFreezeLogResponse {
261 repeated WalletFreezeLogEntry entries = 1;
262}
263
264message TenantQuota {
265 int64 tenant_id = 1;
266 string subject_type = 2;
267 int64 subject_id = 3;
268 int64 monthly_quota_points = 4;
269 int64 monthly_quota_input_tokens = 5;
270 int64 monthly_quota_output_tokens = 6;
271 bool hard_limit = 7;
272 bool enabled = 8;
273 int64 updated_at_unix = 9;
274}
275
276message AffinityStats {
277 int64 bind_total = 1;
278 int64 hit_count = 2;
279 int64 miss_count = 3;
280}
281
282message UpstreamSummary {
283 string upstream_code = 1;
284 bool any_key_available = 2;
285 int32 keys_active = 3;
286 int32 keys_cooldown = 4;
287 double error_rate_5m = 5;
288}
289
290message ListTenantQuotasRequest {
291 int64 tenant_id = 1;
292 string subject_type = 2;
293 int64 subject_id = 3;
294}
295
296message ListTenantQuotasResponse {
297 repeated TenantQuota quotas = 1;
298}
299
300message UpsertTenantQuotaRequest {
301 TenantQuota quota = 1;
302}
303
304message RuntimeStatus {
305 string snapshot_version = 1;
306 int32 upstream_total = 2;
307 int32 upstream_active = 3;
308 int32 upstream_cooldown = 4;
309 bool redis_ok = 5;
310 bool db_ok = 6;
311 AffinityStats affinity = 7;
312 UpstreamSummary upstream_summary = 8;
313}
314
315message GetConfigRequest {
316 optional bool mask_secrets = 1;
317}
318
319message GetConfigResponse {
320 string config_path = 1;
321 string sha256 = 2;
322 int64 file_size_bytes = 3;
323 int64 loaded_at_unix = 4;
324 string content_yaml = 5;
325 repeated string warnings = 6;
326}
327
328message UpdateConfigRequest {
329 oneof payload {
330 string full_content_yaml = 1;
331 ConfigPatch patch = 2;
332 ConfigSectionPatch section_patch = 3;
333 }
334 bool reload_after_write = 4;
335 bool restart_after_write = 5;
336 string operator = 6;
337}
338
339message ConfigPatch {
340 repeated ConfigSectionPatch sections = 1;
341}
342
343message ConfigSectionPatch {
344 string path = 1;
345 google.protobuf.Struct value = 2;
346}
347
348message UpdateConfigResponse {
349 string sha256 = 1;
350 string backup_path = 2;
351 bool reloaded = 3;
352 bool restarted = 4;
353 repeated string warnings = 5;
354}
355
356message ReloadConfigRequest {
357 string operator = 1;
358}
359
360message ReloadConfigResponse {
361 string previous_sha256 = 1;
362 string current_sha256 = 2;
363 repeated string warnings = 3;
364 RuntimeStatus status = 4;
365}
366
367message RestartServiceRequest {
368 string operator = 1;
369 int32 drain_timeout_seconds = 2;
370 bool self_reexec = 3;
371 bool rolling = 4;
372 string target_instance_id = 5;
373}
374
375message RestartServiceResponse {
376 bool accepted = 1;
377 string reason = 2;
378 string instance_id = 3;
379}
380
381message InstanceInfo {
382 string instance_id = 1;
383 string cluster_name = 2;
384 string listen_addr = 3;
385 int64 started_at_unix = 4;
386 int32 pid = 5;
387 string config_sha256 = 6;
388 bool is_config_writer = 7;
389}
390
391message ListInstancesRequest {
392 string cluster_name = 1;
393}
394
395message ListInstancesResponse {
396 repeated InstanceInfo instances = 1;
397}
398
399// --- 以下请求/响应消息在文档 §6.2 中未逐字定义，按语义最小化补齐 ---
400
401message ListUpstreamsRequest {
402 repeated string codes = 1; // 空 = 全部
403}
404
405message ListUpstreamsResponse {
406 repeated Upstream upstreams = 1;
407}
408
409message GetUpstreamRequest {
410 string code = 1;
411}
412
413message ListUpstreamKeysRequest {
414 string upstream_code = 1; // 空 = 全部上游
415}
416
417message ListUpstreamKeysResponse {
418 repeated UpstreamKey keys = 1;
419}
420
421message GetUpstreamKeyRequest {
422 string upstream_code = 1;
423 string label = 2;
424}
425
426message ListBindingsRequest {
427 int64 tenant_id = 1; // 0 = 全部租户
428 string model_alias = 2; // 空 = 全部别名
429 string session_key = 3; // 空 = 全部会话
430}
431
432message ListBindingsResponse {
433 repeated SessionBinding bindings = 1;
434}
435
436message InvalidateBindingRequest {
437 int64 tenant_id = 1;
438 string model_alias = 2;
439 string session_key = 3;
440}
441
442message InvalidateBindingResponse {
443 bool invalidated = 1;
444}
445
446message QueryAuditRequest {
447 int64 tenant_id = 1; // 0 = 全部
448 string subject_type = 2;
449 int64 subject_id = 3; // 0 = 全部
450 string model_alias = 4;
451 string upstream_code = 5;
452 string status = 6; // 空 = 全部
453 int64 from_ms = 7;
454 int64 to_ms = 8;
455 int32 limit = 9; // 默认 100
456 int32 offset = 10;
457}
458
459message QueryAuditResponse {
460 repeated CallRecord records = 1;
461 int64 total = 2;
462}
463
464message AggregateBillingRequest {
465 int64 tenant_id = 1; // 0 = 全部
466 string subject_type = 2;
467 int64 subject_id = 3; // 0 = 全部
468 string period_kind = 4; // hour / day / month
469 int64 period_start_ms = 5;
470 int64 period_end_ms = 6;
471 string model_alias = 7;
472 string upstream_code = 8;
473}
474
475message AggregateBillingResponse {
476 repeated BillingAggregate aggregates = 1;
477}
478
479message GetRuntimeStatusRequest {}
Added agent/lprovider/admin_grpc.rs agent/lprovider/admin_grpc.rs +1531 −0 View agent/lprovider/admin_grpc.rs
@@ -0,0 +1,1531 @@
1use std::sync::Arc;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
4
5use redis::AsyncCommands;
6use time::OffsetDateTime;
7use tokio_stream::StreamExt;
8use tonic::{Request, Response, Status};
9use tracing::warn;
10
11use crate::affinity::AffinityManager;
12use crate::ai_grpc::status_from_err;
13use crate::audit::AuditService;
14use crate::config::RuntimeConfig;
15use crate::error::LproviderError;
16use crate::lprovider::admin::v1 as pb;
17use crate::quota::{QuotaService, TenantQuota};
18use crate::state::AppState;
19use crate::upstream::UpstreamManager;
20use crate::wallet::{WalletInfo, WalletService};
21
22const RELEASE_LOCK_SCRIPT: &str = r#"
23if redis.call('GET', KEYS[1]) == ARGV[1] then
24 return redis.call('DEL', KEYS[1])
25end
26return 0
27"#;
28
29const RENEW_LOCK_SCRIPT: &str = r#"
30if redis.call('GET', KEYS[1]) == ARGV[1] then
31 return redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]))
32end
33return 0
34"#;
35
36#[derive(Clone)]
37pub struct AdminGrpc {
38 pub state: Arc<AppState>,
39 pub manager: Arc<UpstreamManager>,
40 pub affinity: Arc<AffinityManager>,
41 pub wallets: Arc<WalletService>,
42 pub quotas: Arc<QuotaService>,
43 pub audit: Arc<AuditService>,
44}
45
46struct ConfigLockGuard {
47 key: String,
48 token: String,
49 redis: Arc<cache::RedisCache>,
50 released: Arc<AtomicBool>,
51}
52
53impl Drop for ConfigLockGuard {
54 fn drop(&mut self) {
55 self.released.store(true, Ordering::Relaxed);
56 let redis = self.redis.clone();
57 let key = self.key.clone();
58 let token = self.token.clone();
59 tokio::spawn(async move {
60 let mut c = redis.connection();
61 let script = redis::Script::new(RELEASE_LOCK_SCRIPT);
62 let _: redis::RedisResult<i64> =
63 script.key(&key).arg(&token).invoke_async(&mut c).await;
64 });
65 }
66}
67
68impl AdminGrpc {
69 pub fn new(
70 state: Arc<AppState>,
71 manager: Arc<UpstreamManager>,
72 affinity: Arc<AffinityManager>,
73 ) -> Self {
74 let db = state.db.clone().map(Arc::new);
75 Self {
76 state,
77 manager,
78 affinity,
79 wallets: Arc::new(WalletService::new(db.clone())),
80 quotas: Arc::new(QuotaService::new(db.clone())),
81 audit: Arc::new(AuditService::new(db)),
82 }
83 }
84
85 async fn acquire_config_lock(&self) -> Result<Option<ConfigLockGuard>, LproviderError> {
86 let Some(redis) = &self.state.redis else {
87 return Ok(None);
88 };
89 let key = format!("{}:lock:config_writer", self.affinity.key_prefix());
90 let token = format!("{}:{:016x}", self.state.instance_id, rand::random::<u64>());
91 let mut conn = redis.connection();
92 let ok: bool = conn.set_nx(&key, &token).await.map_err(redis_err)?;
93 if !ok {
94 return Err(LproviderError::ResourceLocked(
95 "config_writer lock held by another instance".into(),
96 ));
97 }
98 let _: () = conn.expire(&key, 30).await.map_err(redis_err)?;
99 let released = Arc::new(AtomicBool::new(false));
100 let renew_flag = released.clone();
101 let renew_redis = redis.clone();
102 let renew_key = key.clone();
103 let renew_token = token.clone();
104 tokio::spawn(async move {
105 loop {
106 tokio::time::sleep(Duration::from_secs(10)).await;
107 if renew_flag.load(Ordering::Relaxed) {
108 break;
109 }
110 let mut c = renew_redis.connection();
111 let script = redis::Script::new(RENEW_LOCK_SCRIPT);
112 let res: redis::RedisResult<i64> = script
113 .key(&renew_key)
114 .arg(&renew_token)
115 .arg(30i64)
116 .invoke_async(&mut c)
117 .await;
118 if !matches!(res, Ok(1)) {
119 break;
120 }
121 }
122 });
123 Ok(Some(ConfigLockGuard {
124 key,
125 token,
126 redis: redis.clone(),
127 released,
128 }))
129 }
130
131 #[allow(clippy::too_many_arguments)]
132 async fn audit_write(
133 &self,
134 operator: &str,
135 rpc: &str,
136 target: Option<&str>,
137 summary: Option<&str>,
138 effect_mode: Option<&str>,
139 result: &str,
140 error: Option<&str>,
141 ) {
142 if let Err(e) = self
143 .audit
144 .admin_audit(
145 operator,
146 rpc,
147 Some(&self.state.instance_id),
148 target,
149 summary,
150 None,
151 effect_mode,
152 result,
153 error,
154 )
155 .await
156 {
157 warn!(rpc, error = %e, "admin audit write failed");
158 }
159 }
160
161 async fn build_runtime_status(&self) -> pb::RuntimeStatus {
162 let snapshot = self.manager.snapshot().await;
163 let mut active = 0i32;
164 let mut cooldown = 0i32;
165 for up in &snapshot.upstreams {
166 for k in &up.keys {
167 let in_cooldown = match *k.cooldown_until.read().await {
168 Some(i) => !i.saturating_duration_since(Instant::now()).is_zero(),
169 None => false,
170 };
171 if in_cooldown {
172 cooldown += 1;
173 } else {
174 active += 1;
175 }
176 }
177 }
178 pb::RuntimeStatus {
179 snapshot_version: self.state.snapshot_version(),
180 upstream_total: snapshot.upstreams.len() as i32,
181 upstream_active: active,
182 upstream_cooldown: cooldown,
183 redis_ok: self.state.redis_ok().await,
184 db_ok: self.state.db_writable().await,
185 affinity: Some(pb::AffinityStats {
186 bind_total: self.state.affinity.bind_total.load(Ordering::Relaxed),
187 hit_count: self.state.affinity.hit_count.load(Ordering::Relaxed),
188 miss_count: self.state.affinity.miss_count.load(Ordering::Relaxed),
189 }),
190 upstream_summary: Some(pb::UpstreamSummary {
191 upstream_code: "<all>".into(),
192 any_key_available: active > 0,
193 keys_active: active,
194 keys_cooldown: cooldown,
195 error_rate_5m: 0.0,
196 }),
197 }
198 }
199
200 async fn reload_runtime(&self) -> Result<(), LproviderError> {
201 let mut cli = self.state.cli.clone();
202 cli.config = self.state.runtime.read().await.config_path.clone();
203 let new_rt = RuntimeConfig::load(&cli)?;
204 self.manager.reload(new_rt.lp())?;
205 *self.state.runtime.write().await = new_rt;
206 Ok(())
207 }
208}
209
210fn redis_err(e: redis::RedisError) -> LproviderError {
211 LproviderError::Redis(e.to_string())
212}
213
214fn redis_status(e: redis::RedisError) -> Status {
215 status_from_err(&LproviderError::Redis(e.to_string()))
216}
217
218fn operator_of(req_operator: &str) -> String {
219 let t = req_operator.trim();
220 if t.is_empty() {
221 "unknown".to_string()
222 } else {
223 t.to_string()
224 }
225}
226
227fn now_unix() -> i64 {
228 SystemTime::now()
229 .duration_since(UNIX_EPOCH)
230 .map(|d| d.as_secs() as i64)
231 .unwrap_or(0)
232}
233
234fn unix_ms(dt: OffsetDateTime) -> i64 {
235 dt.unix_timestamp() * 1000 + dt.millisecond() as i64
236}
237
238fn mask_secrets(content: &str) -> String {
239 content
240 .lines()
241 .map(|line| {
242 if line.contains("bearer_token") || line.contains("api_key") {
243 match line.find(':') {
244 Some(idx) => format!("{}: \"***\"", &line[..idx]),
245 None => line.to_string(),
246 }
247 } else {
248 line.to_string()
249 }
250 })
251 .collect::<Vec<_>>()
252 .join("\n")
253}
254
255fn infra_changed(old: &config::AppConfig, new: &config::AppConfig) -> bool {
256 let o = &old.lprovider;
257 let n = &new.lprovider;
258 o.bind != n.bind
259 || o.port != n.port
260 || o.admin_bind != n.admin_bind
261 || o.admin_port != n.admin_port
262 || o.enabled != n.enabled
263 || o.max_context_tokens != n.max_context_tokens
264 || o.max_output_tokens != n.max_output_tokens
265 || o.request_timeout_ms != n.request_timeout_ms
266 || o.redis.key_prefix != n.redis.key_prefix
267 || o.redis.session_affinity_ttl_seconds != n.redis.session_affinity_ttl_seconds
268 || o.redis.min_prefix_tokens != n.redis.min_prefix_tokens
269 || o.redis.prefix_token_budget != n.redis.prefix_token_budget
270 || o.redis.cooldown_seconds != n.redis.cooldown_seconds
271 || o.secret.envelope_key_env != n.secret.envelope_key_env
272 || serde_json::to_string(&old.db).ok() != serde_json::to_string(&new.db).ok()
273}
274
275fn set_dotted(
276 doc: &mut serde_yaml::Value,
277 path: &str,
278 value: serde_yaml::Value,
279) -> Result<(), LproviderError> {
280 let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
281 if parts.is_empty() {
282 return Err(LproviderError::InvalidRequest("empty patch path".into()));
283 }
284 let mut cur: &mut serde_yaml::Value = doc;
285 for part in &parts[..parts.len() - 1] {
286 if let Ok(idx) = part.parse::<usize>() {
287 cur = {
288 let seq = cur.as_sequence_mut().ok_or_else(|| {
289 LproviderError::InvalidRequest(format!(
290 "patch segment '{part}' is not a sequence"
291 ))
292 })?;
293 if idx >= seq.len() {
294 seq.resize(idx + 1, serde_yaml::Value::Null);
295 }
296 &mut seq[idx]
297 };
298 } else {
299 cur = {
300 let map = cur.as_mapping_mut().ok_or_else(|| {
301 LproviderError::InvalidRequest(format!(
302 "patch segment '{part}' is not a mapping"
303 ))
304 })?;
305 let key = serde_yaml::Value::String((*part).to_string());
306 if !map.contains_key(&key) {
307 map.insert(key.clone(), serde_yaml::Value::Null);
308 }
309 map.get_mut(&key).expect("inserted")
310 };
311 }
312 }
313 let last = parts[parts.len() - 1];
314 if let Ok(idx) = last.parse::<usize>() {
315 let seq = cur.as_sequence_mut().ok_or_else(|| {
316 LproviderError::InvalidRequest("patch parent is not a sequence".into())
317 })?;
318 if idx >= seq.len() {
319 seq.resize(idx + 1, serde_yaml::Value::Null);
320 }
321 seq[idx] = value;
322 } else {
323 let map = cur.as_mapping_mut().ok_or_else(|| {
324 LproviderError::InvalidRequest("patch parent is not a mapping".into())
325 })?;
326 map.insert(serde_yaml::Value::String(last.to_string()), value);
327 }
328 Ok(())
329}
330
331async fn compose_yaml(
332 req: &pb::UpdateConfigRequest,
333 config_path: &std::path::Path,
334) -> Result<String, LproviderError> {
335 match &req.payload {
336 Some(pb::update_config_request::Payload::FullContentYaml(y)) => Ok(y.clone()),
337 Some(pb::update_config_request::Payload::SectionPatch(p)) => {
338 let current = tokio::fs::read_to_string(config_path)
339 .await
340 .map_err(|e| LproviderError::Internal(format!("read config: {e}")))?;
341 let doc: serde_yaml::Value = serde_yaml::from_str(¤t).map_err(|e| {
342 LproviderError::InvalidRequest(format!("current config not valid yaml: {e}"))
343 })?;
344 let value = p.value.as_ref().ok_or_else(|| {
345 LproviderError::InvalidRequest("section patch value missing".into())
346 })?;
347 let patched = apply_patch(doc, &p.path, value)?;
348 serde_yaml::to_string(&patched)
349 .map_err(|e| LproviderError::InvalidRequest(format!("serialize yaml: {e}")))
350 }
351 Some(pb::update_config_request::Payload::Patch(p)) => {
352 let current = tokio::fs::read_to_string(config_path)
353 .await
354 .map_err(|e| LproviderError::Internal(format!("read config: {e}")))?;
355 let mut doc: serde_yaml::Value = serde_yaml::from_str(¤t).map_err(|e| {
356 LproviderError::InvalidRequest(format!("current config not valid yaml: {e}"))
357 })?;
358 for s in &p.sections {
359 let value = s.value.as_ref().ok_or_else(|| {
360 LproviderError::InvalidRequest("patch section value missing".into())
361 })?;
362 doc = apply_patch(doc, &s.path, value)?;
363 }
364 serde_yaml::to_string(&doc)
365 .map_err(|e| LproviderError::InvalidRequest(format!("serialize yaml: {e}")))
366 }
367 None => Err(LproviderError::InvalidRequest("no payload provided".into())),
368 }
369}
370
371fn apply_patch(
372 mut doc: serde_yaml::Value,
373 path: &str,
374 value: &prost_types::Struct,
375) -> Result<serde_yaml::Value, LproviderError> {
376 let json = crate::ai_grpc::prost_struct_to_json(value);
377 let yaml_val: serde_yaml::Value = serde_yaml::to_value(json)
378 .map_err(|e| LproviderError::InvalidRequest(format!("convert patch value: {e}")))?;
379 set_dotted(&mut doc, path, yaml_val)?;
380 Ok(doc)
381}
382
383fn wallet_to_proto(w: &WalletInfo) -> pb::Wallet {
384 pb::Wallet {
385 wallet_id: w.wallet_id,
386 owner_type: w.owner_type.clone(),
387 owner_id: w.owner_id,
388 tenant_id: w.tenant_id,
389 status: w.status.clone(),
390 balance_points: w.balance_points,
391 total_topped_up: w.total_topped_up,
392 total_consumed: w.total_consumed,
393 total_refunded: w.total_refunded,
394 created_at_unix: w.created_at.unix_timestamp(),
395 updated_at_unix: w.updated_at.unix_timestamp(),
396 }
397}
398
399fn quota_to_proto(q: &TenantQuota) -> pb::TenantQuota {
400 pb::TenantQuota {
401 tenant_id: q.tenant_id,
402 subject_type: q.subject_type.clone(),
403 subject_id: q.subject_id,
404 monthly_quota_points: q.monthly_quota_points,
405 monthly_quota_input_tokens: q.monthly_quota_input_tokens,
406 monthly_quota_output_tokens: q.monthly_quota_output_tokens,
407 hard_limit: q.hard_limit,
408 enabled: q.enabled,
409 updated_at_unix: 0,
410 }
411}
412
413fn proto_to_quota(q: &pb::TenantQuota) -> TenantQuota {
414 TenantQuota {
415 tenant_id: q.tenant_id,
416 subject_type: q.subject_type.clone(),
417 subject_id: q.subject_id,
418 monthly_quota_points: q.monthly_quota_points,
419 monthly_quota_input_tokens: q.monthly_quota_input_tokens,
420 monthly_quota_output_tokens: q.monthly_quota_output_tokens,
421 hard_limit: q.hard_limit,
422 enabled: q.enabled,
423 }
424}
425
426fn upstream_to_proto(u: &config::lprovider::LproviderUpstream) -> pb::Upstream {
427 pb::Upstream {
428 code: u.code.clone(),
429 provider_type: u.provider_type.clone(),
430 base_url: u.base_url.clone(),
431 default_model: u.default_model.clone(),
432 capabilities: u.capabilities.clone(),
433 context_window: u.context_window.min(i32::MAX as usize) as i32,
434 max_output: u.max_output.min(i32::MAX as usize) as i32,
435 enabled: u.enabled,
436 cost_tier: u.cost_tier,
437 quality_tier: u.quality_tier,
438 latency_tier: u.latency_tier,
439 tags: u.tags.clone(),
440 }
441}
442
443fn call_log_to_proto(m: &db::lprovider::llm_call_log::Model) -> pb::CallRecord {
444 pb::CallRecord {
445 id: m.id,
446 request_id: m.request_id.to_string(),
447 trace_id: m.trace_id.map(|u| u.to_string()).unwrap_or_default(),
448 tenant_id: m.tenant_id,
449 api_key_id: m.api_key_id.unwrap_or(0),
450 subject_type: m.subject_type.clone().unwrap_or_default(),
451 subject_id: m.subject_id.unwrap_or(0),
452 wallet_id: m.wallet_id.unwrap_or(0),
453 anonymous: m.anonymous,
454 subject_path: m.subject_path.clone().unwrap_or_default(),
455 session_key: m.session_key.clone().unwrap_or_default(),
456 model_alias: m.model_alias.clone(),
457 upstream_code: m.upstream_code.clone(),
458 requested_model: m.requested_model.clone().unwrap_or_default(),
459 chosen_strategy: m.chosen_strategy.clone().unwrap_or_default(),
460 tier_cost: m.tier_cost,
461 tier_quality: m.tier_quality,
462 tier_latency: m.tier_latency,
463 key_label: m.key_label.clone().unwrap_or_default(),
464 caller: m.caller.clone().unwrap_or_default(),
465 caller_service: m.caller_service.clone().unwrap_or_default(),
466 priority: m.priority,
467 streamed: m.streamed,
468 status: m.status.clone(),
469 error_kind: m.error_kind.clone().unwrap_or_default(),
470 error_message: m.error_message.clone().unwrap_or_default(),
471 prompt_tokens: m.prompt_tokens,
472 completion_tokens: m.completion_tokens,
473 cache_hit_tokens: m.cache_hit_tokens,
474 cache_miss_tokens: m.cache_miss_tokens,
475 billable_input_tokens: m.billable_input_tokens,
476 billable_output_tokens: m.billable_output_tokens,
477 points_input: m.points_input,
478 points_output: m.points_output,
479 points_cache_hit: m.points_cache_hit,
480 points_total: m.points_total,
481 points_rate_input: m.points_rate_input,
482 points_rate_output: m.points_rate_output,
483 points_rate_cache_hit: m.points_rate_cache_hit,
484 request_started_at_ms: unix_ms(m.request_started_at),
485 response_completed_at_ms: m.response_completed_at.map(unix_ms).unwrap_or(0),
486 latency_ms: m.latency_ms,
487 client_ip: m.client_ip.clone().unwrap_or_default(),
488 user_agent: m.user_agent.clone().unwrap_or_default(),
489 created_at_ms: unix_ms(m.created_at),
490 }
491}
492
493fn aggregate_to_proto(m: &db::lprovider::llm_billing_aggregate::Model) -> pb::BillingAggregate {
494 pb::BillingAggregate {
495 id: m.id,
496 tenant_id: m.tenant_id,
497 subject_type: m.subject_type.clone(),
498 subject_id: m.subject_id,
499 period_start_ms: unix_ms(m.period_start),
500 period_end_ms: unix_ms(m.period_end),
501 period_kind: m.period_kind.clone(),
502 model_alias: m.model_alias.clone(),
503 upstream_code: m.upstream_code.clone(),
504 call_count: m.call_count,
505 success_count: m.success_count,
506 error_count: m.error_count,
507 total_input_tokens: m.total_input_tokens,
508 total_output_tokens: m.total_output_tokens,
509 total_cache_hit_tokens: m.total_cache_hit_tokens,
510 total_points: m.total_points,
511 updated_at_ms: unix_ms(m.updated_at),
512 }
513}
514
515async fn build_upstream_key(
516 up: &crate::upstream::UpstreamEntry,
517 k: &crate::upstream::KeyEntry,
518 cfg_credential_ref: Option<String>,
519 cfg_qps: i64,
520 cfg_rpm: i64,
521) -> pb::UpstreamKey {
522 let (status, cd_until) = match *k.cooldown_until.read().await {
523 Some(inst) => {
524 let rem = inst.saturating_duration_since(Instant::now());
525 if rem.is_zero() {
526 ("active".to_string(), 0)
527 } else {
528 ("cooldown".to_string(), now_unix() + rem.as_secs() as i64)
529 }
530 }
531 None => ("active".to_string(), 0),
532 };
533 pb::UpstreamKey {
534 upstream_code: up.code.clone(),
535 label: k.label.clone(),
536 credential_ref: cfg_credential_ref.unwrap_or_default(),
537 key_fingerprint: k.fingerprint.clone(),
538 status,
539 qps_limit: cfg_qps as i32,
540 rpm_limit: cfg_rpm as i32,
541 cooldown_until_unix: cd_until,
542 stats: Some(pb::KeyStats {
543 success_count: 0,
544 error_count: 0,
545 rate_limited_count: 0,
546 in_flight: k.in_flight.load(Ordering::Relaxed) as i32,
547 }),
548 }
549}
550
551fn self_instance_info(state: &AppState) -> pb::InstanceInfo {
552 let lp = state.runtime.blocking_read().lp().clone();
553 pb::InstanceInfo {
554 instance_id: state.instance_id.clone(),
555 cluster_name: state.cluster_name.clone(),
556 listen_addr: format!("{}:{}", lp.bind, lp.port),
557 started_at_unix: state.started_at.unix_timestamp(),
558 pid: state.pid as i32,
559 config_sha256: state.runtime.blocking_read().config_sha256.clone(),
560 is_config_writer: false,
561 }
562}
563
564fn instance_from_json(v: &serde_json::Value) -> Option<pb::InstanceInfo> {
565 Some(pb::InstanceInfo {
566 instance_id: v.get("instance_id")?.as_str()?.to_string(),
567 cluster_name: v.get("cluster_name")?.as_str()?.to_string(),
568 listen_addr: v.get("listen_addr")?.as_str()?.to_string(),
569 started_at_unix: v.get("started_at_unix")?.as_i64()?,
570 pid: v.get("pid")?.as_i64()? as i32,
571 config_sha256: v.get("config_sha256")?.as_str()?.to_string(),
572 is_config_writer: v.get("is_config_writer")?.as_bool()?,
573 })
574}
575
576#[tonic::async_trait]
577impl pb::lprovider_admin_server::LproviderAdmin for AdminGrpc {
578 async fn get_config(
579 &self,
580 request: Request<pb::GetConfigRequest>,
581 ) -> Result<Response<pb::GetConfigResponse>, Status> {
582 let req = request.into_inner();
583 let runtime = self.state.runtime.read().await;
584 let path = runtime.config_path.clone();
585 let sha = runtime.config_sha256.clone();
586 drop(runtime);
587 let content = tokio::fs::read_to_string(&path)
588 .await
589 .map_err(|e| status_from_err(&LproviderError::Internal(format!("read config: {e}"))))?;
590 let size = tokio::fs::metadata(&path)
591 .await
592 .map(|m| m.len() as i64)
593 .unwrap_or(0);
594 let masked = req.mask_secrets.unwrap_or(false);
595 let content_yaml = if masked {
596 mask_secrets(&content)
597 } else {
598 content
599 };
600 Ok(Response::new(pb::GetConfigResponse {
601 config_path: path.to_string_lossy().into_owned(),
602 sha256: sha,
603 file_size_bytes: size,
604 loaded_at_unix: self.state.started_at.unix_timestamp(),
605 content_yaml,
606 warnings: vec![],
607 }))
608 }
609
610 async fn update_config(
611 &self,
612 request: Request<pb::UpdateConfigRequest>,
613 ) -> Result<Response<pb::UpdateConfigResponse>, Status> {
614 let req = request.into_inner();
615 let operator = operator_of(&req.operator);
616 const RPC: &str = "UpdateConfig";
617 let _guard = match self.acquire_config_lock().await {
618 Ok(g) => g,
619 Err(e) => {
620 self.audit_write(
621 &operator,
622 RPC,
623 None,
624 None,
625 None,
626 "error",
627 Some(&e.to_string()),
628 )
629 .await;
630 return Err(status_from_err(&e));
631 }
632 };
633 let config_path = self.state.runtime.read().await.config_path.clone();
634 let new_yaml = match compose_yaml(&req, &config_path).await {
635 Ok(y) => y,
636 Err(e) => {
637 self.audit_write(
638 &operator,
639 RPC,
640 None,
641 None,
642 None,
643 "error",
644 Some(&e.to_string()),
645 )
646 .await;
647 return Err(status_from_err(&e));
648 }
649 };
650
651 let result = async {
652 let tmp = config_path.with_extension("tmp");
653 tokio::fs::write(&tmp, &new_yaml)
654 .await
655 .map_err(|e| LproviderError::Internal(format!("write tmp: {e}")))?;
656 let tmp_file = std::fs::File::open(&tmp)
657 .map_err(|e| LproviderError::Internal(format!("open tmp: {e}")))?;
658 tmp_file
659 .sync_all()
660 .map_err(|e| LproviderError::Internal(format!("fsync tmp: {e}")))?;
661 drop(tmp_file);
662 config::ConfigLoader::new()
663 .config_file(&tmp)
664 .load()
665 .map_err(|e| {
666 LproviderError::InvalidRequest(format!("config validation failed: {e}"))
667 })?;
668
669 let bak = config_path.with_extension("bak");
670 let _ = tokio::fs::copy(&config_path, &bak).await;
671 tokio::fs::rename(&tmp, &config_path)
672 .await
673 .map_err(|e| LproviderError::Internal(format!("rename: {e}")))?;
674
675 let mut cli = self.state.cli.clone();
676 cli.config = config_path.clone();
677 let new_rt = RuntimeConfig::load(&cli)?;
678 let sha256 = new_rt.config_sha256.clone();
679 let old_app = self.state.runtime.read().await.app.clone();
680 let warnings = if infra_changed(&old_app, &new_rt.app) {
681 vec!["infrastructure section changed (bind/port/redis/db/secret/context/timeout); restart required".to_string()]
682 } else {
683 vec![]
684 };
685
686 let mut reloaded = false;
687 let mut restarted = false;
688 if req.restart_after_write {
689 restarted = true;
690 let state = self.state.clone();
691 tokio::spawn(async move {
692 tokio::time::sleep(Duration::from_secs(30)).await;
693 state.shutdown.cancel();
694 });
695 } else {
696 reloaded = true;
697 self.manager.reload(new_rt.lp())?;
698 *self.state.runtime.write().await = new_rt;
699 }
700
701 Ok::<pb::UpdateConfigResponse, LproviderError>(pb::UpdateConfigResponse {
702 sha256,
703 backup_path: bak.to_string_lossy().into_owned(),
704 reloaded,
705 restarted,
706 warnings,
707 })
708 }
709 .await;
710
711 match result {
712 Ok(resp) => {
713 self.audit_write(
714 &operator,
715 RPC,
716 Some(
717 &self
718 .state
719 .runtime
720 .blocking_read()
721 .config_path
722 .to_string_lossy(),
723 ),
724 Some("config updated"),
725 Some("reload"),
726 "success",
727 None,
728 )
729 .await;
730 Ok(Response::new(resp))
731 }
732 Err(e) => {
733 self.audit_write(
734 &operator,
735 RPC,
736 None,
737 None,
738 None,
739 "error",
740 Some(&e.to_string()),
741 )
742 .await;
743 Err(status_from_err(&e))
744 }
745 }
746 }
747
748 async fn reload_config(
749 &self,
750 request: Request<pb::ReloadConfigRequest>,
751 ) -> Result<Response<pb::ReloadConfigResponse>, Status> {
752 let req = request.into_inner();
753 let operator = operator_of(&req.operator);
754 const RPC: &str = "ReloadConfig";
755 let previous = self.state.runtime.read().await.config_sha256.clone();
756 match self.reload_runtime().await {
757 Ok(()) => {
758 let current = self.state.runtime.read().await.config_sha256.clone();
759 let status = self.build_runtime_status().await;
760 self.audit_write(
761 &operator,
762 RPC,
763 None,
764 Some("config reloaded"),
765 Some("reload"),
766 "success",
767 None,
768 )
769 .await;
770 Ok(Response::new(pb::ReloadConfigResponse {
771 previous_sha256: previous,
772 current_sha256: current,
773 warnings: vec![],
774 status: Some(status),
775 }))
776 }
777 Err(e) => {
778 self.audit_write(
779 &operator,
780 RPC,
781 None,
782 None,
783 None,
784 "error",
785 Some(&e.to_string()),
786 )
787 .await;
788 Err(status_from_err(&e))
789 }
790 }
791 }
792
793 async fn restart_service(
794 &self,
795 request: Request<pb::RestartServiceRequest>,
796 ) -> Result<Response<pb::RestartServiceResponse>, Status> {
797 let req = request.into_inner();
798 let operator = operator_of(&req.operator);
799 const RPC: &str = "RestartService";
800 let drain = if req.drain_timeout_seconds <= 0 {
801 30
802 } else {
803 req.drain_timeout_seconds
804 } as u64;
805 let state = self.state.clone();
806 let instance_id = self.state.instance_id.clone();
807 let self_reexec = req.self_reexec;
808 tokio::spawn(async move {
809 tokio::time::sleep(Duration::from_secs(drain)).await;
810 if self_reexec && let Ok(exe) = std::env::current_exe() {
811 let args: Vec<String> = std::env::args().skip(1).collect();
812 let _ = std::process::Command::new(exe).args(args).spawn();
813 }
814 state.shutdown.cancel();
815 });
816 self.audit_write(
817 &operator,
818 RPC,
819 None,
820 Some("restart scheduled"),
821 Some("restart"),
822 "success",
823 None,
824 )
825 .await;
826 Ok(Response::new(pb::RestartServiceResponse {
827 accepted: true,
828 reason: format!("draining {}s then restart", drain),
829 instance_id,
830 }))
831 }
832
833 async fn list_upstreams(
834 &self,
835 request: Request<pb::ListUpstreamsRequest>,
836 ) -> Result<Response<pb::ListUpstreamsResponse>, Status> {
837 let req = request.into_inner();
838 let lp = self.state.runtime.blocking_read().lp().clone();
839 let upstreams = lp
840 .upstreams
841 .iter()
842 .filter(|u| req.codes.is_empty() || req.codes.contains(&u.code))
843 .map(upstream_to_proto)
844 .collect();
845 Ok(Response::new(pb::ListUpstreamsResponse { upstreams }))
846 }
847
848 async fn get_upstream(
849 &self,
850 request: Request<pb::GetUpstreamRequest>,
851 ) -> Result<Response<pb::Upstream>, Status> {
852 let req = request.into_inner();
853 let lp = self.state.runtime.blocking_read().lp().clone();
854 match lp.upstreams.iter().find(|u| u.code == req.code) {
855 Some(u) => Ok(Response::new(upstream_to_proto(u))),
856 None => Err(status_from_err(&LproviderError::NotFound(format!(
857 "upstream {}",
858 req.code
859 )))),
860 }
861 }
862
863 async fn list_upstream_keys(
864 &self,
865 request: Request<pb::ListUpstreamKeysRequest>,
866 ) -> Result<Response<pb::ListUpstreamKeysResponse>, Status> {
867 let req = request.into_inner();
868 let snapshot = self.manager.snapshot().await;
869 let lp = self.state.runtime.blocking_read().lp().clone();
870 let mut keys = Vec::new();
871 for up in &snapshot.upstreams {
872 if !req.upstream_code.is_empty() && up.code != req.upstream_code {
873 continue;
874 }
875 let cfg_up = lp.upstreams.iter().find(|u| u.code == up.code);
876 for k in &up.keys {
877 let cred = cfg_up.and_then(|u| u.credentials.iter().find(|c| c.label == k.label));
878 keys.push(
879 build_upstream_key(
880 up,
881 k,
882 cred.map(|c| c.credential_ref.clone()),
883 cred.map(|c| c.qps_limit).unwrap_or(0),
884 cred.map(|c| c.rpm_limit).unwrap_or(0),
885 )
886 .await,
887 );
888 }
889 }
890 Ok(Response::new(pb::ListUpstreamKeysResponse { keys }))
891 }
892
893 async fn get_upstream_key(
894 &self,
895 request: Request<pb::GetUpstreamKeyRequest>,
896 ) -> Result<Response<pb::UpstreamKey>, Status> {
897 let req = request.into_inner();
898 let snapshot = self.manager.snapshot().await;
899 let lp = self.state.runtime.blocking_read().lp().clone();
900 for up in &snapshot.upstreams {
901 if up.code != req.upstream_code {
902 continue;
903 }
904 if let Some(k) = up.keys.iter().find(|k| k.label == req.label) {
905 let cred = lp
906 .upstreams
907 .iter()
908 .find(|u| u.code == up.code)
909 .and_then(|u| u.credentials.iter().find(|c| c.label == k.label));
910 return Ok(Response::new(
911 build_upstream_key(
912 up,
913 k,
914 cred.map(|c| c.credential_ref.clone()),
915 cred.map(|c| c.qps_limit).unwrap_or(0),
916 cred.map(|c| c.rpm_limit).unwrap_or(0),
917 )
918 .await,
919 ));
920 }
921 }
922 Err(status_from_err(&LproviderError::NotFound(format!(
923 "key {}/{}",
924 req.upstream_code, req.label
925 ))))
926 }
927
928 async fn list_bindings(
929 &self,
930 request: Request<pb::ListBindingsRequest>,
931 ) -> Result<Response<pb::ListBindingsResponse>, Status> {
932 let req = request.into_inner();
933 let Some(redis) = &self.state.redis else {
934 return Ok(Response::new(pb::ListBindingsResponse { bindings: vec![] }));
935 };
936 let prefix = self.affinity.key_prefix().to_string();
937 let pattern = format!("{prefix}:bind:*");
938 let mut conn = redis.connection();
939 let mut raw_keys: Vec<String> = Vec::new();
940 {
941 let mut iter = conn.scan_match(pattern).await.map_err(redis_status)?;
942 while let Some(item) = iter.next().await {
943 if let Ok(k) = item {
944 raw_keys.push(k);
945 }
946 }
947 }
948 let mut bindings = Vec::new();
949 for key in raw_keys {
950 let Some(suffix) = key.strip_prefix(&format!("{prefix}:bind:")) else {
951 continue;
952 };
953 let mut it = suffix.splitn(3, ':');
954 let (Some(tenant_s), Some(alias), Some(skey)) = (it.next(), it.next(), it.next())
955 else {
956 continue;
957 };
958 let Ok(tenant) = tenant_s.parse::<i64>() else {
959 continue;
960 };
961 if req.tenant_id != 0 && tenant != req.tenant_id {
962 continue;
963 }
964 if !req.model_alias.is_empty() && alias != req.model_alias {
965 continue;
966 }
967 if !req.session_key.is_empty() && skey != req.session_key {
968 continue;
969 }
970 let val: Option<String> = conn.get(&key).await.ok().flatten();
971 let ttl: i64 = conn.ttl(&key).await.ok().unwrap_or(0);
972 if let Some(v) = val
973 && let Ok(b) = serde_json::from_str::<crate::affinity::SessionBinding>(&v)
974 {
975 bindings.push(pb::SessionBinding {
976 model_alias: alias.to_string(),
977 session_key: skey.to_string(),
978 upstream_code: b.upstream_code,
979 key_label: b.key_label,
980 prefix_hash: b.prefix_hash,
981 bound_at_unix: b.bound_at,
982 expires_at_unix: if ttl > 0 { b.bound_at + ttl } else { 0 },
983 last_hit_tokens: b.last_hit_tokens as i32,
984 hit_count: b.hit_count,
985 });
986 }
987 }
988 Ok(Response::new(pb::ListBindingsResponse { bindings }))
989 }
990
991 async fn invalidate_binding(
992 &self,
993 request: Request<pb::InvalidateBindingRequest>,
994 ) -> Result<Response<pb::InvalidateBindingResponse>, Status> {
995 let req = request.into_inner();
996 let operator = "unknown".to_string();
997 const RPC: &str = "InvalidateBinding";
998 match self
999 .affinity
1000 .invalidate(req.tenant_id, &req.model_alias, &req.session_key)
1001 .await
1002 {
1003 Ok(()) => {
1004 self.audit_write(
1005 &operator,
1006 RPC,
1007 Some(&format!("{}:{}", req.model_alias, req.session_key)),
1008 Some("binding invalidated"),
1009 None,
1010 "success",
1011 None,
1012 )
1013 .await;
1014 Ok(Response::new(pb::InvalidateBindingResponse {
1015 invalidated: true,
1016 }))
1017 }
1018 Err(e) => Err(status_from_err(&e)),
1019 }
1020 }
1021
1022 async fn query_audit(
1023 &self,
1024 request: Request<pb::QueryAuditRequest>,
1025 ) -> Result<Response<pb::QueryAuditResponse>, Status> {
1026 let req = request.into_inner();
1027 let st = (!req.subject_type.is_empty()).then(|| req.subject_type.clone());
1028 let ma = (!req.model_alias.is_empty()).then(|| req.model_alias.clone());
1029 let uc = (!req.upstream_code.is_empty()).then(|| req.upstream_code.clone());
1030 let status = (!req.status.is_empty()).then(|| req.status.clone());
1031 let limit = if req.limit <= 0 {
1032 100
1033 } else {
1034 req.limit as u64
1035 };
1036 let (models, total) = self
1037 .audit
1038 .query_audit(
1039 (req.tenant_id != 0).then_some(req.tenant_id),
1040 st.as_deref(),
1041 (req.subject_id != 0).then_some(req.subject_id),
1042 ma.as_deref(),
1043 uc.as_deref(),
1044 status.as_deref(),
1045 (req.from_ms > 0).then_some(req.from_ms),
1046 (req.to_ms > 0).then_some(req.to_ms),
1047 limit,
1048 req.offset as u64,
1049 )
1050 .await
1051 .map_err(|e| status_from_err(&e))?;
1052 Ok(Response::new(pb::QueryAuditResponse {
1053 records: models.iter().map(call_log_to_proto).collect(),
1054 total: total as i64,
1055 }))
1056 }
1057
1058 async fn aggregate_billing(
1059 &self,
1060 request: Request<pb::AggregateBillingRequest>,
1061 ) -> Result<Response<pb::AggregateBillingResponse>, Status> {
1062 let req = request.into_inner();
1063 let st = (!req.subject_type.is_empty()).then(|| req.subject_type.clone());
1064 let pk = (!req.period_kind.is_empty()).then(|| req.period_kind.clone());
1065 let ma = (!req.model_alias.is_empty()).then(|| req.model_alias.clone());
1066 let uc = (!req.upstream_code.is_empty()).then(|| req.upstream_code.clone());
1067 let (models, _total) = self
1068 .audit
1069 .aggregate_billing(
1070 (req.tenant_id != 0).then_some(req.tenant_id),
1071 st.as_deref(),
1072 (req.subject_id != 0).then_some(req.subject_id),
1073 pk.as_deref(),
1074 (req.period_start_ms > 0).then_some(req.period_start_ms),
1075 (req.period_end_ms > 0).then_some(req.period_end_ms),
1076 ma.as_deref(),
1077 uc.as_deref(),
1078 500,
1079 0,
1080 )
1081 .await
1082 .map_err(|e| status_from_err(&e))?;
1083 Ok(Response::new(pb::AggregateBillingResponse {
1084 aggregates: models.iter().map(aggregate_to_proto).collect(),
1085 }))
1086 }
1087
1088 async fn list_tenant_quotas(
1089 &self,
1090 request: Request<pb::ListTenantQuotasRequest>,
1091 ) -> Result<Response<pb::ListTenantQuotasResponse>, Status> {
1092 let req = request.into_inner();
1093 let (quotas, _total) = self
1094 .quotas
1095 .list_quotas(
1096 (req.tenant_id != 0).then_some(req.tenant_id),
1097 (!req.subject_type.is_empty()).then(|| req.subject_type.clone()),
1098 (req.subject_id != 0).then_some(req.subject_id),
1099 100,
1100 0,
1101 )
1102 .await
1103 .map_err(|e| status_from_err(&e))?;
1104 Ok(Response::new(pb::ListTenantQuotasResponse {
1105 quotas: quotas.iter().map(quota_to_proto).collect(),
1106 }))
1107 }
1108
1109 async fn upsert_tenant_quota(
1110 &self,
1111 request: Request<pb::UpsertTenantQuotaRequest>,
1112 ) -> Result<Response<pb::TenantQuota>, Status> {
1113 let req = request.into_inner();
1114 let operator = "unknown".to_string();
1115 const RPC: &str = "UpsertTenantQuota";
1116 let Some(q) = &req.quota else {
1117 return Err(status_from_err(&LproviderError::InvalidRequest(
1118 "quota missing".into(),
1119 )));
1120 };
1121 let quota = proto_to_quota(q);
1122 match self.quotas.upsert_quota("a).await {
1123 Ok(()) => {
1124 self.audit_write(
1125 &operator,
1126 RPC,
1127 Some(&format!(
1128 "{}:{}:{}",
1129 quota.tenant_id, quota.subject_type, quota.subject_id
1130 )),
1131 Some("quota upserted"),
1132 None,
1133 "success",
1134 None,
1135 )
1136 .await;
1137 Ok(Response::new(quota_to_proto("a)))
1138 }
1139 Err(e) => Err(status_from_err(&e)),
1140 }
1141 }
1142
1143 async fn get_wallet(
1144 &self,
1145 request: Request<pb::GetWalletRequest>,
1146 ) -> Result<Response<pb::Wallet>, Status> {
1147 let req = request.into_inner();
1148 let wallet = self
1149 .wallets
1150 .get_wallet(req.wallet_id)
1151 .await
1152 .map_err(|e| status_from_err(&e))?;
1153 Ok(Response::new(wallet_to_proto(&wallet)))
1154 }
1155
1156 async fn list_wallets(
1157 &self,
1158 request: Request<pb::ListWalletsRequest>,
1159 ) -> Result<Response<pb::ListWalletsResponse>, Status> {
1160 let req = request.into_inner();
1161 let (wallets, _total) = self
1162 .wallets
1163 .list_wallets(
1164 (req.tenant_id != 0).then_some(req.tenant_id),
1165 (!req.owner_type.is_empty()).then(|| req.owner_type.clone()),
1166 (req.owner_id != 0).then_some(req.owner_id),
1167 None,
1168 100,
1169 0,
1170 )
1171 .await
1172 .map_err(|e| status_from_err(&e))?;
1173 Ok(Response::new(pb::ListWalletsResponse {
1174 wallets: wallets.iter().map(wallet_to_proto).collect(),
1175 }))
1176 }
1177
1178 async fn topup_wallet(
1179 &self,
1180 request: Request<pb::TopupWalletRequest>,
1181 ) -> Result<Response<pb::Wallet>, Status> {
1182 let req = request.into_inner();
1183 let operator = operator_of(&req.operator);
1184 const RPC: &str = "TopupWallet";
1185 if req.idempotency_key.trim().is_empty() {
1186 return Err(status_from_err(&LproviderError::InvalidRequest(
1187 "idempotency_key is required for topup".into(),
1188 )));
1189 }
1190 if req.delta_points <= 0 {
1191 return Err(status_from_err(&LproviderError::InvalidRequest(
1192 "delta_points must be positive".into(),
1193 )));
1194 }
1195 let reason = (!req.reason.is_empty()).then(|| req.reason.clone());
1196 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1197 match self
1198 .wallets
1199 .topup(
1200 req.wallet_id,
1201 req.delta_points,
1202 &req.idempotency_key,
1203 reason,
1204 op,
1205 )
1206 .await
1207 {
1208 Ok(_) => {
1209 let wallet = self
1210 .wallets
1211 .get_wallet(req.wallet_id)
1212 .await
1213 .map_err(|e| status_from_err(&e))?;
1214 self.audit_write(
1215 &operator,
1216 RPC,
1217 Some(&format!("wallet {}", req.wallet_id)),
1218 Some(&format!("+{} points", req.delta_points)),
1219 None,
1220 "success",
1221 None,
1222 )
1223 .await;
1224 Ok(Response::new(wallet_to_proto(&wallet)))
1225 }
1226 Err(e) => Err(status_from_err(&e)),
1227 }
1228 }
1229
1230 async fn refund_wallet(
1231 &self,
1232 request: Request<pb::RefundWalletRequest>,
1233 ) -> Result<Response<pb::Wallet>, Status> {
1234 let req = request.into_inner();
1235 let operator = operator_of(&req.operator);
1236 const RPC: &str = "RefundWallet";
1237 let request_id = match uuid::Uuid::parse_str(&req.request_id) {
1238 Ok(id) => id,
1239 Err(_) => {
1240 return Err(status_from_err(&LproviderError::InvalidRequest(
1241 "request_id is required and must be a valid uuid".into(),
1242 )));
1243 }
1244 };
1245 if req.delta_points <= 0 {
1246 return Err(status_from_err(&LproviderError::InvalidRequest(
1247 "delta_points must be positive".into(),
1248 )));
1249 }
1250 let reason = (!req.reason.is_empty()).then(|| req.reason.clone());
1251 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1252 match self
1253 .wallets
1254 .refund(req.wallet_id, req.delta_points, request_id, reason, op)
1255 .await
1256 {
1257 Ok(_) => {
1258 let wallet = self
1259 .wallets
1260 .get_wallet(req.wallet_id)
1261 .await
1262 .map_err(|e| status_from_err(&e))?;
1263 self.audit_write(
1264 &operator,
1265 RPC,
1266 Some(&format!("wallet {}", req.wallet_id)),
1267 Some(&format!("refund {} points", req.delta_points)),
1268 None,
1269 "success",
1270 None,
1271 )
1272 .await;
1273 Ok(Response::new(wallet_to_proto(&wallet)))
1274 }
1275 Err(e) => Err(status_from_err(&e)),
1276 }
1277 }
1278
1279 async fn adjust_wallet(
1280 &self,
1281 request: Request<pb::AdjustWalletRequest>,
1282 ) -> Result<Response<pb::Wallet>, Status> {
1283 let req = request.into_inner();
1284 let operator = operator_of(&req.operator);
1285 const RPC: &str = "AdjustWallet";
1286 if req.idempotency_key.trim().is_empty() {
1287 return Err(status_from_err(&LproviderError::InvalidRequest(
1288 "idempotency_key is required for adjust".into(),
1289 )));
1290 }
1291 if req.delta_points == 0 {
1292 return Err(status_from_err(&LproviderError::InvalidRequest(
1293 "delta_points must not be zero".into(),
1294 )));
1295 }
1296 let reason = (!req.reason.is_empty()).then(|| req.reason.clone());
1297 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1298 match self
1299 .wallets
1300 .adjust(
1301 req.wallet_id,
1302 req.delta_points,
1303 &req.idempotency_key,
1304 reason,
1305 op,
1306 )
1307 .await
1308 {
1309 Ok(_) => {
1310 let wallet = self
1311 .wallets
1312 .get_wallet(req.wallet_id)
1313 .await
1314 .map_err(|e| status_from_err(&e))?;
1315 self.audit_write(
1316 &operator,
1317 RPC,
1318 Some(&format!("wallet {}", req.wallet_id)),
1319 Some(&format!("adjust {} points", req.delta_points)),
1320 None,
1321 "success",
1322 None,
1323 )
1324 .await;
1325 Ok(Response::new(wallet_to_proto(&wallet)))
1326 }
1327 Err(e) => Err(status_from_err(&e)),
1328 }
1329 }
1330
1331 async fn freeze_wallet(
1332 &self,
1333 request: Request<pb::FreezeWalletRequest>,
1334 ) -> Result<Response<pb::Wallet>, Status> {
1335 let req = request.into_inner();
1336 let operator = operator_of(&req.operator);
1337 const RPC: &str = "FreezeWallet";
1338 let reason = if req.reason.is_empty() {
1339 "admin action".to_string()
1340 } else {
1341 req.reason.clone()
1342 };
1343 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1344 match self.wallets.freeze(req.wallet_id, reason, op).await {
1345 Ok(()) => {
1346 let wallet = self
1347 .wallets
1348 .get_wallet(req.wallet_id)
1349 .await
1350 .map_err(|e| status_from_err(&e))?;
1351 self.audit_write(
1352 &operator,
1353 RPC,
1354 Some(&format!("wallet {}", req.wallet_id)),
1355 None,
1356 None,
1357 "success",
1358 None,
1359 )
1360 .await;
1361 Ok(Response::new(wallet_to_proto(&wallet)))
1362 }
1363 Err(e) => Err(status_from_err(&e)),
1364 }
1365 }
1366
1367 async fn unfreeze_wallet(
1368 &self,
1369 request: Request<pb::UnfreezeWalletRequest>,
1370 ) -> Result<Response<pb::Wallet>, Status> {
1371 let req = request.into_inner();
1372 let operator = operator_of(&req.operator);
1373 const RPC: &str = "UnfreezeWallet";
1374 let reason = if req.reason.is_empty() {
1375 "admin action".to_string()
1376 } else {
1377 req.reason.clone()
1378 };
1379 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1380 match self.wallets.unfreeze(req.wallet_id, reason, op).await {
1381 Ok(()) => {
1382 let wallet = self
1383 .wallets
1384 .get_wallet(req.wallet_id)
1385 .await
1386 .map_err(|e| status_from_err(&e))?;
1387 self.audit_write(
1388 &operator,
1389 RPC,
1390 Some(&format!("wallet {}", req.wallet_id)),
1391 None,
1392 None,
1393 "success",
1394 None,
1395 )
1396 .await;
1397 Ok(Response::new(wallet_to_proto(&wallet)))
1398 }
1399 Err(e) => Err(status_from_err(&e)),
1400 }
1401 }
1402
1403 async fn disable_wallet(
1404 &self,
1405 request: Request<pb::DisableWalletRequest>,
1406 ) -> Result<Response<pb::Wallet>, Status> {
1407 let req = request.into_inner();
1408 let operator = operator_of(&req.operator);
1409 const RPC: &str = "DisableWallet";
1410 let reason = if req.reason.is_empty() {
1411 "admin action".to_string()
1412 } else {
1413 req.reason.clone()
1414 };
1415 let op = (!req.operator.is_empty()).then(|| req.operator.clone());
1416 match self.wallets.disable(req.wallet_id, reason, op).await {
1417 Ok(()) => {
1418 let wallet = self
1419 .wallets
1420 .get_wallet(req.wallet_id)
1421 .await
1422 .map_err(|e| status_from_err(&e))?;
1423 self.audit_write(
1424 &operator,
1425 RPC,
1426 Some(&format!("wallet {}", req.wallet_id)),
1427 None,
1428 None,
1429 "success",
1430 None,
1431 )
1432 .await;
1433 Ok(Response::new(wallet_to_proto(&wallet)))
1434 }
1435 Err(e) => Err(status_from_err(&e)),
1436 }
1437 }
1438
1439 async fn list_wallet_ledger(
1440 &self,
1441 request: Request<pb::ListWalletLedgerRequest>,
1442 ) -> Result<Response<pb::ListWalletLedgerResponse>, Status> {
1443 let req = request.into_inner();
1444 let (entries, _total) = self
1445 .wallets
1446 .list_ledger((req.wallet_id != 0).then_some(req.wallet_id), None, 100, 0)
1447 .await
1448 .map_err(|e| status_from_err(&e))?;
1449 let from = req.from_unix;
1450 let to = req.to_unix;
1451 let entries = entries
1452 .into_iter()
1453 .filter(|e| {
1454 let t = e.created_at.unix_timestamp();
1455 (from <= 0 || t >= from) && (to <= 0 || t <= to)
1456 })
1457 .map(|e| pb::WalletLedgerEntry {
1458 id: 0,
1459 wallet_id: e.wallet_id,
1460 request_id: e.request_id.map(|u| u.to_string()).unwrap_or_default(),
1461 delta_points: e.delta_points,
1462 balance_after: e.balance_after,
1463 kind: e.kind,
1464 reason: e.reason.unwrap_or_default(),
1465 operator: e.operator.unwrap_or_default(),
1466 created_at_unix: e.created_at.unix_timestamp(),
1467 })
1468 .collect();
1469 Ok(Response::new(pb::ListWalletLedgerResponse { entries }))
1470 }
1471
1472 async fn list_wallet_freeze_log(
1473 &self,
1474 request: Request<pb::ListWalletFreezeLogRequest>,
1475 ) -> Result<Response<pb::ListWalletFreezeLogResponse>, Status> {
1476 let req = request.into_inner();
1477 let (models, _total) = self
1478 .wallets
1479 .list_freeze_log((req.wallet_id != 0).then_some(req.wallet_id), 100, 0)
1480 .await
1481 .map_err(|e| status_from_err(&e))?;
1482 let entries = models
1483 .into_iter()
1484 .map(|m| pb::WalletFreezeLogEntry {
1485 id: m.id,
1486 wallet_id: m.wallet_id,
1487 action: m.action,
1488 reason: m.reason,
1489 operator: m.operator.unwrap_or_default(),
1490 created_at_unix: m.created_at.unix_timestamp(),
1491 })
1492 .collect();
1493 Ok(Response::new(pb::ListWalletFreezeLogResponse { entries }))
1494 }
1495
1496 async fn get_runtime_status(
1497 &self,
1498 _request: Request<pb::GetRuntimeStatusRequest>,
1499 ) -> Result<Response<pb::RuntimeStatus>, Status> {
1500 Ok(Response::new(self.build_runtime_status().await))
1501 }
1502
1503 async fn list_instances(
1504 &self,
1505 request: Request<pb::ListInstancesRequest>,
1506 ) -> Result<Response<pb::ListInstancesResponse>, Status> {
1507 let req = request.into_inner();
1508 let mut instances = vec![self_instance_info(&self.state)];
1509 if let Some(redis) = &self.state.redis {
1510 let cluster = if req.cluster_name.is_empty() {
1511 self.state.cluster_name.clone()
1512 } else {
1513 req.cluster_name.clone()
1514 };
1515 let key = format!("{}:instances:{}", self.affinity.key_prefix(), cluster);
1516 let mut conn = redis.connection();
1517 let members: redis::RedisResult<Vec<String>> = conn.smembers(&key).await;
1518 if let Ok(members) = members {
1519 for m in members {
1520 if let Ok(v) = serde_json::from_str::<serde_json::Value>(&m)
1521 && let Some(info) = instance_from_json(&v)
1522 && !instances.iter().any(|i| i.instance_id == info.instance_id)
1523 {
1524 instances.push(info);
1525 }
1526 }
1527 }
1528 }
1529 Ok(Response::new(pb::ListInstancesResponse { instances }))
1530 }
1531}
Added agent/lprovider/affinity.rs agent/lprovider/affinity.rs +312 −0 View agent/lprovider/affinity.rs
@@ -0,0 +1,312 @@
1use std::sync::Arc;
2
3use cache::RedisCache;
4use config::lprovider::LproviderRedisSettings;
5use llm::provider::LlmProvider;
6use llm::types::{ChatMessage, ChatRequest, MessageRole};
7use redis::AsyncCommands;
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10use tokio_stream::StreamExt;
11
12use crate::error::LproviderError;
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct SessionBinding {
16 pub upstream_code: String,
17 pub key_label: String,
18 pub model: String,
19 pub prefix_hash: String,
20 pub bound_at: i64,
21 pub last_hit_tokens: i64,
22 pub hit_count: i64,
23 #[serde(default)]
24 pub miss_count: i64,
25}
26
27pub struct AffinityManager {
28 redis: Option<Arc<RedisCache>>,
29 key_prefix: String,
30 ttl_seconds: i64,
31 min_prefix_tokens: usize,
32 prefix_token_budget: usize,
33 no_hit_max_age_seconds: i64,
34}
35
36const BIND_SCRIPT: &str = r#"
37redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2]))
38redis.call('SADD', KEYS[2], KEYS[1])
39redis.call('SADD', KEYS[3], KEYS[1])
40return 1
41"#;
42
43const UNBIND_SCRIPT: &str = r#"
44redis.call('DEL', KEYS[1])
45redis.call('SREM', KEYS[2], KEYS[1])
46redis.call('SREM', KEYS[3], KEYS[1])
47return 1
48"#;
49
50impl AffinityManager {
51 pub fn new(redis: Option<Arc<RedisCache>>, cfg: &LproviderRedisSettings) -> Self {
52 Self {
53 redis,
54 key_prefix: cfg.key_prefix.clone(),
55 ttl_seconds: cfg.session_affinity_ttl_seconds,
56 min_prefix_tokens: cfg.min_prefix_tokens,
57 prefix_token_budget: cfg.prefix_token_budget,
58 no_hit_max_age_seconds: cfg.session_affinity_ttl_seconds,
59 }
60 }
61
62 pub fn key_prefix(&self) -> &str {
63 &self.key_prefix
64 }
65
66 pub fn ttl_seconds(&self) -> i64 {
67 self.ttl_seconds
68 }
69
70 pub fn bind_key(&self, tenant_id: i64, model_alias: &str, session_key: &str) -> String {
71 format!(
72 "{}:bind:{}:{}:{}",
73 self.key_prefix, tenant_id, model_alias, session_key
74 )
75 }
76
77 fn idx_upstream_key(&self, code: &str) -> String {
78 format!("{}:bind_idx:upstream:{}", self.key_prefix, code)
79 }
80
81 fn idx_key_key(&self, code: &str, label: &str) -> String {
82 format!("{}:bind_idx:key:{}:{}", self.key_prefix, code, label)
83 }
84
85 pub fn compute_session_key(
86 &self,
87 model_alias: &str,
88 messages: &[ChatMessage],
89 estimator: &Arc<dyn LlmProvider>,
90 ) -> Option<(String, usize)> {
91 let mut prefix: Vec<ChatMessage> = Vec::new();
92 for m in messages {
93 match m.role {
94 MessageRole::System | MessageRole::User => prefix.push(m.clone()),
95 _ => break,
96 }
97 }
98 while prefix.len() > 1
99 && estimator.count_tokens(&ChatRequest::new(prefix.clone())) > self.prefix_token_budget
100 {
101 prefix.pop();
102 }
103 let tokens = estimator.count_tokens(&ChatRequest::new(prefix.clone()));
104 if tokens < self.min_prefix_tokens {
105 return None;
106 }
107 let text = prefix
108 .iter()
109 .map(|m| m.text_content())
110 .collect::<Vec<_>>()
111 .join("\n");
112 let mut hasher = Sha256::new();
113 hasher.update(model_alias.as_bytes());
114 hasher.update(b"\n");
115 hasher.update(text.as_bytes());
116 Some((hex::encode(hasher.finalize()), tokens))
117 }
118
119 pub fn validate_session_key_override(&self, key: &str) -> Result<(), LproviderError> {
120 if key.is_empty() || key.len() > 128 {
121 return Err(LproviderError::InvalidRequest(
122 "X-Session-Key 长度须在 1..=128 之间".into(),
123 ));
124 }
125 if !key
126 .chars()
127 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-'))
128 {
129 return Err(LproviderError::InvalidRequest(
130 "X-Session-Key 仅允许 [A-Za-z0-9._-] 字符".into(),
131 ));
132 }
133 Ok(())
134 }
135
136 pub async fn lookup(
137 &self,
138 tenant_id: i64,
139 model_alias: &str,
140 session_key: &str,
141 ) -> Result<Option<SessionBinding>, LproviderError> {
142 let Some(cache) = &self.redis else {
143 return Ok(None);
144 };
145 let mut conn = cache.connection();
146 let key = self.bind_key(tenant_id, model_alias, session_key);
147 let raw: Option<String> = conn
148 .get(&key)
149 .await
150 .map_err(|e| LproviderError::Redis(e.to_string()))?;
151 match raw {
152 Some(raw) => serde_json::from_str(&raw)
153 .map(Some)
154 .map_err(|e| LproviderError::Internal(format!("绑定 JSON 解析失败: {e}"))),
155 None => Ok(None),
156 }
157 }
158
159 pub async fn bind(
160 &self,
161 tenant_id: i64,
162 model_alias: &str,
163 session_key: &str,
164 binding: &SessionBinding,
165 ) -> Result<(), LproviderError> {
166 let Some(cache) = &self.redis else {
167 return Ok(());
168 };
169 let mut conn = cache.connection();
170 let json = serde_json::to_string(binding)
171 .map_err(|e| LproviderError::Internal(format!("绑定 JSON 序列化失败: {e}")))?;
172 redis::Script::new(BIND_SCRIPT)
173 .key(self.bind_key(tenant_id, model_alias, session_key))
174 .key(self.idx_upstream_key(&binding.upstream_code))
175 .key(self.idx_key_key(&binding.upstream_code, &binding.key_label))
176 .arg(json)
177 .arg(self.ttl_seconds)
178 .invoke_async::<()>(&mut conn)
179 .await
180 .map_err(|e| LproviderError::Redis(e.to_string()))?;
181 Ok(())
182 }
183
184 pub async fn record_feedback(
185 &self,
186 tenant_id: i64,
187 model_alias: &str,
188 session_key: &str,
189 hit: bool,
190 hit_tokens: i64,
191 ) -> Result<(), LproviderError> {
192 let Some(cache) = &self.redis else {
193 return Ok(());
194 };
195 let mut conn = cache.connection();
196 let key = self.bind_key(tenant_id, model_alias, session_key);
197 let raw: Option<String> = conn
198 .get(&key)
199 .await
200 .map_err(|e| LproviderError::Redis(e.to_string()))?;
201 let Some(raw) = raw else {
202 return Ok(());
203 };
204 let mut b: SessionBinding = serde_json::from_str(&raw)
205 .map_err(|e| LproviderError::Internal(format!("绑定 JSON 解析失败: {e}")))?;
206 if hit {
207 b.hit_count += 1;
208 b.last_hit_tokens = hit_tokens;
209 } else {
210 b.miss_count += 1;
211 }
212 let total = b.hit_count + b.miss_count;
213 let ratio = if total > 0 {
214 b.hit_count as f64 / total as f64
215 } else {
216 0.0
217 };
218 let new_ttl = if hit && ratio >= 0.5 {
219 self.ttl_seconds
220 } else if b.hit_count == 0 && b.miss_count > self.min_prefix_tokens as i64 {
221 (self.ttl_seconds / 4).max(1)
222 } else {
223 let rem: i64 = conn
224 .ttl(&key)
225 .await
226 .map_err(|e| LproviderError::Redis(e.to_string()))?;
227 if rem > 0 { rem } else { self.ttl_seconds }
228 };
229 let json = serde_json::to_string(&b)
230 .map_err(|e| LproviderError::Internal(format!("绑定 JSON 序列化失败: {e}")))?;
231 let _: () = conn
232 .set_ex(&key, json, new_ttl as u64)
233 .await
234 .map_err(|e| LproviderError::Redis(e.to_string()))?;
235 Ok(())
236 }
237
238 pub async fn invalidate(
239 &self,
240 tenant_id: i64,
241 model_alias: &str,
242 session_key: &str,
243 ) -> Result<(), LproviderError> {
244 let Some(cache) = &self.redis else {
245 return Ok(());
246 };
247 let mut conn = cache.connection();
248 let key = self.bind_key(tenant_id, model_alias, session_key);
249 let raw: Option<String> = conn
250 .get(&key)
251 .await
252 .map_err(|e| LproviderError::Redis(e.to_string()))?;
253 let Some(raw) = raw else {
254 return Ok(());
255 };
256 let b: SessionBinding = serde_json::from_str(&raw)
257 .map_err(|e| LproviderError::Internal(format!("绑定 JSON 解析失败: {e}")))?;
258 redis::Script::new(UNBIND_SCRIPT)
259 .key(&key)
260 .key(self.idx_upstream_key(&b.upstream_code))
261 .key(self.idx_key_key(&b.upstream_code, &b.key_label))
262 .invoke_async::<()>(&mut conn)
263 .await
264 .map_err(|e| LproviderError::Redis(e.to_string()))?;
265 Ok(())
266 }
267
268 pub async fn cleanup_once(&self) -> Result<usize, LproviderError> {
269 let Some(cache) = &self.redis else {
270 return Ok(0);
271 };
272 let mut conn = cache.connection();
273 let pattern = format!("{}:bind:*", self.key_prefix);
274 let mut keys: Vec<String> = Vec::new();
275 {
276 let mut iter = conn
277 .scan_match(pattern)
278 .await
279 .map_err(|e| LproviderError::Redis(e.to_string()))?;
280 loop {
281 match iter.next().await {
282 Some(Ok(k)) => keys.push(k),
283 Some(Err(e)) => return Err(LproviderError::Redis(e.to_string())),
284 None => break,
285 }
286 }
287 }
288 let now = time::OffsetDateTime::now_utc().unix_timestamp();
289 let mut removed = 0usize;
290 for key in keys {
291 let raw: Option<String> = conn
292 .get(&key)
293 .await
294 .map_err(|e| LproviderError::Redis(e.to_string()))?;
295 let Some(raw) = raw else { continue };
296 let Ok(b) = serde_json::from_str::<SessionBinding>(&raw) else {
297 continue;
298 };
299 if b.hit_count == 0 && now - b.bound_at > self.no_hit_max_age_seconds {
300 redis::Script::new(UNBIND_SCRIPT)
301 .key(&key)
302 .key(self.idx_upstream_key(&b.upstream_code))
303 .key(self.idx_key_key(&b.upstream_code, &b.key_label))
304 .invoke_async::<()>(&mut conn)
305 .await
306 .map_err(|e| LproviderError::Redis(e.to_string()))?;
307 removed += 1;
308 }
309 }
310 Ok(removed)
311 }
312}
Added agent/lprovider/ai.proto agent/lprovider/ai.proto +179 −0 View agent/lprovider/ai.proto
@@ -0,0 +1,179 @@
1syntax = "proto3";
2package lprovider.ai.v1;
3
4import "google/protobuf/struct.proto";
5
6service LproviderAI {
7 rpc Chat(ChatRequestMsg) returns (ChatResponseMsg);
8
9 rpc ChatStream(ChatRequestMsg) returns (stream StreamEventMsg);
10
11 // request_id 语义：与整个流绑定（取流中第一条消息的 request_id），
12 // 同一流内后续消息不得变更 request_id；Cancel 按该 request_id 终止整个流。
13 rpc ChatStreamBidirectional(stream ChatRequestMsg) returns (stream StreamEventMsg);
14
15 rpc CountTokens(CountTokensRequest) returns (CountTokensResponse);
16
17 rpc Cancel(CancelRequest) returns (CancelResponse);
18}
19
20message ChatRequestMsg {
21 string model_alias = 1;
22 repeated ChatMessageMsg messages = 2;
23 repeated ToolMsg tools = 3;
24 optional int32 max_output_tokens = 4;
25 optional double temperature = 5;
26 optional double top_p = 6;
27 optional ResponseFormatMsg response_format = 7;
28 repeated string stop_sequences = 8;
29 google.protobuf.Struct provider_options = 9;
30
31 fixed64 tenant_id = 20;
32 int64 api_key_id = 21;
33 string caller_service = 22;
34 string request_id = 23;
35 int32 priority = 24;
36 string trace_id = 25;
37 map<string, string> labels = 26;
38 WalletSelector wallet_selector = 27;
39 ModelSelection model_selection = 28;
40 string session_key_override = 30;
41}
42
43message ModelSelection {
44 enum Strategy {
45 STRATEGY_UNSPECIFIED = 0;
46 BALANCED = 1;
47 COST = 2;
48 QUALITY = 3;
49 LATENCY = 4;
50 EXPLICIT = 5;
51 }
52 Strategy strategy = 1;
53 repeated string allowed_upstreams = 2;
54 repeated string excluded_upstreams = 3;
55 int32 balance_bias = 4;
56 repeated string required_capabilities = 5;
57}
58
59message WalletSelector {
60 oneof mode {
61 WalletRef wallet = 1;
62 AnonymousRef anonymous = 2;
63 }
64}
65
66message WalletRef {
67 int64 wallet_id = 1;
68}
69
70message AnonymousRef {}
71
72message ChatMessageMsg {
73 string role = 1;
74 oneof content {
75 string text = 2;
76 ContentPartsMsg parts = 3;
77 }
78 string tool_call_id = 4;
79 repeated ToolCallMsg tool_calls = 5;
80 string name = 6;
81 string reasoning_text = 7;
82}
83
84message ContentPartsMsg {
85 message Part {
86 string type = 1;
87 string text = 2;
88 string url = 3;
89 bytes data = 4;
90 string mime_type = 5;
91 }
92 repeated Part parts = 1;
93}
94
95message ToolMsg {
96 string tool_type = 1;
97 google.protobuf.Struct function = 2;
98}
99
100message ToolCallMsg {
101 string id = 1;
102 string tool_type = 2;
103 google.protobuf.Struct function = 3;
104}
105
106message ResponseFormatMsg {
107 string type = 1;
108 string json_schema = 2;
109}
110
111message ChatResponseMsg {
112 string id = 1;
113 string model_alias = 2;
114 string upstream_code = 3;
115 string key_label = 4;
116 string content = 5;
117 string reasoning_text = 6;
118 repeated ToolCallMsg tool_calls = 7;
119 string finish_reason = 8;
120 UsageMsg usage = 9;
121 google.protobuf.Struct raw_provider_response = 10;
122 int64 latency_ms = 11;
123 string request_id = 12;
124}
125
126message StreamEventMsg {
127 string event_type = 1;
128 string text_delta = 2;
129 string reasoning_delta = 3;
130 ToolCallDeltaMsg tool_call_delta = 4;
131 UsageMsg usage = 5;
132 string finish_reason = 6;
133 string error_kind = 7;
134 string error_message = 8;
135 string request_id = 9;
136 string session_key = 10;
137}
138
139message ToolCallDeltaMsg {
140 string id = 1;
141 int32 index = 2;
142 string tool_type = 3;
143 string name_delta = 4;
144 string arguments_delta = 5;
145}
146
147message UsageMsg {
148 int32 prompt_tokens = 1;
149 int32 completion_tokens = 2;
150 int32 total_tokens = 3;
151 int32 cache_hit_tokens = 4;
152 int32 cache_miss_tokens = 5;
153 int64 points_input = 6;
154 int64 points_output = 7;
155 int64 points_cache_hit = 8;
156 int64 points_total = 9;
157}
158
159message CountTokensRequest {
160 string model_alias = 1;
161 repeated ChatMessageMsg messages = 2;
162 repeated ToolMsg tools = 3;
163}
164
165message CountTokensResponse {
166 int32 estimated_input_tokens = 1;
167 int32 estimated_output_tokens = 2;
168 int32 max_context_tokens = 3;
169 int32 max_output_tokens = 4;
170 bool fits_context = 5;
171}
172
173message CancelRequest {
174 string request_id = 1;
175}
176
177message CancelResponse {
178 bool cancelled = 1;
179}
Added agent/lprovider/ai_grpc.rs agent/lprovider/ai_grpc.rs +658 −0 View agent/lprovider/ai_grpc.rs
@@ -0,0 +1,658 @@
1use std::collections::{BTreeMap, HashMap};
2use std::sync::{Arc, Mutex};
3
4use futures::StreamExt;
5use llm::types::{
6 ChatMessage, ChatRequest, ContentPart, FinishReason, ImageDetail, MessageContent, MessageRole,
7 ResponseFormat, StreamEvent, ToolCall, ToolDefinition,
8};
9use tokio::sync::mpsc;
10use tokio_stream::wrappers::ReceiverStream;
11use tonic::{Request, Response, Status};
12
13use crate::error::{LproviderError, WalletError};
14use crate::lprovider::ai::v1::{self, lprovider_ai_server::LproviderAi};
15use crate::router::{ChatContext, ChatSelectionSpec, Router};
16use crate::state::AppState;
17use crate::wallet::WalletService;
18
19#[derive(Clone)]
20pub struct AiGrpc {
21 pub router: Arc<Router>,
22 pub state: Arc<AppState>,
23 pub wallets: Arc<WalletService>,
24 cancels: Arc<Mutex<HashMap<String, mpsc::Sender<()>>>>,
25}
26
27impl AiGrpc {
28 pub fn new(router: Arc<Router>, state: Arc<AppState>) -> Self {
29 let wallets = Arc::new(WalletService::new(state.db.clone().map(Arc::new)));
30 Self {
31 router,
32 state,
33 wallets,
34 cancels: Arc::new(Mutex::new(HashMap::new())),
35 }
36 }
37
38 fn register_cancel(&self, request_id: &str) -> mpsc::Receiver<()> {
39 let (tx, rx) = mpsc::channel::<()>(1);
40 self.cancels
41 .lock()
42 .unwrap()
43 .insert(request_id.to_string(), tx);
44 rx
45 }
46
47 fn tenant_id(msg: &v1::ChatRequestMsg) -> Result<i64, Status> {
48 if msg.tenant_id == 0 {
49 return Err(status_from_err(&LproviderError::InvalidRequest(
50 "tenant_id is required".into(),
51 )));
52 }
53 Ok(msg.tenant_id as i64)
54 }
55
56 fn request_id(msg: &v1::ChatRequestMsg) -> Result<uuid::Uuid, Status> {
57 if msg.request_id.is_empty() {
58 return Ok(uuid::Uuid::now_v7());
59 }
60 uuid::Uuid::parse_str(&msg.request_id).map_err(|_| {
61 status_from_err(&LproviderError::InvalidRequest(
62 "request_id is not a valid UUID".into(),
63 ))
64 })
65 }
66
67 async fn build_context(
68 &self,
69 msg: v1::ChatRequestMsg,
70 stream: bool,
71 ) -> Result<ChatContext, Status> {
72 let tenant_id = Self::tenant_id(&msg)?;
73 let request_id = Self::request_id(&msg)?;
74 let trace_id = (!msg.trace_id.is_empty())
75 .then(|| uuid::Uuid::parse_str(&msg.trace_id).ok())
76 .flatten();
77
78 let (wallet_id, anonymous, subject_type, subject_id) =
79 match msg.wallet_selector.and_then(|s| s.mode) {
80 None => {
81 return Err(status_from_err(&LproviderError::Wallet(
82 WalletError::WalletRequired,
83 )));
84 }
85 Some(v1::wallet_selector::Mode::Wallet(w)) => {
86 let wallet = self
87 .wallets
88 .get_wallet(w.wallet_id)
89 .await
90 .map_err(|e| status_from_err(&e))?;
91 (
92 Some(wallet.wallet_id),
93 false,
94 wallet.owner_type,
95 wallet.owner_id,
96 )
97 }
98 Some(v1::wallet_selector::Mode::Anonymous(_)) => {
99 (None, true, "tenant".to_string(), tenant_id)
100 }
101 };
102
103 let (request, _alias, max_output) = proto_to_chat_request(&msg)?;
104 let lp = self.state.runtime.read().await.lp().clone();
105 let max_output = max_output.min(lp.max_output_tokens);
106
107 let selection = msg.model_selection.map(|m| ChatSelectionSpec {
108 strategy: strategy_name(m.strategy),
109 allowed: m.allowed_upstreams,
110 excluded: m.excluded_upstreams,
111 balance_bias: m.balance_bias as f64,
112 required_capabilities: m.required_capabilities,
113 });
114
115 let session_key_override =
116 (!msg.session_key_override.is_empty()).then_some(msg.session_key_override);
117
118 Ok(ChatContext {
119 request_id,
120 trace_id,
121 tenant_id,
122 wallet_id,
123 subject_type,
124 subject_id,
125 anonymous,
126 model_alias: if msg.model_alias.is_empty() {
127 "auto".into()
128 } else {
129 msg.model_alias
130 },
131 request,
132 max_output,
133 stream,
134 priority: msg.priority,
135 caller: None,
136 caller_service: (!msg.caller_service.is_empty()).then_some(msg.caller_service),
137 session_key_override,
138 selection: selection.unwrap_or_default(),
139 client_ip: None,
140 user_agent: None,
141 })
142 }
143}
144
145fn strategy_name(raw: i32) -> Option<String> {
146 match v1::model_selection::Strategy::try_from(raw) {
147 Ok(v1::model_selection::Strategy::Balanced) => Some("BALANCED".into()),
148 Ok(v1::model_selection::Strategy::Cost) => Some("COST".into()),
149 Ok(v1::model_selection::Strategy::Quality) => Some("QUALITY".into()),
150 Ok(v1::model_selection::Strategy::Latency) => Some("LATENCY".into()),
151 Ok(v1::model_selection::Strategy::Explicit) => Some("EXPLICIT".into()),
152 _ => None,
153 }
154}
155
156pub fn status_from_err(err: &LproviderError) -> Status {
157 match err.status_code() {
158 http::StatusCode::BAD_REQUEST => Status::invalid_argument(err.to_string()),
159 http::StatusCode::UNAUTHORIZED => Status::unauthenticated(err.to_string()),
160 http::StatusCode::PAYMENT_REQUIRED => Status::failed_precondition(err.to_string()),
161 http::StatusCode::FORBIDDEN => Status::permission_denied(err.to_string()),
162 http::StatusCode::NOT_FOUND => Status::not_found(err.to_string()),
163 http::StatusCode::CONFLICT => Status::aborted(err.to_string()),
164 http::StatusCode::TOO_MANY_REQUESTS => Status::resource_exhausted(err.to_string()),
165 http::StatusCode::BAD_GATEWAY => Status::unavailable(err.to_string()),
166 http::StatusCode::SERVICE_UNAVAILABLE => Status::unavailable(err.to_string()),
167 _ => Status::internal(err.to_string()),
168 }
169}
170
171fn proto_to_chat_request(msg: &v1::ChatRequestMsg) -> Result<(ChatRequest, String, usize), Status> {
172 let messages = msg
173 .messages
174 .iter()
175 .map(proto_message_to_llm)
176 .collect::<Result<Vec<_>, Status>>()?;
177 let tools = msg
178 .tools
179 .iter()
180 .map(proto_tool_to_llm)
181 .collect::<Result<Vec<_>, Status>>()?;
182
183 let mut request = ChatRequest::new(messages);
184 if !tools.is_empty() {
185 request = request.with_tools(tools);
186 }
187 if let Some(m) = msg.max_output_tokens
188 && m > 0
189 {
190 request = request.with_max_tokens(m as usize);
191 }
192 if let Some(t) = msg.temperature {
193 request = request.with_temperature(t as f32);
194 }
195 if let Some(p) = msg.top_p {
196 request = request.with_top_p(p as f32);
197 }
198 if let Some(rf) = &msg.response_format
199 && (rf.r#type == "json_object" || rf.r#type == "json_schema")
200 {
201 request = request.with_response_format(ResponseFormat::Json);
202 }
203 if !msg.stop_sequences.is_empty() {
204 request = request.with_stop_sequences(msg.stop_sequences.clone());
205 }
206 if let Some(opts) = &msg.provider_options {
207 request = request.with_provider_options(prost_struct_to_json(opts));
208 }
209 Ok((
210 request,
211 msg.model_alias.clone(),
212 msg.max_output_tokens.unwrap_or(0).max(0) as usize,
213 ))
214}
215
216fn proto_message_to_llm(m: &v1::ChatMessageMsg) -> Result<ChatMessage, Status> {
217 let role = match m.role.to_ascii_lowercase().as_str() {
218 "system" => MessageRole::System,
219 "user" => MessageRole::User,
220 "assistant" => MessageRole::Assistant,
221 "tool" => MessageRole::Tool,
222 other => {
223 return Err(status_from_err(&LproviderError::InvalidRequest(format!(
224 "unknown role: {other}"
225 ))));
226 }
227 };
228 let content = match &m.content {
229 None => MessageContent::Text(String::new()),
230 Some(v1::chat_message_msg::Content::Text(t)) => MessageContent::Text(t.clone()),
231 Some(v1::chat_message_msg::Content::Parts(p)) => {
232 let parts = p
233 .parts
234 .iter()
235 .map(|part| {
236 if part.r#type == "text" {
237 ContentPart::Text(part.text.clone())
238 } else {
239 ContentPart::Image {
240 url: part.url.clone(),
241 detail: ImageDetail::Auto,
242 }
243 }
244 })
245 .collect();
246 MessageContent::Parts(parts)
247 }
248 };
249 let tool_calls = m
250 .tool_calls
251 .iter()
252 .map(|tc| {
253 let (name, arguments) = tool_call_fields(&tc.function);
254 Ok::<_, Status>(ToolCall {
255 id: tc.id.clone(),
256 name,
257 arguments,
258 })
259 })
260 .collect::<Result<Vec<_>, _>>()?;
261 Ok(ChatMessage {
262 role,
263 content,
264 name: (!m.name.is_empty()).then(|| m.name.clone()),
265 tool_calls,
266 tool_call_id: (!m.tool_call_id.is_empty()).then(|| m.tool_call_id.clone()),
267 })
268}
269
270fn tool_call_fields(function: &Option<prost_types::Struct>) -> (String, serde_json::Value) {
271 let fields = function.as_ref().map(|s| &s.fields).unwrap_or_else(|| {
272 static EMPTY: std::sync::OnceLock<BTreeMap<String, prost_types::Value>> =
273 std::sync::OnceLock::new();
274 EMPTY.get_or_init(BTreeMap::new)
275 });
276 let name = fields
277 .get("name")
278 .and_then(|v| v.kind.as_ref())
279 .and_then(|k| match k {
280 prost_types::value::Kind::StringValue(s) => Some(s.clone()),
281 _ => None,
282 })
283 .unwrap_or_default();
284 let arguments = fields
285 .get("arguments")
286 .and_then(|v| v.kind.as_ref())
287 .and_then(|k| match k {
288 prost_types::value::Kind::StringValue(s) => serde_json::from_str(s).ok(),
289 _ => None,
290 })
291 .or_else(|| function.as_ref().map(prost_struct_to_json))
292 .unwrap_or(serde_json::Value::Object(Default::default()));
293 (name, arguments)
294}
295
296pub fn prost_value_to_json(v: &prost_types::Value) -> serde_json::Value {
297 match &v.kind {
298 Some(prost_types::value::Kind::NullValue(_)) => serde_json::Value::Null,
299 Some(prost_types::value::Kind::NumberValue(n)) => serde_json::Value::from(*n),
300 Some(prost_types::value::Kind::StringValue(s)) => serde_json::Value::String(s.clone()),
301 Some(prost_types::value::Kind::BoolValue(b)) => serde_json::Value::Bool(*b),
302 Some(prost_types::value::Kind::StructValue(s)) => prost_struct_to_json(s),
303 Some(prost_types::value::Kind::ListValue(l)) => {
304 serde_json::Value::Array(l.values.iter().map(prost_value_to_json).collect())
305 }
306 None => serde_json::Value::Null,
307 }
308}
309
310pub fn prost_struct_to_json(s: &prost_types::Struct) -> serde_json::Value {
311 let mut map = serde_json::Map::new();
312 for (k, v) in &s.fields {
313 map.insert(k.clone(), prost_value_to_json(v));
314 }
315 serde_json::Value::Object(map)
316}
317
318fn proto_tool_to_llm(t: &v1::ToolMsg) -> Result<ToolDefinition, Status> {
319 let fields = t.function.as_ref().map(|s| &s.fields).unwrap_or_else(|| {
320 static EMPTY: std::sync::OnceLock<BTreeMap<String, prost_types::Value>> =
321 std::sync::OnceLock::new();
322 EMPTY.get_or_init(BTreeMap::new)
323 });
324 let name = fields
325 .get("name")
326 .and_then(|v| v.kind.as_ref())
327 .and_then(|k| match k {
328 prost_types::value::Kind::StringValue(s) => Some(s.clone()),
329 _ => None,
330 })
331 .ok_or_else(|| Status::invalid_argument("tool missing name"))?;
332 let description = fields
333 .get("description")
334 .and_then(|v| v.kind.as_ref())
335 .and_then(|k| match k {
336 prost_types::value::Kind::StringValue(s) => Some(s.clone()),
337 _ => None,
338 })
339 .unwrap_or_default();
340 let parameters = fields
341 .get("parameters")
342 .and_then(|v| v.kind.as_ref())
343 .and_then(|k| match k {
344 prost_types::value::Kind::StructValue(s) => Some(prost_struct_to_json(s)),
345 _ => None,
346 })
347 .unwrap_or_else(|| serde_json::json!({"type": "object", "properties": {}}));
348 Ok(ToolDefinition::new(name, description, parameters))
349}
350
351fn finish_reason_str(r: FinishReason) -> &'static str {
352 match r {
353 FinishReason::Stop => "stop",
354 FinishReason::ToolUse => "tool_use",
355 FinishReason::MaxTokens => "max_tokens",
356 FinishReason::ContentFilter => "content_filter",
357 FinishReason::Error => "error",
358 }
359}
360
361fn tool_calls_to_proto(calls: &[ToolCall]) -> Vec<v1::ToolCallMsg> {
362 calls
363 .iter()
364 .map(|c| v1::ToolCallMsg {
365 id: c.id.clone(),
366 tool_type: "function".into(),
367 function: {
368 let mut fields = BTreeMap::new();
369 fields.insert(
370 "name".into(),
371 prost_types::Value {
372 kind: Some(prost_types::value::Kind::StringValue(c.name.clone())),
373 },
374 );
375 fields.insert(
376 "arguments".into(),
377 prost_types::Value {
378 kind: Some(prost_types::value::Kind::StringValue(
379 c.arguments.to_string(),
380 )),
381 },
382 );
383 Some(prost_types::Struct { fields })
384 },
385 })
386 .collect()
387}
388
389fn stream_event_to_proto(ev: StreamEvent) -> v1::StreamEventMsg {
390 let mut out = v1::StreamEventMsg::default();
391 match ev {
392 StreamEvent::TextDelta { delta } => {
393 out.event_type = "text_delta".into();
394 out.text_delta = delta;
395 }
396 StreamEvent::ReasoningDelta { delta } => {
397 out.event_type = "reasoning_delta".into();
398 out.reasoning_delta = delta;
399 }
400 StreamEvent::ToolCallDelta {
401 index,
402 id,
403 name,
404 args_delta,
405 } => {
406 out.event_type = "tool_call_delta".into();
407 out.tool_call_delta = Some(v1::ToolCallDeltaMsg {
408 id,
409 index: index.unwrap_or(0) as i32,
410 tool_type: "function".into(),
411 name_delta: name.unwrap_or_default(),
412 arguments_delta: args_delta,
413 });
414 }
415 StreamEvent::Usage(u) => {
416 out.event_type = "usage".into();
417 out.usage = Some(v1::UsageMsg {
418 prompt_tokens: u.input_tokens as i32,
419 completion_tokens: u.output_tokens as i32,
420 total_tokens: u.total() as i32,
421 cache_hit_tokens: 0,
422 cache_miss_tokens: 0,
423 points_input: 0,
424 points_output: 0,
425 points_cache_hit: 0,
426 points_total: 0,
427 });
428 }
429 StreamEvent::Done { finish_reason } => {
430 out.event_type = "done".into();
431 out.finish_reason = finish_reason_str(finish_reason).into();
432 }
433 }
434 out
435}
436
437#[tonic::async_trait]
438impl LproviderAi for AiGrpc {
439 type ChatStreamStream = ReceiverStream<Result<v1::StreamEventMsg, Status>>;
440 type ChatStreamBidirectionalStream = ReceiverStream<Result<v1::StreamEventMsg, Status>>;
441
442 async fn chat(
443 &self,
444 request: Request<v1::ChatRequestMsg>,
445 ) -> Result<Response<v1::ChatResponseMsg>, Status> {
446 let msg = request.into_inner();
447 let ctx = self.build_context(msg, false).await?;
448 let request_id = ctx.request_id;
449 let started = std::time::Instant::now();
450 let outcome = self
451 .router
452 .handle_chat(ctx)
453 .await
454 .map_err(|e| status_from_err(&e))?;
455 let latency_ms = started.elapsed().as_millis() as i64;
456
457 let bill = outcome.bill;
458 let usage = v1::UsageMsg {
459 prompt_tokens: outcome.response.usage.input_tokens as i32,
460 completion_tokens: outcome.response.usage.output_tokens as i32,
461 total_tokens: outcome.response.usage.total() as i32,
462 cache_hit_tokens: 0,
463 cache_miss_tokens: 0,
464 points_input: bill.points_input,
465 points_output: bill.points_output,
466 points_cache_hit: bill.points_cache_hit,
467 points_total: bill.points_total,
468 };
469 Ok(Response::new(v1::ChatResponseMsg {
470 id: request_id.to_string(),
471 model_alias: outcome.model,
472 upstream_code: outcome.upstream_code,
473 key_label: outcome.key_label,
474 content: outcome.response.content,
475 reasoning_text: String::new(),
476 tool_calls: tool_calls_to_proto(&outcome.response.tool_calls),
477 finish_reason: finish_reason_str(outcome.response.finish_reason).into(),
478 usage: Some(usage),
479 raw_provider_response: None,
480 latency_ms,
481 request_id: request_id.to_string(),
482 }))
483 }
484
485 async fn chat_stream(
486 &self,
487 request: Request<v1::ChatRequestMsg>,
488 ) -> Result<Response<Self::ChatStreamStream>, Status> {
489 let msg = request.into_inner();
490 let ctx = self.build_context(msg, true).await?;
491 let request_id = ctx.request_id.to_string();
492 let cancel = self.register_cancel(&request_id);
493
494 let mut billed = self
495 .router
496 .handle_stream(ctx)
497 .await
498 .map_err(|e| status_from_err(&e))?;
499
500 let (tx, rx) = mpsc::channel::<Result<v1::StreamEventMsg, Status>>(64);
501 let unregister = self.cancels.clone();
502 let rid = request_id.clone();
503 tokio::spawn(async move {
504 let mut cancel = cancel;
505 loop {
506 tokio::select! {
507 _ = cancel.recv() => {
508 let _ = tx.send(Err(Status::cancelled("cancelled"))).await;
509 break;
510 }
511 item = billed.next() => {
512 match item {
513 Some(Ok(ev)) => {
514 if tx.send(Ok(stream_event_to_proto(ev))).await.is_err() {
515 break;
516 }
517 }
518 Some(Err(e)) => {
519 let _ = tx.send(Err(Status::internal(e.to_string()))).await;
520 break;
521 }
522 None => break,
523 }
524 }
525 }
526 }
527 unregister.lock().unwrap().remove(&rid);
528 });
529 Ok(Response::new(ReceiverStream::new(rx)))
530 }
531
532 async fn chat_stream_bidirectional(
533 &self,
534 request: Request<tonic::Streaming<v1::ChatRequestMsg>>,
535 ) -> Result<Response<Self::ChatStreamBidirectionalStream>, Status> {
536 let mut inbound = request.into_inner();
537 let first = inbound
538 .next()
539 .await
540 .ok_or_else(|| Status::invalid_argument("empty stream"))??;
541 let request_id = Self::request_id(&first)?;
542 let ctx = self.build_context(first, true).await?;
543 let rid = request_id.to_string();
544 let cancel = self.register_cancel(&rid);
545 let cancels = self.cancels.clone();
546
547 let mut billed = self
548 .router
549 .handle_stream(ctx)
550 .await
551 .map_err(|e| status_from_err(&e))?;
552
553 let (tx, rx) = mpsc::channel::<Result<v1::StreamEventMsg, Status>>(64);
554 tokio::spawn(async move {
555 let out_tx = tx.clone();
556 let guard_task = tokio::spawn(async move {
557 while let Some(msg) = inbound.next().await {
558 match msg {
559 Ok(m) => {
560 let new_rid = (!m.request_id.is_empty())
561 .then(|| uuid::Uuid::parse_str(&m.request_id).ok())
562 .flatten();
563 if let Some(new_rid) = new_rid
564 && new_rid != request_id
565 {
566 let _ = out_tx
567 .send(Err(Status::invalid_argument(
568 "request_id must not change mid-stream",
569 )))
570 .await;
571 }
572 }
573 Err(e) => {
574 let _ = out_tx.send(Err(e)).await;
575 return;
576 }
577 }
578 }
579 });
580 let mut cancel = cancel;
581 loop {
582 tokio::select! {
583 _ = cancel.recv() => {
584 let _ = tx.send(Err(Status::cancelled("cancelled"))).await;
585 break;
586 }
587 item = billed.next() => {
588 match item {
589 Some(Ok(ev)) => {
590 if tx.send(Ok(stream_event_to_proto(ev))).await.is_err() {
591 break;
592 }
593 }
594 Some(Err(e)) => {
595 let _ = tx.send(Err(Status::internal(e.to_string()))).await;
596 break;
597 }
598 None => break,
599 }
600 }
601 }
602 }
603 guard_task.abort();
604 cancels.lock().unwrap().remove(&rid);
605 });
606 Ok(Response::new(ReceiverStream::new(rx)))
607 }
608
609 async fn count_tokens(
610 &self,
611 request: Request<v1::CountTokensRequest>,
612 ) -> Result<Response<v1::CountTokensResponse>, Status> {
613 let req = request.into_inner();
614 let (chat_request, _alias, _) = {
615 let msg = v1::ChatRequestMsg {
616 model_alias: req.model_alias,
617 messages: req.messages,
618 tools: req.tools,
619 ..Default::default()
620 };
621 proto_to_chat_request(&msg)?
622 };
623 let lp = self.state.runtime.read().await.lp().clone();
624 let input = self
625 .router
626 .count_tokens(&chat_request)
627 .await
628 .map_err(|e| status_from_err(&e))?;
629 let max_output = chat_request
630 .max_tokens
631 .unwrap_or(lp.max_output_tokens)
632 .min(lp.max_output_tokens);
633 let max_context = lp.max_context_tokens;
634 let fits = input.saturating_add(max_output) <= max_context;
635 Ok(Response::new(v1::CountTokensResponse {
636 estimated_input_tokens: input as i32,
637 estimated_output_tokens: max_output as i32,
638 max_context_tokens: max_context as i32,
639 max_output_tokens: max_output as i32,
640 fits_context: fits,
641 }))
642 }
643
644 async fn cancel(
645 &self,
646 request: Request<v1::CancelRequest>,
647 ) -> Result<Response<v1::CancelResponse>, Status> {
648 let req = request.into_inner();
649 let tx = self.cancels.lock().unwrap().get(&req.request_id).cloned();
650 let cancelled = if let Some(tx) = tx {
651 let _ = tx.send(()).await;
652 true
653 } else {
654 false
655 };
656 Ok(Response::new(v1::CancelResponse { cancelled }))
657 }
658}
Added agent/lprovider/audit.rs agent/lprovider/audit.rs +496 −0 View agent/lprovider/audit.rs
@@ -0,0 +1,496 @@
1use std::sync::Arc;
2use std::time::{SystemTime, UNIX_EPOCH};
3
4use db::conn::{Db, DbError};
5use sea_orm::{
6 ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseBackend, EntityTrait, PaginatorTrait,
7 QueryFilter, QueryOrder, QuerySelect, Statement, Value,
8};
9use time::OffsetDateTime;
10use tracing;
11use uuid::Uuid;
12
13use crate::error::LproviderError;
14use db::lprovider::llm_call_log::{self, Model as CallLogModel};
15use db::lprovider::lprovider_admin_audit;
16
17const ERROR_MAX_LEN: usize = 512;
18
19fn now() -> OffsetDateTime {
20 OffsetDateTime::from_unix_timestamp(
21 SystemTime::now()
22 .duration_since(UNIX_EPOCH)
23 .unwrap_or_default()
24 .as_secs() as i64,
25 )
26 .unwrap_or(OffsetDateTime::UNIX_EPOCH)
27}
28
29pub fn sanitize_error(raw: &str) -> String {
30 let mut s: String = raw
31 .chars()
32 .map(|c| if c.is_control() { ' ' } else { c })
33 .take(ERROR_MAX_LEN)
34 .collect();
35 redact_marker(&mut s, "sk-", "sk-***");
36 redact_marker(&mut s, "Bearer ", "Bearer ***");
37 redact_marker(&mut s, "api_key", "api_key=***");
38 redact_marker(&mut s, "api-key", "api-key=***");
39 redact_marker(&mut s, "apikey", "apikey=***");
40 s
41}
42
43fn redact_marker(s: &mut String, marker: &str, repl: &str) {
44 let mut search_from = 0usize;
45 while let Some(rel) = s[search_from..].find(marker) {
46 let pos = search_from + rel;
47 let mut end = pos + marker.len();
48 let bytes = s.as_bytes();
49 if end < bytes.len() && (bytes[end] == b'=' || bytes[end] == b':') {
50 end += 1;
51 }
52 while end < bytes.len() {
53 let c = bytes[end];
54 if c.is_ascii_alphanumeric() || c == b'_' || c == b'-' || c == b'.' {
55 end += 1;
56 } else {
57 break;
58 }
59 }
60 if end <= pos + marker.len() {
61 search_from = pos + marker.len();
62 continue;
63 }
64 s.replace_range(pos..end, repl);
65 search_from = pos + repl.len();
66 }
67}
68
69#[derive(Clone, Debug)]
70pub struct CallLogEntry {
71 pub request_id: Uuid,
72 pub trace_id: Option<Uuid>,
73 pub tenant_id: i64,
74 pub api_key_id: Option<i64>,
75 pub subject_type: Option<String>,
76 pub subject_id: Option<i64>,
77 pub wallet_id: Option<i64>,
78 pub anonymous: bool,
79 pub subject_path: Option<String>,
80 pub session_key: Option<String>,
81 pub model_alias: String,
82 pub upstream_code: String,
83 pub requested_model: Option<String>,
84 pub chosen_strategy: Option<String>,
85 pub tier_cost: i32,
86 pub tier_quality: i32,
87 pub tier_latency: i32,
88 pub key_label: Option<String>,
89 pub caller: Option<String>,
90 pub caller_service: Option<String>,
91 pub priority: i32,
92 pub streamed: bool,
93 pub status: String,
94 pub error_kind: Option<String>,
95 pub error_message: Option<String>,
96 pub prompt_tokens: i32,
97 pub completion_tokens: i32,
98 pub cache_hit_tokens: i32,
99 pub cache_miss_tokens: i32,
100 pub billable_input_tokens: i32,
101 pub billable_output_tokens: i32,
102 pub points_input: i64,
103 pub points_output: i64,
104 pub points_cache_hit: i64,
105 pub points_total: i64,
106 pub points_rate_input: i64,
107 pub points_rate_output: i64,
108 pub points_rate_cache_hit: i64,
109 pub request_started_at: OffsetDateTime,
110 pub response_completed_at: Option<OffsetDateTime>,
111 pub latency_ms: i32,
112 pub client_ip: Option<String>,
113 pub user_agent: Option<String>,
114}
115
116impl Default for CallLogEntry {
117 fn default() -> Self {
118 Self {
119 request_id: Uuid::nil(),
120 trace_id: None,
121 tenant_id: 0,
122 api_key_id: None,
123 subject_type: None,
124 subject_id: None,
125 wallet_id: None,
126 anonymous: false,
127 subject_path: None,
128 session_key: None,
129 model_alias: String::new(),
130 upstream_code: String::new(),
131 requested_model: None,
132 chosen_strategy: None,
133 tier_cost: 0,
134 tier_quality: 0,
135 tier_latency: 0,
136 key_label: None,
137 caller: None,
138 caller_service: None,
139 priority: 0,
140 streamed: false,
141 status: String::new(),
142 error_kind: None,
143 error_message: None,
144 prompt_tokens: 0,
145 completion_tokens: 0,
146 cache_hit_tokens: 0,
147 cache_miss_tokens: 0,
148 billable_input_tokens: 0,
149 billable_output_tokens: 0,
150 points_input: 0,
151 points_output: 0,
152 points_cache_hit: 0,
153 points_total: 0,
154 points_rate_input: 0,
155 points_rate_output: 0,
156 points_rate_cache_hit: 0,
157 request_started_at: OffsetDateTime::UNIX_EPOCH,
158 response_completed_at: None,
159 latency_ms: 0,
160 client_ip: None,
161 user_agent: None,
162 }
163 }
164}
165
166#[derive(Clone)]
167pub struct AuditService {
168 db: Option<Arc<Db>>,
169}
170
171impl AuditService {
172 pub fn new(db: Option<Arc<Db>>) -> Self {
173 Self { db }
174 }
175
176 fn conn(&self) -> Result<&Db, LproviderError> {
177 self.db
178 .as_deref()
179 .ok_or_else(|| LproviderError::Db(DbError::InvalidConfig("db not configured".into())))
180 }
181
182 pub async fn insert_call_log(&self, e: &CallLogEntry) -> Result<(), LproviderError> {
183 let db = self.conn()?;
184 let am: llm_call_log::ActiveModel = llm_call_log::ActiveModel::builder()
185 .set_request_id(e.request_id)
186 .set_trace_id(e.trace_id)
187 .set_tenant_id(e.tenant_id)
188 .set_api_key_id(e.api_key_id)
189 .set_subject_type(e.subject_type.clone())
190 .set_subject_id(e.subject_id)
191 .set_wallet_id(e.wallet_id)
192 .set_anonymous(e.anonymous)
193 .set_subject_path(e.subject_path.clone())
194 .set_session_key(e.session_key.clone())
195 .set_model_alias(e.model_alias.clone())
196 .set_upstream_code(e.upstream_code.clone())
197 .set_requested_model(e.requested_model.clone())
198 .set_chosen_strategy(e.chosen_strategy.clone())
199 .set_tier_cost(e.tier_cost)
200 .set_tier_quality(e.tier_quality)
201 .set_tier_latency(e.tier_latency)
202 .set_key_label(e.key_label.clone())
203 .set_caller(e.caller.clone())
204 .set_caller_service(e.caller_service.clone())
205 .set_priority(e.priority)
206 .set_streamed(e.streamed)
207 .set_status(e.status.clone())
208 .set_error_kind(e.error_kind.clone())
209 .set_error_message(e.error_message.clone())
210 .set_prompt_tokens(e.prompt_tokens)
211 .set_completion_tokens(e.completion_tokens)
212 .set_cache_hit_tokens(e.cache_hit_tokens)
213 .set_cache_miss_tokens(e.cache_miss_tokens)
214 .set_billable_input_tokens(e.billable_input_tokens)
215 .set_billable_output_tokens(e.billable_output_tokens)
216 .set_points_input(e.points_input)
217 .set_points_output(e.points_output)
218 .set_points_cache_hit(e.points_cache_hit)
219 .set_points_total(e.points_total)
220 .set_points_rate_input(e.points_rate_input)
221 .set_points_rate_output(e.points_rate_output)
222 .set_points_rate_cache_hit(e.points_rate_cache_hit)
223 .set_request_started_at(e.request_started_at)
224 .set_response_completed_at(e.response_completed_at)
225 .set_latency_ms(e.latency_ms)
226 .set_client_ip(e.client_ip.clone())
227 .set_user_agent(e.user_agent.clone())
228 .set_created_at(now())
229 .into();
230 am.insert(db.primary_connection()).await.map_err(|err| {
231 LproviderError::Db(DbError::InvalidConfig(format!(
232 "insert llm_call_log failed: {err}"
233 )))
234 })?;
235 Ok(())
236 }
237
238 #[allow(clippy::too_many_arguments)]
239 pub async fn admin_audit(
240 &self,
241 operator: &str,
242 rpc_method: &str,
243 instance_id: Option<&str>,
244 target: Option<&str>,
245 change_summary: Option<&str>,
246 config_sha256: Option<&str>,
247 effect_mode: Option<&str>,
248 result: &str,
249 error: Option<&str>,
250 ) -> Result<(), LproviderError> {
251 let db = self.conn()?;
252 let am: lprovider_admin_audit::ActiveModel = lprovider_admin_audit::ActiveModel::builder()
253 .set_operator(operator.to_string())
254 .set_rpc_method(rpc_method.to_string())
255 .set_instance_id(instance_id.map(|s| s.to_string()))
256 .set_target(target.map(|s| s.to_string()))
257 .set_change_summary(change_summary.map(|s| s.to_string()))
258 .set_config_sha256(config_sha256.map(|s| s.to_string()))
259 .set_effect_mode(effect_mode.map(|s| s.to_string()))
260 .set_result(result.to_string())
261 .set_error_message(error.map(sanitize_error))
262 .set_created_at(now())
263 .into();
264 am.insert(db.primary_connection()).await.map_err(|err| {
265 LproviderError::Db(DbError::InvalidConfig(format!(
266 "insert lprovider_admin_audit failed: {err}"
267 )))
268 })?;
269 Ok(())
270 }
271
272 #[allow(clippy::too_many_arguments)]
273 pub async fn query_audit(
274 &self,
275 tenant_id: Option<i64>,
276 subject_type: Option<&str>,
277 subject_id: Option<i64>,
278 model_alias: Option<&str>,
279 upstream_code: Option<&str>,
280 status: Option<&str>,
281 from_ms: Option<i64>,
282 to_ms: Option<i64>,
283 limit: u64,
284 offset: u64,
285 ) -> Result<(Vec<CallLogModel>, u64), LproviderError> {
286 let db = self.conn()?;
287 let mut select = llm_call_log::Entity::find();
288 if let Some(v) = tenant_id {
289 select = select.filter(llm_call_log::Column::TenantId.eq(v));
290 }
291 if let Some(v) = subject_type {
292 select = select.filter(llm_call_log::Column::SubjectType.eq(v.to_string()));
293 }
294 if let Some(v) = subject_id {
295 select = select.filter(llm_call_log::Column::SubjectId.eq(v));
296 }
297 if let Some(v) = model_alias {
298 select = select.filter(llm_call_log::Column::ModelAlias.eq(v.to_string()));
299 }
300 if let Some(v) = upstream_code {
301 select = select.filter(llm_call_log::Column::UpstreamCode.eq(v.to_string()));
302 }
303 if let Some(v) = status {
304 select = select.filter(llm_call_log::Column::Status.eq(v.to_string()));
305 }
306 if let Some(ms) = from_ms {
307 let ts = OffsetDateTime::from_unix_timestamp(ms / 1000)
308 .unwrap_or(OffsetDateTime::UNIX_EPOCH);
309 select = select.filter(llm_call_log::Column::CreatedAt.gte(ts));
310 }
311 if let Some(ms) = to_ms {
312 let ts = OffsetDateTime::from_unix_timestamp(ms / 1000)
313 .unwrap_or(OffsetDateTime::UNIX_EPOCH);
314 select = select.filter(llm_call_log::Column::CreatedAt.lte(ts));
315 }
316 let total = select
317 .clone()
318 .count(db.primary_connection())
319 .await
320 .map_err(db_err)?;
321 let rows = select
322 .order_by_desc(llm_call_log::Column::CreatedAt)
323 .offset(offset)
324 .limit(limit)
325 .all(db.primary_connection())
326 .await
327 .map_err(db_err)?;
328 Ok((rows, total))
329 }
330
331 #[allow(clippy::too_many_arguments)]
332 pub async fn aggregate_billing(
333 &self,
334 tenant_id: Option<i64>,
335 subject_type: Option<&str>,
336 subject_id: Option<i64>,
337 period_kind: Option<&str>,
338 period_start_ms: Option<i64>,
339 period_end_ms: Option<i64>,
340 model_alias: Option<&str>,
341 upstream_code: Option<&str>,
342 limit: u64,
343 offset: u64,
344 ) -> Result<(Vec<db::lprovider::llm_billing_aggregate::Model>, u64), LproviderError> {
345 use db::lprovider::llm_billing_aggregate;
346 let db = self.conn()?;
347 let mut select = llm_billing_aggregate::Entity::find();
348 if let Some(v) = tenant_id {
349 select = select.filter(llm_billing_aggregate::Column::TenantId.eq(v));
350 }
351 if let Some(v) = subject_type {
352 select = select.filter(llm_billing_aggregate::Column::SubjectType.eq(v.to_string()));
353 }
354 if let Some(v) = subject_id {
355 select = select.filter(llm_billing_aggregate::Column::SubjectId.eq(v));
356 }
357 if let Some(v) = period_kind {
358 select = select.filter(llm_billing_aggregate::Column::PeriodKind.eq(v.to_string()));
359 }
360 if let Some(ms) = period_start_ms {
361 let ts = OffsetDateTime::from_unix_timestamp(ms / 1000)
362 .unwrap_or(OffsetDateTime::UNIX_EPOCH);
363 select = select.filter(llm_billing_aggregate::Column::PeriodStart.gte(ts));
364 }
365 if let Some(ms) = period_end_ms {
366 let ts = OffsetDateTime::from_unix_timestamp(ms / 1000)
367 .unwrap_or(OffsetDateTime::UNIX_EPOCH);
368 select = select.filter(llm_billing_aggregate::Column::PeriodEnd.lte(ts));
369 }
370 if let Some(v) = model_alias {
371 select = select.filter(llm_billing_aggregate::Column::ModelAlias.eq(v.to_string()));
372 }
373 if let Some(v) = upstream_code {
374 select = select.filter(llm_billing_aggregate::Column::UpstreamCode.eq(v.to_string()));
375 }
376 let total = select
377 .clone()
378 .count(db.primary_connection())
379 .await
380 .map_err(db_err)?;
381 let rows = select
382 .order_by_desc(llm_billing_aggregate::Column::PeriodStart)
383 .offset(offset)
384 .limit(limit)
385 .all(db.primary_connection())
386 .await
387 .map_err(db_err)?;
388 Ok((rows, total))
389 }
390
391 pub async fn aggregate_incremental(&self) -> Result<u64, LproviderError> {
392 let db = self.conn()?;
393 let last: Option<time::OffsetDateTime> = {
394 let row = db
395 .writer()
396 .query_one_raw(Statement::from_string(
397 DatabaseBackend::Postgres,
398 "SELECT MAX(period_end) AS max_end FROM llm_billing_aggregate",
399 ))
400 .await
401 .map_err(db_err)?;
402 row.and_then(|r| r.try_get("", "max_end").ok())
403 };
404 let since = last.unwrap_or_else(|| now() - time::Duration::hours(1));
405 let sql = r#"
406INSERT INTO llm_billing_aggregate
407 (tenant_id, subject_type, subject_id, period_start, period_end, period_kind,
408 model_alias, upstream_code, call_count, success_count, error_count,
409 total_input_tokens, total_output_tokens, total_cache_hit_tokens, total_points,
410 created_at, updated_at)
411SELECT tenant_id, subject_type, subject_id,
412 date_trunc('hour', created_at),
413 date_trunc('hour', created_at) + interval '1 hour',
414 'hour', model_alias, upstream_code,
415 count(*) AS c, count(*) FILTER (WHERE status = 'success'),
416 count(*) FILTER (WHERE status = 'error'),
417 COALESCE(sum(prompt_tokens), 0), COALESCE(sum(completion_tokens), 0),
418 COALESCE(sum(cache_hit_tokens), 0), COALESCE(sum(points_total), 0),
419 now(), now()
420FROM llm_call_log
421WHERE created_at >= $1 AND subject_type IS NOT NULL AND subject_id IS NOT NULL
422GROUP BY tenant_id, subject_type, subject_id, date_trunc('hour', created_at),
423 model_alias, upstream_code
424ON CONFLICT (tenant_id, subject_type, subject_id, period_start, period_kind,
425 model_alias, upstream_code)
426DO UPDATE SET
427 call_count = llm_billing_aggregate.call_count + EXCLUDED.call_count,
428 success_count = llm_billing_aggregate.success_count + EXCLUDED.success_count,
429 error_count = llm_billing_aggregate.error_count + EXCLUDED.error_count,
430 total_input_tokens = llm_billing_aggregate.total_input_tokens + EXCLUDED.total_input_tokens,
431 total_output_tokens = llm_billing_aggregate.total_output_tokens + EXCLUDED.total_output_tokens,
432 total_cache_hit_tokens = llm_billing_aggregate.total_cache_hit_tokens + EXCLUDED.total_cache_hit_tokens,
433 total_points = llm_billing_aggregate.total_points + EXCLUDED.total_points,
434 updated_at = now()
435"#;
436 let since_iso = since
437 .format(&time::format_description::well_known::Rfc3339)
438 .map_err(|e| LproviderError::Internal(format!("format since: {e}")))?;
439 let sql = sql.replace("created_at >= $1", "created_at >= $1::timestamptz");
440 let res = db
441 .writer()
442 .execute_raw(Statement::from_sql_and_values(
443 DatabaseBackend::Postgres,
444 &sql,
445 [Value::String(Some(since_iso))],
446 ))
447 .await
448 .map_err(db_err)?;
449 Ok(res.rows_affected())
450 }
451
452 pub async fn reconcile_daily(&self) -> Result<f64, LproviderError> {
453 let db = self.conn()?;
454 let a: i64 = db
455 .writer()
456 .query_one_raw(Statement::from_string(
457 DatabaseBackend::Postgres,
458 "SELECT COALESCE(SUM(points_total), 0) AS a FROM llm_call_log
459 WHERE created_at >= date_trunc('day', now())
460 AND subject_type IS NOT NULL AND subject_id IS NOT NULL",
461 ))
462 .await
463 .map_err(db_err)?
464 .and_then(|r| r.try_get("", "a").ok())
465 .unwrap_or(0);
466 let b: i64 = db
467 .writer()
468 .query_one_raw(Statement::from_string(
469 DatabaseBackend::Postgres,
470 "SELECT COALESCE(SUM(total_points), 0) AS b FROM llm_billing_aggregate
471 WHERE period_start >= date_trunc('day', now())",
472 ))
473 .await
474 .map_err(db_err)?
475 .and_then(|r| r.try_get("", "b").ok())
476 .unwrap_or(0);
477 let ratio = if a == 0 {
478 0.0
479 } else {
480 (b as f64 - a as f64).abs() / a as f64
481 };
482 if ratio > 0.001 {
483 tracing::error!(
484 ratio,
485 call_log_points = a,
486 aggregate_points = b,
487 "lprovider billing reconcile drift exceeds 0.1%"
488 );
489 }
490 Ok(ratio)
491 }
492}
493
494fn db_err(err: sea_orm::DbErr) -> LproviderError {
495 LproviderError::Db(DbError::InvalidConfig(format!("sea-orm: {err}")))
496}
Added agent/lprovider/billing.rs agent/lprovider/billing.rs +72 −0 View agent/lprovider/billing.rs
@@ -0,0 +1,72 @@
1use config::lprovider::{LproviderPricingSettings, PointsPer1k};
2
3#[derive(Clone, Copy, Debug, Default)]
4pub struct Bill {
5 pub billable_input_tokens: i64,
6 pub billable_output_tokens: i64,
7 pub billable_cache_hit_tokens: i64,
8 pub points_input: i64,
9 pub points_output: i64,
10 pub points_cache_hit: i64,
11 pub points_total: i64,
12}
13
14pub fn ceil_div(numerator: i64, denominator: i64) -> i64 {
15 if numerator <= 0 {
16 return 0;
17 }
18 (numerator + denominator - 1) / denominator
19}
20
21pub fn points_for_tokens(tokens: i64, points_per_1k: i64) -> i64 {
22 if points_per_1k <= 0 || tokens <= 0 {
23 return 0;
24 }
25 ceil_div(tokens * points_per_1k, 1000)
26}
27
28pub fn rates_for<'a>(
29 upstream_code: &str,
30 pricing: &'a LproviderPricingSettings,
31) -> &'a PointsPer1k {
32 pricing
33 .rates
34 .get(upstream_code)
35 .unwrap_or(&pricing.default_per_1k)
36}
37
38pub fn compute_bill(
39 prompt_tokens: i64,
40 completion_tokens: i64,
41 cache_hit_tokens: i64,
42 rate: &PointsPer1k,
43) -> Bill {
44 let billable_input = (prompt_tokens - cache_hit_tokens).max(0);
45 let billable_output = completion_tokens.max(0);
46 let billable_cache_hit = cache_hit_tokens.max(0);
47 let points_input = points_for_tokens(billable_input, rate.points_per_1k_input_tokens);
48 let points_output = points_for_tokens(billable_output, rate.points_per_1k_output_tokens);
49 let points_cache_hit =
50 points_for_tokens(billable_cache_hit, rate.points_per_1k_cache_hit_tokens);
51 Bill {
52 billable_input_tokens: billable_input,
53 billable_output_tokens: billable_output,
54 billable_cache_hit_tokens: billable_cache_hit,
55 points_input,
56 points_output,
57 points_cache_hit,
58 points_total: points_input + points_output + points_cache_hit,
59 }
60}
61
62pub fn worst_case_estimate(
63 estimated_input_tokens: i64,
64 max_output_tokens: i64,
65 rate: &PointsPer1k,
66) -> i64 {
67 let worst_rate = rate
68 .points_per_1k_input_tokens
69 .max(rate.points_per_1k_output_tokens)
70 .max(rate.points_per_1k_cache_hit_tokens);
71 points_for_tokens(estimated_input_tokens + max_output_tokens, worst_rate)
72}
Added agent/lprovider/build.rs agent/lprovider/build.rs +11 −0 View agent/lprovider/build.rs
@@ -0,0 +1,11 @@
1fn main() -> Result<(), Box<dyn std::error::Error>> {
2 println!("cargo:rerun-if-changed=ai.proto");
3 println!("cargo:rerun-if-changed=admin.proto");
4
5 tonic_prost_build::configure()
6 .build_client(true)
7 .build_server(true)
8 .compile_protos(&["ai.proto", "admin.proto"], &["."])?;
9
10 Ok(())
11}
Modified agent/lprovider/Cargo.toml agent/lprovider/Cargo.toml +45 −0 View agent/lprovider/Cargo.toml
@@ -11,6 +11,51 @@
1111[lib]
1212path = "lib.rs"
1313name = "lprovider"
14
15[[bin]]
16name = "lprovider"
17path = "main.rs"
18
19[build-dependencies]
20tonic-prost-build.workspace = true
21
1422[dependencies]
23aes-gcm.workspace = true
24async-trait.workspace = true
1525axum.workspace = true
26base64.workspace = true
27bytes.workspace = true
28cache.workspace = true
29clap.workspace = true
30config.workspace = true
31db.workspace = true
32futures.workspace = true
33hex.workspace = true
34http.workspace = true
35llm = { path = "../llm" }
36metrics.workspace = true
37prost.workspace = true
38prost-types.workspace = true
39rand.workspace = true
40redis.workspace = true
41reqwest.workspace = true
42sea-orm.workspace = true
43secret.workspace = true
44serde.workspace = true
45serde_json.workspace = true
46serde_yaml.workspace = true
47sha2.workspace = true
48thiserror.workspace = true
49time.workspace = true
1650tokio.workspace = true
51tokio-stream.workspace = true
52tokio-util.workspace = true
53tonic.workspace = true
54tonic-health.workspace = true
55tonic-prost.workspace = true
56tracing.workspace = true
57tracing-subscriber.workspace = true
58uuid.workspace = true
59
60[dev-dependencies]
61tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] }
Added agent/lprovider/config.rs agent/lprovider/config.rs +72 −0 View agent/lprovider/config.rs
@@ -0,0 +1,72 @@
1use std::path::PathBuf;
2
3use config::{AppConfig, ConfigLoader};
4use sha2::{Digest, Sha256};
5
6use crate::error::LproviderError;
7
8#[derive(Debug, Clone, Default)]
9pub struct CliOverrides {
10 pub config: PathBuf,
11 pub bind: Option<String>,
12 pub port: Option<u16>,
13 pub admin_bind: Option<String>,
14 pub admin_port: Option<u16>,
15 pub admin_cert: Option<PathBuf>,
16 pub admin_key: Option<PathBuf>,
17 pub instance_id: Option<String>,
18 pub cluster_name: Option<String>,
19}
20
21#[derive(Debug)]
22pub struct RuntimeConfig {
23 pub app: AppConfig,
24 pub config_path: PathBuf,
25 pub config_sha256: String,
26}
27
28impl RuntimeConfig {
29 pub fn load(cli: &CliOverrides) -> Result<Self, LproviderError> {
30 let mut app = ConfigLoader::new()
31 .config_file(cli.config.clone())
32 .load()
33 .map_err(|e| LproviderError::Config(e.to_string()))?;
34
35 if let Some(bind) = &cli.bind {
36 app.lprovider.bind = bind.clone();
37 }
38 if let Some(port) = cli.port {
39 app.lprovider.port = port;
40 }
41 if let Some(bind) = &cli.admin_bind {
42 app.lprovider.admin_bind = bind.clone();
43 }
44 if let Some(port) = cli.admin_port {
45 app.lprovider.admin_port = port;
46 }
47
48 app.lprovider.is_valid().map_err(LproviderError::Config)?;
49
50 let config_sha256 = sha256_file(&cli.config)?;
51
52 Ok(Self {
53 app,
54 config_path: cli.config.clone(),
55 config_sha256,
56 })
57 }
58
59 pub fn lp(&self) -> &config::lprovider::LproviderConfig {
60 &self.app.lprovider
61 }
62}
63
64pub fn sha256_file(path: &std::path::Path) -> Result<String, LproviderError> {
65 let bytes = std::fs::read(path).map_err(|e| LproviderError::Config(e.to_string()))?;
66 Ok(sha256_hex(&bytes))
67}
68
69pub fn sha256_hex(bytes: &[u8]) -> String {
70 let digest = Sha256::digest(bytes);
71 hex::encode(digest)
72}
Added agent/lprovider/credentials.rs agent/lprovider/credentials.rs +46 −0 View agent/lprovider/credentials.rs
@@ -0,0 +1,46 @@
1use config::lprovider::{LproviderConfig, LproviderUpstream, LproviderUpstreamCredential};
2use llm::provider::openai::SecretString;
3use secret::{
4 SecretError, envelope_key_from_env, is_envelope, key_fingerprint, resolve,
5 resolve_credential_ref,
6};
7
8pub struct ResolvedCredential {
9 pub upstream_code: String,
10 pub label: String,
11 pub api_key: SecretString,
12 pub fingerprint: String,
13}
14
15pub fn resolve_all(cfg: &LproviderConfig) -> Result<Vec<ResolvedCredential>, SecretError> {
16 let mut out = Vec::new();
17 for upstream in &cfg.upstreams {
18 if !upstream.enabled {
19 continue;
20 }
21 for cred in &upstream.credentials {
22 out.push(resolve_one(upstream, cred)?);
23 }
24 }
25 Ok(out)
26}
27
28fn resolve_one(
29 upstream: &LproviderUpstream,
30 cred: &LproviderUpstreamCredential,
31) -> Result<ResolvedCredential, SecretError> {
32 let raw = resolve_credential_ref(&cred.credential_ref)?;
33 let api_key = if is_envelope(&raw) {
34 let key = envelope_key_from_env()?;
35 resolve(&key, &raw)?
36 } else {
37 raw
38 };
39 let fingerprint = key_fingerprint(&api_key);
40 Ok(ResolvedCredential {
41 upstream_code: upstream.code.clone(),
42 label: cred.label.clone(),
43 api_key: SecretString::new(api_key),
44 fingerprint,
45 })
46}
Added agent/lprovider/error.rs agent/lprovider/error.rs +165 −0 View agent/lprovider/error.rs
@@ -0,0 +1,165 @@
1use axum::Json;
2use axum::http::StatusCode;
3use axum::response::{IntoResponse, Response};
4use serde_json::json;
5
6#[derive(Debug, thiserror::Error)]
7pub enum LproviderError {
8 #[error("config error: {0}")]
9 Config(String),
10
11 #[error("database error: {0}")]
12 Db(#[from] db::conn::DbError),
13
14 #[error("redis error: {0}")]
15 Redis(String),
16
17 #[error("invalid request: {0}")]
18 InvalidRequest(String),
19
20 #[error("auth: {0}")]
21 Auth(AuthError),
22
23 #[error("wallet: {0}")]
24 Wallet(WalletError),
25
26 #[error("quota exceeded: {0}")]
27 QuotaExceeded(String),
28
29 #[error("upstream: {0}")]
30 Upstream(UpstreamError),
31
32 #[error("resource locked: {0}")]
33 ResourceLocked(String),
34
35 #[error("not found: {0}")]
36 NotFound(String),
37
38 #[error("conflict: {0}")]
39 Conflict(String),
40
41 #[error("internal error: {0}")]
42 Internal(String),
43}
44
45#[derive(Debug, thiserror::Error)]
46pub enum AuthError {
47 #[error("missing bearer token")]
48 MissingToken,
49
50 #[error("invalid bearer token")]
51 InvalidToken,
52
53 #[error("permission denied: {0}")]
54 PermissionDenied(String),
55}
56
57#[derive(Debug, thiserror::Error)]
58pub enum WalletError {
59 #[error("wallet id required")]
60 WalletRequired,
61
62 #[error("wallet not found: {0}")]
63 WalletNotFound(i64),
64
65 #[error("wallet not active: {0}")]
66 WalletNotActive(i64),
67
68 #[error("insufficient balance")]
69 InsufficientBalance,
70
71 #[error("wallet debit failed")]
72 DebitFailed,
73
74 #[error("wallet owner mismatch")]
75 OwnerMismatch,
76}
77
78#[derive(Debug, thiserror::Error)]
79pub enum UpstreamError {
80 #[error("no upstream available")]
81 NoUpstreamAvailable,
82
83 #[error("no key available for upstream {0}")]
84 NoKeyAvailable(String),
85
86 #[error("context length exceeded")]
87 ContextLengthExceeded,
88
89 #[error("all upstream attempts failed: {0}")]
90 AllFailed(String),
91}
92
93impl LproviderError {
94 pub fn status_code(&self) -> StatusCode {
95 match self {
96 LproviderError::Config(_) | LproviderError::Db(_) | LproviderError::Internal(_) => {
97 StatusCode::INTERNAL_SERVER_ERROR
98 }
99 LproviderError::Redis(_)
100 | LproviderError::Upstream(UpstreamError::NoUpstreamAvailable)
101 | LproviderError::Upstream(UpstreamError::NoKeyAvailable(_)) => {
102 StatusCode::SERVICE_UNAVAILABLE
103 }
104 LproviderError::Upstream(UpstreamError::AllFailed(_)) => StatusCode::BAD_GATEWAY,
105 LproviderError::InvalidRequest(_)
106 | LproviderError::Upstream(UpstreamError::ContextLengthExceeded) => {
107 StatusCode::BAD_REQUEST
108 }
109 LproviderError::Auth(AuthError::MissingToken)
110 | LproviderError::Auth(AuthError::InvalidToken) => StatusCode::UNAUTHORIZED,
111 LproviderError::Auth(AuthError::PermissionDenied(_))
112 | LproviderError::Wallet(WalletError::WalletNotActive(_))
113 | LproviderError::Wallet(WalletError::OwnerMismatch) => StatusCode::FORBIDDEN,
114 LproviderError::Wallet(WalletError::WalletRequired)
115 | LproviderError::Wallet(WalletError::InsufficientBalance)
116 | LproviderError::Wallet(WalletError::DebitFailed) => StatusCode::PAYMENT_REQUIRED,
117 LproviderError::Wallet(WalletError::WalletNotFound(_)) => StatusCode::NOT_FOUND,
118 LproviderError::NotFound(_) => StatusCode::NOT_FOUND,
119 LproviderError::ResourceLocked(_) | LproviderError::Conflict(_) => StatusCode::CONFLICT,
120 LproviderError::QuotaExceeded(_) => StatusCode::TOO_MANY_REQUESTS,
121 }
122 }
123}
124
125impl IntoResponse for LproviderError {
126 fn into_response(self) -> Response {
127 let status = self.status_code();
128 let body = json!({
129 "error": {
130 "message": self.to_string(),
131 "type": error_type(&self),
132 "param": null,
133 "code": status.as_u16(),
134 }
135 });
136 (status, Json(body)).into_response()
137 }
138}
139
140fn error_type(err: &LproviderError) -> &'static str {
141 match err {
142 LproviderError::Config(_) => "config_error",
143 LproviderError::Db(_) => "database_error",
144 LproviderError::Redis(_) => "redis_error",
145 LproviderError::InvalidRequest(_) => "invalid_request_error",
146 LproviderError::Auth(AuthError::MissingToken) => "missing_token",
147 LproviderError::Auth(AuthError::InvalidToken) => "invalid_token",
148 LproviderError::Auth(AuthError::PermissionDenied(_)) => "permission_denied",
149 LproviderError::Wallet(WalletError::WalletRequired) => "wallet_required",
150 LproviderError::Wallet(WalletError::WalletNotFound(_)) => "wallet_not_found",
151 LproviderError::Wallet(WalletError::WalletNotActive(_)) => "wallet_not_active",
152 LproviderError::Wallet(WalletError::InsufficientBalance) => "insufficient_balance",
153 LproviderError::Wallet(WalletError::DebitFailed) => "wallet_debit_failed",
154 LproviderError::Wallet(WalletError::OwnerMismatch) => "wallet_owner_mismatch",
155 LproviderError::QuotaExceeded(_) => "quota_exceeded",
156 LproviderError::Upstream(UpstreamError::NoUpstreamAvailable) => "no_upstream_available",
157 LproviderError::Upstream(UpstreamError::NoKeyAvailable(_)) => "no_key_available",
158 LproviderError::Upstream(UpstreamError::ContextLengthExceeded) => "context_length_exceeded",
159 LproviderError::Upstream(UpstreamError::AllFailed(_)) => "all_upstreams_failed",
160 LproviderError::ResourceLocked(_) => "resource_locked",
161 LproviderError::NotFound(_) => "not_found",
162 LproviderError::Conflict(_) => "conflict",
163 LproviderError::Internal(_) => "internal_error",
164 }
165}
Added agent/lprovider/http.rs agent/lprovider/http.rs +591 −0 View agent/lprovider/http.rs
@@ -0,0 +1,591 @@
1use std::convert::Infallible;
2use std::sync::Arc;
3
4use axum::extract::State;
5use axum::http::{HeaderMap, StatusCode};
6use axum::response::sse::{Event, KeepAlive, Sse};
7use axum::response::{IntoResponse, Response};
8use axum::routing::{get, post};
9use axum::{Json, Router as AxumRouter};
10use futures::StreamExt;
11use llm::types::{
12 ChatMessage, ChatRequest, ContentPart, FinishReason, ImageDetail, MessageContent, MessageRole,
13 ResponseFormat, StreamEvent, ToolCall, ToolDefinition,
14};
15use serde::Deserialize;
16use serde_json::{Value, json};
17use uuid::Uuid;
18
19use crate::affinity::AffinityManager;
20use crate::error::{AuthError, LproviderError, WalletError};
21use crate::router::Router;
22use crate::router::{CallOutcomeFull, ChatContext, ChatSelectionSpec};
23use crate::state::AppState;
24use crate::wallet::WalletService;
25
26pub struct RouterContext {
27 pub state: Arc<AppState>,
28 pub router: Arc<Router>,
29 pub wallets: Arc<WalletService>,
30 pub affinity: Arc<AffinityManager>,
31}
32
33pub fn router(
34 state: Arc<AppState>,
35 router: Arc<Router>,
36 affinity: Arc<AffinityManager>,
37) -> AxumRouter {
38 let wallets = Arc::new(WalletService::new(state.db.clone().map(Arc::new)));
39 let ctx = Arc::new(RouterContext {
40 state,
41 router,
42 wallets,
43 affinity,
44 });
45 AxumRouter::new()
46 .route("/health/live", get(health_live))
47 .route("/health/ready", get(health_ready))
48 .route("/v1/models", get(list_models))
49 .route("/v1/chat/completions", post(chat_completions))
50 .with_state(ctx)
51}
52
53async fn health_live() -> impl IntoResponse {
54 (StatusCode::OK, Json(json!({ "status": "ok" })))
55}
56
57async fn health_ready(State(ctx): State<Arc<RouterContext>>) -> impl IntoResponse {
58 match ctx.state.ready().await {
59 Ok(()) => (
60 StatusCode::OK,
61 Json(json!({
62 "status": "ready",
63 "snapshot_version": ctx.state.snapshot_version(),
64 })),
65 ),
66 Err(reason) => (
67 StatusCode::SERVICE_UNAVAILABLE,
68 Json(json!({
69 "status": "not_ready",
70 "reason": reason,
71 })),
72 ),
73 }
74}
75
76async fn list_models(State(ctx): State<Arc<RouterContext>>) -> impl IntoResponse {
77 let runtime = ctx.state.runtime.read().await;
78 let lp = runtime.lp();
79
80 let created = ctx.state.started_at.unix_timestamp();
81
82 let mut capability_union: Vec<&str> = Vec::new();
83 let mut ctx_window = lp.max_context_tokens;
84 let mut max_out = lp.max_output_tokens;
85 for upstream in lp.upstreams.iter().filter(|u| u.enabled) {
86 for cap in &upstream.capabilities {
87 if !capability_union.contains(&cap.as_str()) {
88 capability_union.push(cap.as_str());
89 }
90 }
91 ctx_window = ctx_window.min(upstream.context_window);
92 max_out = max_out.min(upstream.max_output);
93 }
94
95 let data: Vec<_> = lp
96 .models
97 .iter()
98 .map(|m| {
99 json!({
100 "id": m.alias,
101 "object": "model",
102 "created": created,
103 "owned_by": "lprovider",
104 "description": m.description,
105 "default_strategy": m.default_strategy,
106 "context_window": ctx_window,
107 "max_output_tokens": max_out,
108 "capabilities": capability_union,
109 })
110 })
111 .collect();
112
113 (
114 StatusCode::OK,
115 Json(json!({ "object": "list", "data": data })),
116 )
117}
118
119#[derive(Deserialize)]
120struct ChatBody {
121 model: Option<String>,
122 messages: Vec<InMessage>,
123 #[serde(default)]
124 tools: Vec<InTool>,
125 max_tokens: Option<usize>,
126 temperature: Option<f32>,
127 top_p: Option<f32>,
128 response_format: Option<InResponseFormat>,
129 stop: Option<Vec<String>>,
130 #[serde(default)]
131 stream: bool,
132}
133
134#[derive(Deserialize)]
135struct InMessage {
136 role: String,
137 #[serde(default)]
138 content: Option<Value>,
139 #[serde(default)]
140 tool_calls: Option<Vec<InToolCall>>,
141 tool_call_id: Option<String>,
142 name: Option<String>,
143}
144
145#[derive(Deserialize)]
146struct InToolCall {
147 id: String,
148 function: InFunctionCall,
149}
150
151#[derive(Deserialize)]
152struct InFunctionCall {
153 name: String,
154 #[serde(default)]
155 arguments: Value,
156}
157
158#[derive(Deserialize)]
159struct InTool {
160 function: InToolDef,
161}
162
163#[derive(Deserialize)]
164struct InToolDef {
165 name: String,
166 description: Option<String>,
167 parameters: Option<Value>,
168}
169
170#[derive(Deserialize)]
171struct InResponseFormat {
172 #[serde(rename = "type")]
173 kind: Option<String>,
174}
175
176async fn chat_completions(
177 State(ctx): State<Arc<RouterContext>>,
178 headers: HeaderMap,
179 Json(body): Json<ChatBody>,
180) -> Response {
181 match build_chat_context(&ctx, &headers, &body).await {
182 Ok(chat_ctx) => {
183 let request_id = chat_ctx.request_id;
184 let model_alias = chat_ctx.model_alias.clone();
185 if body.stream {
186 match ctx.router.handle_stream(chat_ctx).await {
187 Ok(stream) => {
188 let request_id_str = request_id.to_string();
189 let sse_stream = stream.map(move |item| {
190 let data = match item {
191 Ok(ev) => sse_chunk(&ev, &request_id_str, &model_alias),
192 Err(e) => format!(
193 "data: {}\n\n",
194 json!({
195 "error": {
196 "message": e.to_string(),
197 "type": "stream_error",
198 }
199 })
200 ),
201 };
202 Ok::<Event, Infallible>(Event::default().data(data))
203 });
204 let mut resp = Sse::new(sse_stream)
205 .keep_alive(KeepAlive::default())
206 .into_response();
207 resp.headers_mut()
208 .insert("X-Request-Id", request_id.to_string().parse().unwrap());
209 resp.headers_mut().insert(
210 axum::http::header::CONTENT_TYPE,
211 "text/event-stream".parse().unwrap(),
212 );
213 resp
214 }
215 Err(e) => e.into_response(),
216 }
217 } else {
218 match ctx.router.handle_chat(chat_ctx).await {
219 Ok(outcome) => {
220 let mut resp =
221 (StatusCode::OK, Json(openai_response(&outcome))).into_response();
222 resp.headers_mut().insert(
223 "X-Request-Id",
224 outcome.request_id.to_string().parse().unwrap(),
225 );
226 resp
227 }
228 Err(e) => e.into_response(),
229 }
230 }
231 }
232 Err(e) => e.into_response(),
233 }
234}
235
236async fn build_chat_context(
237 ctx: &RouterContext,
238 headers: &HeaderMap,
239 body: &ChatBody,
240) -> Result<ChatContext, LproviderError> {
241 let lp = ctx.state.runtime.read().await.lp().clone();
242
243 if let Some(expected) = &lp.auth.bearer_token {
244 let provided = headers
245 .get(axum::http::header::AUTHORIZATION)
246 .and_then(|v| v.to_str().ok())
247 .and_then(|v| v.strip_prefix("Bearer "))
248 .unwrap_or("");
249 if provided != expected {
250 return Err(LproviderError::Auth(AuthError::InvalidToken));
251 }
252 }
253
254 let tenant_id = headers
255 .get("x-tenant-id")
256 .and_then(|v| v.to_str().ok())
257 .and_then(|v| v.trim().parse::<i64>().ok())
258 .ok_or_else(|| LproviderError::InvalidRequest("X-Tenant-Id header required".into()))?;
259
260 let wallet_id = headers
261 .get("x-wallet-id")
262 .and_then(|v| v.to_str().ok())
263 .and_then(|v| v.trim().parse::<i64>().ok())
264 .ok_or(LproviderError::Wallet(WalletError::WalletRequired))?;
265
266 let wallet = ctx.wallets.get_wallet(wallet_id).await?;
267 let subject_type = wallet.owner_type;
268 let subject_id = wallet.owner_id;
269
270 let session_key_override = match headers.get("x-session-key").and_then(|v| v.to_str().ok()) {
271 Some(key) => {
272 ctx.affinity.validate_session_key_override(key)?;
273 Some(key.to_string())
274 }
275 None => None,
276 };
277
278 let request_id = headers
279 .get("x-request-id")
280 .and_then(|v| v.to_str().ok())
281 .and_then(|v| Uuid::parse_str(v).ok())
282 .unwrap_or_else(Uuid::now_v7);
283
284 let client_ip = headers
285 .get("x-forwarded-for")
286 .and_then(|v| v.to_str().ok())
287 .and_then(|v| v.split(',').next())
288 .map(|s| s.trim().to_string());
289
290 let user_agent = headers
291 .get(axum::http::header::USER_AGENT)
292 .and_then(|v| v.to_str().ok())
293 .map(|s| s.to_string());
294
295 let model_alias = body.model.clone().unwrap_or_else(|| "auto".to_string());
296 let request = map_chat_request(body)?;
297 let max_output = body
298 .max_tokens
299 .unwrap_or(lp.max_output_tokens)
300 .min(lp.max_output_tokens);
301
302 Ok(ChatContext {
303 request_id,
304 trace_id: None,
305 tenant_id,
306 wallet_id: Some(wallet_id),
307 subject_type,
308 subject_id,
309 anonymous: false,
310 model_alias,
311 request,
312 max_output,
313 stream: body.stream,
314 priority: 0,
315 caller: None,
316 caller_service: None,
317 session_key_override,
318 selection: ChatSelectionSpec::default(),
319 client_ip,
320 user_agent,
321 })
322}
323
324fn map_chat_request(body: &ChatBody) -> Result<ChatRequest, LproviderError> {
325 let mut messages = Vec::with_capacity(body.messages.len());
326 for m in &body.messages {
327 let content = map_content(m.content.clone())?;
328 let tool_calls = map_tool_calls(m.tool_calls.as_ref());
329 let msg = match m.role.as_str() {
330 "system" => ChatMessage {
331 role: MessageRole::System,
332 content,
333 name: m.name.clone(),
334 tool_calls: Vec::new(),
335 tool_call_id: None,
336 },
337 "user" => ChatMessage {
338 role: MessageRole::User,
339 content,
340 name: m.name.clone(),
341 tool_calls: Vec::new(),
342 tool_call_id: None,
343 },
344 "assistant" => ChatMessage {
345 role: MessageRole::Assistant,
346 content,
347 name: m.name.clone(),
348 tool_calls,
349 tool_call_id: None,
350 },
351 "tool" => ChatMessage {
352 role: MessageRole::Tool,
353 content,
354 name: m.name.clone(),
355 tool_calls: Vec::new(),
356 tool_call_id: m.tool_call_id.clone(),
357 },
358 other => {
359 return Err(LproviderError::InvalidRequest(format!(
360 "unsupported message role: {other}"
361 )));
362 }
363 };
364 messages.push(msg);
365 }
366
367 let tools: Vec<ToolDefinition> = body
368 .tools
369 .iter()
370 .map(|t| ToolDefinition {
371 name: t.function.name.clone(),
372 description: t.function.description.clone().unwrap_or_default(),
373 parameters: t
374 .function
375 .parameters
376 .clone()
377 .unwrap_or_else(|| json!({ "type": "object", "properties": {} })),
378 concurrency: Default::default(),
379 })
380 .collect();
381
382 let mut request = ChatRequest::new(messages);
383 if !tools.is_empty() {
384 request = request.with_tools(tools);
385 }
386 if let Some(mt) = body.max_tokens {
387 request = request.with_max_tokens(mt);
388 }
389 if let Some(t) = body.temperature {
390 request = request.with_temperature(t);
391 }
392 if let Some(p) = body.top_p {
393 request = request.with_top_p(p);
394 }
395 if body
396 .response_format
397 .as_ref()
398 .and_then(|rf| rf.kind.as_deref())
399 == Some("json_object")
400 {
401 request = request.with_response_format(ResponseFormat::Json);
402 }
403 if let Some(stop) = &body.stop {
404 request = request.with_stop_sequences(stop.clone());
405 }
406 Ok(request)
407}
408
409fn map_content(v: Option<Value>) -> Result<MessageContent, LproviderError> {
410 match v {
411 None => Ok(MessageContent::Text(String::new())),
412 Some(Value::String(s)) => Ok(MessageContent::Text(s)),
413 Some(Value::Array(parts)) => {
414 let mut out = Vec::with_capacity(parts.len());
415 for p in parts {
416 match p.get("type").and_then(|t| t.as_str()) {
417 Some("text") => {
418 if let Some(t) = p.get("text").and_then(|t| t.as_str()) {
419 out.push(ContentPart::Text(t.to_string()));
420 }
421 }
422 Some("image_url") => {
423 let url = p
424 .get("image_url")
425 .and_then(|i| i.get("url"))
426 .and_then(|u| u.as_str())
427 .unwrap_or_default()
428 .to_string();
429 let detail = match p
430 .get("image_url")
431 .and_then(|i| i.get("detail"))
432 .and_then(|d| d.as_str())
433 .unwrap_or("auto")
434 {
435 "low" => ImageDetail::Low,
436 "high" => ImageDetail::High,
437 _ => ImageDetail::Auto,
438 };
439 out.push(ContentPart::Image { url, detail });
440 }
441 _ => {}
442 }
443 }
444 Ok(MessageContent::Parts(out))
445 }
446 Some(_) => Err(LproviderError::InvalidRequest(
447 "invalid message content".into(),
448 )),
449 }
450}
451
452fn map_tool_calls(tc: Option<&Vec<InToolCall>>) -> Vec<ToolCall> {
453 tc.map(|list| {
454 list.iter()
455 .map(|t| ToolCall {
456 id: t.id.clone(),
457 name: t.function.name.clone(),
458 arguments: match &t.function.arguments {
459 Value::String(s) => serde_json::from_str::<Value>(s).unwrap_or(Value::Null),
460 other => other.clone(),
461 },
462 })
463 .collect()
464 })
465 .unwrap_or_default()
466}
467
468fn finish_reason_str(fr: &FinishReason) -> &'static str {
469 match fr {
470 FinishReason::Stop => "stop",
471 FinishReason::ToolUse => "tool_calls",
472 FinishReason::MaxTokens => "length",
473 FinishReason::ContentFilter => "content_filter",
474 FinishReason::Error => "error",
475 }
476}
477
478fn created_now() -> u64 {
479 std::time::SystemTime::now()
480 .duration_since(std::time::UNIX_EPOCH)
481 .unwrap_or_default()
482 .as_secs()
483}
484
485fn openai_response(o: &CallOutcomeFull) -> Value {
486 let mut message = json!({
487 "role": "assistant",
488 "content": o.response.content,
489 });
490 if !o.response.tool_calls.is_empty() {
491 message["tool_calls"] = Value::Array(
492 o.response
493 .tool_calls
494 .iter()
495 .map(|t| {
496 json!({
497 "id": t.id,
498 "type": "function",
499 "function": {
500 "name": t.name,
501 "arguments": t.arguments.to_string(),
502 },
503 })
504 })
505 .collect(),
506 );
507 }
508 json!({
509 "id": format!("chatcmpl-{}", o.request_id),
510 "object": "chat.completion",
511 "created": created_now(),
512 "model": o.model,
513 "choices": [{
514 "index": 0,
515 "message": message,
516 "finish_reason": finish_reason_str(&o.response.finish_reason),
517 }],
518 "usage": {
519 "prompt_tokens": o.response.usage.input_tokens,
520 "completion_tokens": o.response.usage.output_tokens,
521 "total_tokens": o.response.usage.total(),
522 "prompt_cache_hit_tokens": o.bill.billable_cache_hit_tokens,
523 },
524 })
525}
526
527fn sse_chunk(ev: &StreamEvent, id: &str, model: &str) -> String {
528 let base = || {
529 json!({
530 "id": id,
531 "object": "chat.completion.chunk",
532 "created": created_now(),
533 "model": model,
534 })
535 };
536 let mut chunk = base();
537 match ev {
538 StreamEvent::TextDelta { delta } => {
539 chunk["choices"] = json!([{
540 "index": 0,
541 "delta": { "content": delta },
542 "finish_reason": Value::Null,
543 }]);
544 }
545 StreamEvent::ReasoningDelta { delta } => {
546 chunk["choices"] = json!([{
547 "index": 0,
548 "delta": { "reasoning_content": delta },
549 "finish_reason": Value::Null,
550 }]);
551 }
552 StreamEvent::ToolCallDelta {
553 index,
554 id: call_id,
555 name,
556 args_delta,
557 } => {
558 chunk["choices"] = json!([{
559 "index": 0,
560 "delta": {
561 "tool_calls": [{
562 "index": index.unwrap_or(0),
563 "id": call_id,
564 "type": "function",
565 "function": {
566 "name": name.clone().unwrap_or_default(),
567 "arguments": args_delta,
568 },
569 }],
570 },
571 "finish_reason": Value::Null,
572 }]);
573 }
574 StreamEvent::Usage(u) => {
575 chunk["choices"] = json!([]);
576 chunk["usage"] = json!({
577 "prompt_tokens": u.input_tokens,
578 "completion_tokens": u.output_tokens,
579 "total_tokens": u.total(),
580 });
581 }
582 StreamEvent::Done { finish_reason } => {
583 chunk["choices"] = json!([{
584 "index": 0,
585 "delta": {},
586 "finish_reason": finish_reason_str(finish_reason),
587 }]);
588 }
589 }
590 format!("data: {chunk}\n\n")
591}
Modified agent/lprovider/lib.rs agent/lprovider/lib.rs +28 −0 View agent/lprovider/lib.rs
@@ -1 +1,29 @@
1pub mod admin_grpc;
2pub mod affinity;
3pub mod ai_grpc;
4pub mod audit;
5pub mod billing;
6pub mod config;
7pub mod credentials;
8pub mod error;
9pub mod http;
10pub mod quota;
11pub mod rotation;
12pub mod router;
13pub mod selector;
14pub mod state;
15pub mod upstream;
16pub mod wallet;
117
18pub mod lprovider {
19 pub mod ai {
20 pub mod v1 {
21 tonic::include_proto!("lprovider.ai.v1");
22 }
23 }
24 pub mod admin {
25 pub mod v1 {
26 tonic::include_proto!("lprovider.admin.v1");
27 }
28 }
29}
Added agent/lprovider/main.rs agent/lprovider/main.rs +363 −0 View agent/lprovider/main.rs
@@ -0,0 +1,363 @@
1use std::path::PathBuf;
2use std::sync::Arc;
3use std::time::Duration;
4
5use cache::RedisCache;
6use clap::Parser;
7use db::conn::Db;
8use redis::AsyncCommands;
9use tonic_health::ServingStatus;
10use tonic_health::server::health_reporter;
11use tracing::{info, warn};
12use tracing_subscriber::EnvFilter;
13
14use lprovider::admin_grpc::AdminGrpc;
15use lprovider::affinity::AffinityManager;
16use lprovider::ai_grpc::AiGrpc;
17use lprovider::audit::AuditService;
18use lprovider::config::{CliOverrides, RuntimeConfig};
19use lprovider::http;
20use lprovider::lprovider::admin::v1::lprovider_admin_server::LproviderAdminServer;
21use lprovider::lprovider::ai::v1::lprovider_ai_server::LproviderAiServer;
22use lprovider::router::Router;
23use lprovider::state::AppState;
24use lprovider::upstream::UpstreamManager;
25
26#[derive(Parser, Debug)]
27#[command(name = "lprovider", about = "LLM Provider Gateway (OpenAI-compatible)")]
28struct Cli {
29 /// 主配置文件路径
30 #[arg(long, default_value = "config/default.yaml")]
31 config: PathBuf,
32
33 /// HTTP 服务绑定地址（覆盖配置文件）
34 #[arg(long)]
35 bind: Option<String>,
36
37 /// HTTP 服务端口（覆盖配置文件）
38 #[arg(long)]
39 port: Option<u16>,
40
41 /// Admin gRPC 服务绑定地址（覆盖配置文件）
42 #[arg(long)]
43 admin_bind: Option<String>,
44
45 /// Admin gRPC 服务端口（覆盖配置文件）
46 #[arg(long)]
47 admin_port: Option<u16>,
48
49 /// Admin gRPC mTLS 服务端证书（PEM，与 --admin-key 成对提供）
50 #[arg(long)]
51 admin_cert: Option<PathBuf>,
52
53 /// Admin gRPC mTLS 服务端私钥（PEM，与 --admin-cert 成对提供）
54 #[arg(long)]
55 admin_key: Option<PathBuf>,
56
57 /// 实例 id（用于集群注册与审计）
58 #[arg(long)]
59 instance_id: Option<String>,
60
61 /// 集群名（用于 Redis 实例注册键）
62 #[arg(long)]
63 cluster_name: Option<String>,
64}
65
66#[tokio::main]
67async fn main() -> Result<(), Box<dyn std::error::Error>> {
68 tracing_subscriber::fmt()
69 .with_env_filter(
70 EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
71 )
72 .init();
73
74 let cli = Cli::parse();
75 let cli_overrides = CliOverrides {
76 config: cli.config,
77 bind: cli.bind,
78 port: cli.port,
79 admin_bind: cli.admin_bind,
80 admin_port: cli.admin_port,
81 admin_cert: cli.admin_cert,
82 admin_key: cli.admin_key,
83 instance_id: cli.instance_id,
84 cluster_name: cli.cluster_name,
85 };
86
87 let runtime = RuntimeConfig::load(&cli_overrides)?;
88
89 let instance_id = cli_overrides
90 .instance_id
91 .clone()
92 .unwrap_or_else(|| format!("{}-{}", hostname(), std::process::id()));
93 let cluster_name = cli_overrides
94 .cluster_name
95 .clone()
96 .unwrap_or_else(|| "default".to_string());
97
98 let db = match Db::from_config(&runtime.app.db).await {
99 Ok(db) => {
100 info!("database connected, schema synced");
101 Some(db)
102 }
103 Err(e) => {
104 warn!(error = %e, "database connection failed; /health/ready will report not ready");
105 None
106 }
107 };
108
109 let redis: Option<Arc<RedisCache>> = match RedisCache::new(&runtime.app.cache).await {
110 Ok(redis) => {
111 info!("redis connected");
112 Some(Arc::new(redis))
113 }
114 Err(e) => {
115 warn!(error = %e, "redis connection failed; /health/ready will report not ready");
116 None
117 }
118 };
119
120 let state = AppState::new(
121 runtime,
122 db,
123 redis.clone(),
124 instance_id.clone(),
125 cluster_name.clone(),
126 cli_overrides.clone(),
127 );
128
129 let lp = state.runtime.blocking_read().lp().clone();
130 let manager = Arc::new(UpstreamManager::new(&lp, redis.clone())?);
131 let affinity = Arc::new(AffinityManager::new(redis.clone(), &lp.redis));
132 let router = Router::new(state.clone(), manager.clone(), affinity.clone());
133 let audit = Arc::new(AuditService::new(state.db.clone().map(Arc::new)));
134
135 register_instance(&state).await;
136
137 let (health_reporter, health_service) = health_reporter();
138 health_reporter
139 .set_service_status("lprovider.ai.v1.LproviderAI", ServingStatus::Serving)
140 .await;
141 health_reporter
142 .set_service_status("lprovider.admin.v1.LproviderAdmin", ServingStatus::Serving)
143 .await;
144 let health_service_admin = health_service.clone();
145
146 let shutdown_token = state.shutdown.clone();
147
148 let http_app = http::router(state.clone(), router.clone(), affinity.clone());
149 let http_addr = format!("{}:{}", lp.bind, lp.port);
150 let http_listener = tokio::net::TcpListener::bind(&http_addr).await?;
151 let http_token = shutdown_token.clone();
152 let http_task = tokio::spawn(async move {
153 let _ = axum::serve(http_listener, http_app)
154 .with_graceful_shutdown(async move { http_token.cancelled().await })
155 .await;
156 });
157
158 let ai_grpc = AiGrpc::new(router.clone(), state.clone());
159 let ai_addr = format!("{}:9090", lp.bind);
160 let ai_addr_log = ai_addr.clone();
161 let ai_token = shutdown_token.clone();
162 let ai_task = tokio::spawn(async move {
163 let svc = LproviderAiServer::new(ai_grpc);
164 let _ = tonic::transport::Server::builder()
165 .add_service(health_service)
166 .add_service(svc)
167 .serve_with_shutdown(ai_addr.parse().expect("valid ai grpc addr"), async move {
168 ai_token.cancelled().await
169 })
170 .await;
171 });
172
173 let admin_grpc = AdminGrpc::new(state.clone(), manager.clone(), affinity.clone());
174 let admin_addr = format!("{}:{}", lp.admin_bind, lp.admin_port);
175 let admin_addr_log = admin_addr.clone();
176 let admin_token = shutdown_token.clone();
177 let admin_tls = build_admin_tls(&cli_overrides);
178 let admin_task = tokio::spawn(async move {
179 let svc = LproviderAdminServer::new(admin_grpc);
180 let mut builder = tonic::transport::Server::builder();
181 if let Some(tls) = admin_tls {
182 builder = builder.tls_config(tls).expect("valid admin tls config");
183 }
184 let _ = builder
185 .add_service(health_service_admin)
186 .add_service(svc)
187 .serve_with_shutdown(
188 admin_addr.parse().expect("valid admin grpc addr"),
189 async move { admin_token.cancelled().await },
190 )
191 .await;
192 });
193
194 let aggregate_task = tokio::spawn(aggregate_worker(audit.clone(), shutdown_token.clone()));
195 let reconcile_task = tokio::spawn(reconcile_worker(audit.clone(), shutdown_token.clone()));
196 let cleanup_task = tokio::spawn(cleanup_worker(affinity.clone(), shutdown_token.clone()));
197
198 info!(
199 instance_id = %instance_id,
200 cluster_name = %cluster_name,
201 http_addr = %http_addr,
202 ai_grpc_addr = %ai_addr_log,
203 admin_grpc_addr = %admin_addr_log,
204 "lprovider started"
205 );
206
207 tokio::select! {
208 _ = shutdown_signal() => {
209 info!("shutdown signal received, draining connections");
210 state.shutdown.cancel();
211 }
212 _ = state.shutdown.cancelled() => {
213 info!("shutdown requested via admin");
214 }
215 }
216
217 let _ = tokio::join!(
218 http_task,
219 ai_task,
220 admin_task,
221 aggregate_task,
222 reconcile_task,
223 cleanup_task,
224 );
225
226 Ok(())
227}
228
229async fn register_instance(state: &AppState) {
230 let Some(redis) = &state.redis else {
231 return;
232 };
233 let runtime = state.runtime.blocking_read();
234 let lp = runtime.lp().clone();
235 let json = serde_json::json!({
236 "instance_id": state.instance_id,
237 "cluster_name": state.cluster_name,
238 "listen_addr": format!("{}:{}", lp.bind, lp.port),
239 "started_at_unix": state.started_at.unix_timestamp(),
240 "pid": state.pid,
241 "config_sha256": runtime.config_sha256,
242 "is_config_writer": false,
243 })
244 .to_string();
245 let prefix = lp.redis.key_prefix.clone();
246 drop(runtime);
247 let mut conn = redis.connection();
248 let key = format!("{prefix}:instances:{}", state.cluster_name);
249 let _: redis::RedisResult<i64> = conn.sadd(&key, &json).await;
250 let _: redis::RedisResult<()> = conn.expire(&key, 86400).await;
251 let sha_key = format!("{prefix}:config_sha:{}", state.instance_id);
252 let _: redis::RedisResult<()> = conn
253 .set(sha_key, state.runtime.blocking_read().config_sha256.clone())
254 .await;
255}
256
257async fn aggregate_worker(audit: Arc<AuditService>, token: tokio_util::sync::CancellationToken) {
258 loop {
259 tokio::select! {
260 _ = tokio::time::sleep(Duration::from_secs(300)) => {
261 match audit.aggregate_incremental().await {
262 Ok(n) => info!(rows = n, "billing aggregate incremental done"),
263 Err(e) => warn!(error = %e, "billing aggregate incremental failed"),
264 }
265 }
266 _ = token.cancelled() => break,
267 }
268 }
269}
270
271async fn reconcile_worker(audit: Arc<AuditService>, token: tokio_util::sync::CancellationToken) {
272 loop {
273 tokio::select! {
274 _ = tokio::time::sleep(until_next_0030()) => {
275 match audit.reconcile_daily().await {
276 Ok(ratio) => info!(ratio, "daily reconcile done"),
277 Err(e) => warn!(error = %e, "daily reconcile failed"),
278 }
279 }
280 _ = token.cancelled() => break,
281 }
282 }
283}
284
285async fn cleanup_worker(
286 affinity: Arc<AffinityManager>,
287 token: tokio_util::sync::CancellationToken,
288) {
289 loop {
290 tokio::select! {
291 _ = tokio::time::sleep(Duration::from_secs(60)) => {
292 match affinity.cleanup_once().await {
293 Ok(n) if n > 0 => info!(removed = n, "affinity binding cleanup"),
294 Ok(_) => {}
295 Err(e) => warn!(error = %e, "affinity binding cleanup failed"),
296 }
297 }
298 _ = token.cancelled() => break,
299 }
300 }
301}
302
303fn until_next_0030() -> Duration {
304 let now = time::OffsetDateTime::now_utc();
305 let today = now.date();
306 let mut target = today.with_hms(0, 30, 0).expect("valid time").assume_utc();
307 if target <= now {
308 target += time::Duration::days(1);
309 }
310 let diff = target - now;
311 Duration::from_secs(diff.whole_seconds().max(1) as u64)
312}
313
314fn build_admin_tls(cli: &CliOverrides) -> Option<tonic::transport::ServerTlsConfig> {
315 let (cert, key) = match (&cli.admin_cert, &cli.admin_key) {
316 (Some(c), Some(k)) => (c, k),
317 _ => {
318 warn!(
319 "--admin-cert/--admin-key not both provided; admin gRPC serving without mTLS (dev mode)"
320 );
321 return None;
322 }
323 };
324 match secret::load_identity(cert, key) {
325 Ok(id) => {
326 let identity = tonic::transport::Identity::from_pem(id.cert_pem, id.key_pem);
327 Some(tonic::transport::ServerTlsConfig::new().identity(identity))
328 }
329 Err(e) => {
330 warn!(error = %e, "failed to load admin mTLS identity; serving without mTLS");
331 None
332 }
333 }
334}
335
336fn hostname() -> String {
337 std::env::var("HOSTNAME")
338 .ok()
339 .filter(|s| !s.is_empty())
340 .unwrap_or_else(|| "localhost".to_string())
341}
342
343async fn shutdown_signal() {
344 let ctrl_c = async {
345 tokio::signal::ctrl_c().await.ok();
346 };
347
348 #[cfg(unix)]
349 let terminate = async {
350 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
351 .expect("failed to install SIGTERM handler")
352 .recv()
353 .await;
354 };
355
356 #[cfg(not(unix))]
357 let terminate = std::future::pending::<()>();
358
359 tokio::select! {
360 _ = ctrl_c => {},
361 _ = terminate => {},
362 }
363}
Added agent/lprovider/quota.rs agent/lprovider/quota.rs +258 −0 View agent/lprovider/quota.rs
@@ -0,0 +1,258 @@
1use std::sync::Arc;
2
3use db::conn::Db;
4use db::lprovider::llm_tenant_quota;
5use sea_orm::{
6 ColumnTrait, Condition, ConnectionTrait, DatabaseBackend, EntityTrait, PaginatorTrait,
7 QueryFilter, QueryOrder, QuerySelect, Statement, Value,
8};
9use tracing::{info, warn};
10
11use crate::error::LproviderError;
12
13#[derive(Clone, Debug)]
14pub struct TenantQuota {
15 pub tenant_id: i64,
16 pub subject_type: String,
17 pub subject_id: i64,
18 pub monthly_quota_points: i64,
19 pub monthly_quota_input_tokens: i64,
20 pub monthly_quota_output_tokens: i64,
21 pub hard_limit: bool,
22 pub enabled: bool,
23}
24
25#[derive(Clone)]
26pub struct QuotaService {
27 db: Option<Arc<Db>>,
28}
29
30impl QuotaService {
31 pub fn new(db: Option<Arc<Db>>) -> Self {
32 Self { db }
33 }
34
35 fn db(&self) -> Result<&Db, LproviderError> {
36 self.db.as_deref().ok_or_else(|| {
37 LproviderError::Db(db::conn::DbError::InvalidConfig("db not configured".into()))
38 })
39 }
40
41 pub async fn get_quota(
42 &self,
43 tenant_id: i64,
44 subject_type: &str,
45 subject_id: i64,
46 ) -> Result<Option<TenantQuota>, LproviderError> {
47 let db = self.db()?;
48 let model = llm_tenant_quota::Entity::find()
49 .filter(
50 Condition::all()
51 .add(llm_tenant_quota::Column::TenantId.eq(tenant_id))
52 .add(llm_tenant_quota::Column::SubjectType.eq(subject_type))
53 .add(llm_tenant_quota::Column::SubjectId.eq(subject_id)),
54 )
55 .one(db.primary_connection())
56 .await
57 .map_err(db_err)?;
58 Ok(model.map(|m| TenantQuota {
59 tenant_id: m.tenant_id,
60 subject_type: m.subject_type,
61 subject_id: m.subject_id,
62 monthly_quota_points: m.monthly_quota_points,
63 monthly_quota_input_tokens: m.monthly_quota_input_tokens,
64 monthly_quota_output_tokens: m.monthly_quota_output_tokens,
65 hard_limit: m.hard_limit,
66 enabled: m.enabled,
67 }))
68 }
69
70 pub async fn list_quotas(
71 &self,
72 tenant_id: Option<i64>,
73 subject_type: Option<String>,
74 subject_id: Option<i64>,
75 limit: u64,
76 offset: u64,
77 ) -> Result<(Vec<TenantQuota>, u64), LproviderError> {
78 let db = self.db()?;
79 let mut cond = sea_orm::Condition::all();
80 if let Some(t) = tenant_id {
81 cond = cond.add(llm_tenant_quota::Column::TenantId.eq(t));
82 }
83 if let Some(st) = subject_type {
84 cond = cond.add(llm_tenant_quota::Column::SubjectType.eq(st));
85 }
86 if let Some(si) = subject_id {
87 cond = cond.add(llm_tenant_quota::Column::SubjectId.eq(si));
88 }
89 let total = llm_tenant_quota::Entity::find()
90 .filter(cond.clone())
91 .count(db.primary_connection())
92 .await
93 .map_err(db_err)?;
94 let models = llm_tenant_quota::Entity::find()
95 .filter(cond)
96 .order_by_asc(llm_tenant_quota::Column::TenantId)
97 .limit(limit)
98 .offset(offset)
99 .all(db.primary_connection())
100 .await
101 .map_err(db_err)?;
102 Ok((
103 models
104 .into_iter()
105 .map(|m| TenantQuota {
106 tenant_id: m.tenant_id,
107 subject_type: m.subject_type,
108 subject_id: m.subject_id,
109 monthly_quota_points: m.monthly_quota_points,
110 monthly_quota_input_tokens: m.monthly_quota_input_tokens,
111 monthly_quota_output_tokens: m.monthly_quota_output_tokens,
112 hard_limit: m.hard_limit,
113 enabled: m.enabled,
114 })
115 .collect(),
116 total,
117 ))
118 }
119
120 pub async fn upsert_quota(&self, quota: &TenantQuota) -> Result<(), LproviderError> {
121 let db = self.db()?;
122 db.writer()
123 .execute_raw(Statement::from_sql_and_values(
124 DatabaseBackend::Postgres,
125 "INSERT INTO llm_tenant_quota (tenant_id, subject_type, subject_id, monthly_quota_points, monthly_quota_input_tokens, monthly_quota_output_tokens, hard_limit, enabled, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now(), now()) ON CONFLICT (tenant_id, subject_type, subject_id) DO UPDATE SET monthly_quota_points = $4, monthly_quota_input_tokens = $5, monthly_quota_output_tokens = $6, hard_limit = $7, enabled = $8, updated_at = now()",
126 [
127 Value::BigInt(Some(quota.tenant_id)),
128 Value::String(Some(quota.subject_type.clone())),
129 Value::BigInt(Some(quota.subject_id)),
130 Value::BigInt(Some(quota.monthly_quota_points)),
131 Value::BigInt(Some(quota.monthly_quota_input_tokens)),
132 Value::BigInt(Some(quota.monthly_quota_output_tokens)),
133 Value::Bool(Some(quota.hard_limit)),
134 Value::Bool(Some(quota.enabled)),
135 ],
136 ))
137 .await
138 .map_err(db_err)?;
139 Ok(())
140 }
141
142 async fn used_this_month(&self, quota: &TenantQuota) -> Result<i64, LproviderError> {
143 let db = self.db()?;
144 let row = db
145 .writer()
146 .query_one_raw(Statement::from_sql_and_values(
147 DatabaseBackend::Postgres,
148 "SELECT COALESCE(SUM(points_total), 0) AS used_points FROM llm_call_log WHERE tenant_id = $1 AND subject_type = $2 AND subject_id = $3 AND created_at >= date_trunc('month', now())",
149 [
150 Value::BigInt(Some(quota.tenant_id)),
151 Value::String(Some(quota.subject_type.clone())),
152 Value::BigInt(Some(quota.subject_id)),
153 ],
154 ))
155 .await
156 .map_err(db_err)?;
157 Ok(row
158 .and_then(|r| r.try_get("", "used_points").ok())
159 .unwrap_or(0))
160 }
161
162 pub async fn check(
163 &self,
164 tenant_id: i64,
165 subject_type: &str,
166 subject_id: i64,
167 estimated_cost: i64,
168 ) -> Result<(), LproviderError> {
169 let Some(quota) = self.get_quota(tenant_id, subject_type, subject_id).await? else {
170 return Ok(());
171 };
172 if !quota.enabled {
173 return Ok(());
174 }
175 if quota.monthly_quota_points == 0 {
176 return Ok(());
177 }
178 let used = self.used_this_month("a).await?;
179 if used + estimated_cost > quota.monthly_quota_points {
180 if quota.hard_limit {
181 return Err(LproviderError::QuotaExceeded(format!(
182 "monthly quota exceeded for {subject_type}:{subject_id}: used {used}, quota {}",
183 quota.monthly_quota_points
184 )));
185 }
186 warn!(
187 tenant_id,
188 subject_type,
189 subject_id,
190 used,
191 quota_points = quota.monthly_quota_points,
192 "llm quota soft-limit exceeded (warning only)"
193 );
194 }
195 Ok(())
196 }
197
198 pub async fn increment_used_tokens(
199 &self,
200 tenant_id: i64,
201 subject_type: Option<&str>,
202 subject_id: Option<i64>,
203 input_tokens: i64,
204 output_tokens: i64,
205 ) -> Result<(), LproviderError> {
206 let (Some(st), Some(si)) = (subject_type, subject_id) else {
207 return Ok(());
208 };
209 let Some(quota) = self.get_quota(tenant_id, st, si).await? else {
210 return Ok(());
211 };
212 if !quota.enabled {
213 return Ok(());
214 }
215 if quota.monthly_quota_input_tokens == 0 && quota.monthly_quota_output_tokens == 0 {
216 return Ok(());
217 }
218 let input_used: i64 = self
219 .db()?
220 .writer()
221 .query_one_raw(Statement::from_sql_and_values(
222 DatabaseBackend::Postgres,
223 "SELECT COALESCE(SUM(prompt_tokens), 0) AS s FROM llm_call_log WHERE tenant_id = $1 AND subject_type = $2 AND subject_id = $3 AND created_at >= date_trunc('month', now())",
224 [
225 Value::BigInt(Some(tenant_id)),
226 Value::String(Some(st.to_string())),
227 Value::BigInt(Some(si)),
228 ],
229 ))
230 .await
231 .map_err(db_err)?
232 .and_then(|r| r.try_get("", "s").ok())
233 .unwrap_or(0);
234 if quota.monthly_quota_input_tokens > 0
235 && input_used + input_tokens > quota.monthly_quota_input_tokens
236 {
237 if quota.hard_limit {
238 return Err(LproviderError::QuotaExceeded(format!(
239 "monthly input token quota exceeded for {st}:{si}"
240 )));
241 }
242 warn!(tenant_id, "llm input token quota soft-limit exceeded");
243 }
244 info!(
245 tenant_id,
246 input_used, output_tokens, "llm quota token accounting"
247 );
248 Ok(())
249 }
250
251 pub async fn aggregate_month_used(&self, quota: &TenantQuota) -> Result<i64, LproviderError> {
252 self.used_this_month(quota).await
253 }
254}
255
256fn db_err(err: sea_orm::DbErr) -> LproviderError {
257 LproviderError::Db(db::conn::DbError::InvalidConfig(format!("sea-orm: {err}")))
258}
Added agent/lprovider/rotation.rs agent/lprovider/rotation.rs +178 −0 View agent/lprovider/rotation.rs
@@ -0,0 +1,178 @@
1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use futures::Stream;
5use llm::error::LlmError;
6use llm::types::{ChatRequest, ChatResponse, StreamEvent};
7use tokio_stream::StreamExt;
8
9use crate::error::{LproviderError, UpstreamError};
10use crate::selector::Selected;
11use crate::upstream::{KeyLease, UpstreamManager, UpstreamSnapshot};
12
13pub const MAX_TOTAL_ATTEMPTS: usize = 3;
14
15pub struct CallOutcome {
16 pub response: ChatResponse,
17 pub upstream_idx: usize,
18 pub key_idx: usize,
19 pub upstream_code: String,
20 pub key_label: String,
21 pub model: String,
22}
23
24#[derive(Clone)]
25pub struct StreamMeta {
26 pub upstream_idx: usize,
27 pub key_idx: usize,
28 pub upstream_code: String,
29 pub key_label: String,
30 pub model: String,
31}
32
33#[allow(clippy::type_complexity)]
34pub struct RotatingStream {
35 inner: Option<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>>,
36 first: Option<Result<StreamEvent, LlmError>>,
37 _lease: Option<KeyLease>,
38 pub meta: Option<StreamMeta>,
39}
40
41impl Stream for RotatingStream {
42 type Item = Result<StreamEvent, LlmError>;
43
44 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
45 if let Some(first) = self.first.take() {
46 return Poll::Ready(Some(first));
47 }
48 match self.inner.as_mut() {
49 None => Poll::Ready(None),
50 Some(inner) => inner.as_mut().poll_next(cx),
51 }
52 }
53}
54
55pub fn map_llm_error(e: &LlmError) -> LproviderError {
56 match e {
57 LlmError::ContextLengthExceeded { .. } => {
58 LproviderError::Upstream(UpstreamError::ContextLengthExceeded)
59 }
60 other => {
61 LproviderError::Upstream(UpstreamError::AllFailed(format!("上游调用失败: {other}")))
62 }
63 }
64}
65
66pub async fn chat_with_rotation(
67 manager: &UpstreamManager,
68 snapshot: &UpstreamSnapshot,
69 ranked: &[Selected],
70 request: &ChatRequest,
71) -> Result<CallOutcome, LproviderError> {
72 let mut failures = 0usize;
73 for sel in ranked {
74 let Some(up) = snapshot.upstreams.get(sel.upstream_idx) else {
75 continue;
76 };
77 loop {
78 if failures >= MAX_TOTAL_ATTEMPTS {
79 return Err(LproviderError::Upstream(UpstreamError::AllFailed(format!(
80 "连续 {failures} 次可重试失败，放弃轮转"
81 ))));
82 }
83 let Some(key_idx) = manager.select_key(up).await else {
84 break;
85 };
86 let key = &up.keys[key_idx];
87 let Some(_lease) = manager.begin_key(sel.upstream_idx, key_idx).await else {
88 break;
89 };
90 match key.provider.chat(request.clone()).await {
91 Ok(response) => {
92 return Ok(CallOutcome {
93 response,
94 upstream_idx: sel.upstream_idx,
95 key_idx,
96 upstream_code: up.code.clone(),
97 key_label: key.label.clone(),
98 model: sel.model.clone(),
99 });
100 }
101 Err(e) if e.is_retryable() => {
102 failures += 1;
103 manager.mark_cooldown(&up.code, &key.label).await;
104 }
105 Err(e) => return Err(map_llm_error(&e)),
106 }
107 }
108 }
109 Err(LproviderError::Upstream(UpstreamError::AllFailed(
110 "无可用上游 Key 可完成调用".into(),
111 )))
112}
113
114pub async fn chat_stream_with_rotation(
115 manager: &UpstreamManager,
116 snapshot: &UpstreamSnapshot,
117 ranked: &[Selected],
118 request: &ChatRequest,
119) -> Result<RotatingStream, LproviderError> {
120 let mut failures = 0usize;
121 for sel in ranked {
122 let Some(up) = snapshot.upstreams.get(sel.upstream_idx) else {
123 continue;
124 };
125 loop {
126 if failures >= MAX_TOTAL_ATTEMPTS {
127 return Err(LproviderError::Upstream(UpstreamError::AllFailed(format!(
128 "连续 {failures} 次可重试失败，放弃轮转"
129 ))));
130 }
131 let Some(key_idx) = manager.select_key(up).await else {
132 break;
133 };
134 let key = &up.keys[key_idx];
135 let Some(lease) = manager.begin_key(sel.upstream_idx, key_idx).await else {
136 break;
137 };
138 let stream = match key.provider.chat_stream(request.clone()).await {
139 Ok(stream) => stream,
140 Err(e) if e.is_retryable() => {
141 failures += 1;
142 manager.mark_cooldown(&up.code, &key.label).await;
143 continue;
144 }
145 Err(e) => return Err(map_llm_error(&e)),
146 };
147 let mut stream = stream;
148 match stream.next().await {
149 Some(Ok(event)) => {
150 return Ok(RotatingStream {
151 inner: Some(stream),
152 first: Some(Ok(event)),
153 _lease: Some(lease),
154 meta: Some(StreamMeta {
155 upstream_idx: sel.upstream_idx,
156 key_idx,
157 upstream_code: up.code.clone(),
158 key_label: key.label.clone(),
159 model: sel.model.clone(),
160 }),
161 });
162 }
163 Some(Err(e)) if e.is_retryable() => {
164 failures += 1;
165 manager.mark_cooldown(&up.code, &key.label).await;
166 }
167 Some(Err(e)) => return Err(map_llm_error(&e)),
168 None => {
169 failures += 1;
170 manager.mark_cooldown(&up.code, &key.label).await;
171 }
172 }
173 }
174 }
175 Err(LproviderError::Upstream(UpstreamError::AllFailed(
176 "无可用上游 Key 可开启流式调用".into(),
177 )))
178}
Added agent/lprovider/router.rs agent/lprovider/router.rs +791 −0 View agent/lprovider/router.rs
@@ -0,0 +1,791 @@
1use std::pin::Pin;
2use std::sync::Arc;
3use std::task::{Context as TaskContext, Poll};
4
5use futures::Stream;
6use time::OffsetDateTime;
7use tracing::warn;
8use uuid::Uuid;
9
10use llm::token::TokenUsageStats;
11use llm::types::{ChatMessage, ChatRequest, ContentPart, MessageContent, StreamEvent};
12
13use config::lprovider::PointsPer1k;
14
15use crate::affinity::{AffinityManager, SessionBinding};
16use crate::audit::{AuditService, CallLogEntry};
17use crate::billing::{compute_bill, rates_for, worst_case_estimate};
18use crate::error::LproviderError;
19use crate::quota::QuotaService;
20use crate::rotation::{
21 CallOutcome, RotatingStream, StreamMeta, chat_stream_with_rotation, chat_with_rotation,
22};
23use crate::selector::{CandidateFilter, Selected, rank, resolve_strategy};
24use crate::state::AppState;
25use crate::upstream::{UpstreamManager, UpstreamSnapshot};
26use crate::wallet::WalletService;
27
28#[derive(Default, Clone)]
29pub struct ChatSelectionSpec {
30 pub strategy: Option<String>,
31 pub allowed: Vec<String>,
32 pub excluded: Vec<String>,
33 pub balance_bias: f64,
34 pub required_capabilities: Vec<String>,
35}
36
37pub struct ChatContext {
38 pub request_id: Uuid,
39 pub trace_id: Option<Uuid>,
40 pub tenant_id: i64,
41 pub wallet_id: Option<i64>,
42 pub subject_type: String,
43 pub subject_id: i64,
44 pub anonymous: bool,
45 pub model_alias: String,
46 pub request: ChatRequest,
47 pub max_output: usize,
48 pub stream: bool,
49 pub priority: i32,
50 pub caller: Option<String>,
51 pub caller_service: Option<String>,
52 pub session_key_override: Option<String>,
53 pub selection: ChatSelectionSpec,
54 pub client_ip: Option<String>,
55 pub user_agent: Option<String>,
56}
57
58pub struct CallOutcomeFull {
59 pub request_id: Uuid,
60 pub upstream_code: String,
61 pub key_label: String,
62 pub model: String,
63 pub chosen_strategy: String,
64 pub response: llm::types::ChatResponse,
65 pub bill: crate::billing::Bill,
66 pub balance_after: Option<i64>,
67}
68
69struct Prepared {
70 input_tokens: usize,
71 strategy: String,
72 ranked: Vec<Selected>,
73 session_key: Option<String>,
74 binding: Option<SessionBinding>,
75}
76
77struct StreamFinishPlan {
78 wallet_id: Option<i64>,
79 rate: PointsPer1k,
80 entry: CallLogEntry,
81 wallets: Arc<WalletService>,
82 audit_service: Arc<AuditService>,
83}
84
85pub struct Router {
86 state: Arc<AppState>,
87 manager: Arc<UpstreamManager>,
88 affinity: Arc<AffinityManager>,
89 wallets: Arc<WalletService>,
90 quotas: Arc<QuotaService>,
91 audit: Arc<AuditService>,
92}
93
94impl Router {
95 pub fn new(
96 state: Arc<AppState>,
97 manager: Arc<UpstreamManager>,
98 affinity: Arc<AffinityManager>,
99 ) -> Arc<Self> {
100 let db = state.db.clone().map(Arc::new);
101 Arc::new(Self {
102 state,
103 manager,
104 affinity,
105 wallets: Arc::new(WalletService::new(db.clone())),
106 quotas: Arc::new(QuotaService::new(db.clone())),
107 audit: Arc::new(AuditService::new(db)),
108 })
109 }
110
111 fn lp(&self) -> config::lprovider::LproviderConfig {
112 self.state.runtime.blocking_read().app.lprovider.clone()
113 }
114
115 fn worst_rate(pricing: &config::lprovider::LproviderPricingSettings) -> PointsPer1k {
116 let mut worst = pricing.default_per_1k;
117 for r in pricing.rates.values() {
118 if max_component(r) > max_component(&worst) {
119 worst = *r;
120 }
121 }
122 worst
123 }
124
125 async fn prepare(&self, ctx: &ChatContext) -> Result<Prepared, LproviderError> {
126 let lp = self.lp();
127 let snapshot = self.manager.snapshot().await;
128
129 let input_tokens = snapshot.estimate_tokens(&ctx.request)?;
130 if input_tokens.saturating_add(ctx.max_output) > lp.max_context_tokens {
131 return Err(LproviderError::Upstream(
132 crate::error::UpstreamError::ContextLengthExceeded,
133 ));
134 }
135
136 let worst = Self::worst_rate(&lp.pricing);
137 let estimated_max_cost =
138 worst_case_estimate(input_tokens as i64, ctx.max_output as i64, &worst);
139
140 let wallet_balance = if let Some(wallet_id) = ctx.wallet_id {
141 let wallet = self.wallets.get_wallet(wallet_id).await?;
142 self.wallets.verify_active(&wallet)?;
143 self.wallets
144 .verify_owner(&wallet, ctx.tenant_id, &ctx.subject_type, ctx.subject_id)?;
145 if wallet.balance_points < estimated_max_cost {
146 return Err(LproviderError::Wallet(
147 crate::error::WalletError::InsufficientBalance,
148 ));
149 }
150 Some(wallet.balance_points)
151 } else {
152 None
153 };
154
155 self.quotas
156 .check(
157 ctx.tenant_id,
158 &ctx.subject_type,
159 ctx.subject_id,
160 estimated_max_cost,
161 )
162 .await?;
163
164 let (model_default, model_overrides) = lp
165 .models
166 .iter()
167 .find(|m| m.alias == ctx.model_alias)
168 .map(|m| {
169 (
170 m.default_strategy.clone(),
171 m.priority_strategy_overrides.clone(),
172 )
173 })
174 .unwrap_or_else(|| ("BALANCED".to_string(), Vec::new()));
175 let strategy = resolve_strategy(
176 &model_default,
177 &model_overrides,
178 ctx.priority,
179 wallet_balance.unwrap_or(i64::MAX),
180 ctx.selection.strategy.as_deref(),
181 lp.low_balance_force_cost_points,
182 );
183
184 let mut required = ctx.selection.required_capabilities.clone();
185 if ctx.stream && !required.iter().any(|c| c == "streaming") {
186 required.push("streaming".to_string());
187 }
188 if !ctx.request.tools.is_empty() && !required.iter().any(|c| c == "tools") {
189 required.push("tools".to_string());
190 }
191 if messages_have_images(&ctx.request.messages) && !required.iter().any(|c| c == "vision") {
192 required.push("vision".to_string());
193 }
194 if matches!(
195 ctx.request.response_format,
196 Some(llm::types::ResponseFormat::Json)
197 ) && !required.iter().any(|c| c == "structured")
198 {
199 required.push("structured".to_string());
200 }
201
202 let filter = CandidateFilter {
203 required_capabilities: required,
204 input_tokens,
205 max_output: ctx.max_output,
206 };
207 let mut ranked = rank(
208 &snapshot,
209 &ctx.model_alias,
210 &filter,
211 &strategy,
212 &ctx.selection.allowed,
213 &ctx.selection.excluded,
214 ctx.selection.balance_bias,
215 lp.max_context_tokens,
216 )?;
217
218 let (session_key, binding) = if ctx.stream {
219 (None, None)
220 } else if let Some(k) = &ctx.session_key_override {
221 self.affinity.validate_session_key_override(k)?;
222 let binding = self
223 .affinity
224 .lookup(ctx.tenant_id, &ctx.model_alias, k)
225 .await?;
226 (Some(k.clone()), binding)
227 } else {
228 match self.affinity.compute_session_key(
229 &ctx.model_alias,
230 &ctx.request.messages,
231 &snapshot.estimator,
232 ) {
233 Some((k, _)) => {
234 let binding = self
235 .affinity
236 .lookup(ctx.tenant_id, &ctx.model_alias, &k)
237 .await?;
238 (Some(k), binding)
239 }
240 None => (None, None),
241 }
242 };
243
244 if let Some(b) = &binding
245 && let Some(pos) = ranked
246 .iter()
247 .position(|c| snapshot.upstreams[c.upstream_idx].code == b.upstream_code)
248 {
249 let item = ranked.remove(pos);
250 ranked.insert(0, item);
251 }
252
253 Ok(Prepared {
254 input_tokens,
255 strategy,
256 ranked,
257 session_key,
258 binding,
259 })
260 }
261
262 #[allow(clippy::too_many_arguments)]
263 fn build_entry(
264 &self,
265 ctx: &ChatContext,
266 outcome: Option<&CallOutcome>,
267 meta: Option<&StreamMeta>,
268 strategy: &str,
269 bill: Option<&crate::billing::Bill>,
270 rate: PointsPer1k,
271 input_tokens: usize,
272 started: OffsetDateTime,
273 completed: OffsetDateTime,
274 error_kind: Option<&str>,
275 error_message: Option<&str>,
276 snapshot: &UpstreamSnapshot,
277 ) -> CallLogEntry {
278 let upstream_idx = outcome
279 .map(|o| o.upstream_idx)
280 .or_else(|| meta.map(|m| m.upstream_idx));
281 let upstream = upstream_idx.and_then(|i| snapshot.upstreams.get(i));
282 let upstream_code = outcome
283 .map(|o| o.upstream_code.clone())
284 .or_else(|| meta.map(|m| m.upstream_code.clone()))
285 .unwrap_or_default();
286 let key_label = outcome
287 .map(|o| o.key_label.clone())
288 .or_else(|| meta.map(|m| m.key_label.clone()));
289 let prompt = outcome
290 .map(|o| o.response.usage.input_tokens as i32)
291 .unwrap_or(input_tokens as i32);
292 let completion = outcome
293 .map(|o| o.response.usage.output_tokens as i32)
294 .unwrap_or(0);
295 CallLogEntry {
296 request_id: ctx.request_id,
297 trace_id: ctx.trace_id,
298 tenant_id: ctx.tenant_id,
299 api_key_id: None,
300 subject_type: Some(ctx.subject_type.clone()),
301 subject_id: Some(ctx.subject_id),
302 wallet_id: ctx.wallet_id,
303 anonymous: ctx.anonymous,
304 subject_path: None,
305 session_key: None,
306 model_alias: ctx.model_alias.clone(),
307 upstream_code,
308 requested_model: ctx.request.model_override.clone(),
309 chosen_strategy: Some(strategy.to_string()),
310 tier_cost: upstream.map(|u| u.cost_tier).unwrap_or(0),
311 tier_quality: upstream.map(|u| u.quality_tier).unwrap_or(0),
312 tier_latency: upstream.map(|u| u.latency_tier).unwrap_or(0),
313 key_label,
314 caller: ctx.caller.clone(),
315 caller_service: ctx.caller_service.clone(),
316 priority: ctx.priority,
317 streamed: ctx.stream,
318 status: if error_kind.is_some() {
319 "error".to_string()
320 } else {
321 "success".to_string()
322 },
323 error_kind: error_kind.map(String::from),
324 error_message: error_message.map(String::from),
325 prompt_tokens: prompt,
326 completion_tokens: completion,
327 cache_hit_tokens: 0,
328 cache_miss_tokens: 0,
329 billable_input_tokens: bill
330 .map(|b| b.billable_input_tokens as i32)
331 .unwrap_or(prompt),
332 billable_output_tokens: bill
333 .map(|b| b.billable_output_tokens as i32)
334 .unwrap_or(completion),
335 points_input: bill.map(|b| b.points_input).unwrap_or(0),
336 points_output: bill.map(|b| b.points_output).unwrap_or(0),
337 points_cache_hit: bill.map(|b| b.points_cache_hit).unwrap_or(0),
338 points_total: bill.map(|b| b.points_total).unwrap_or(0),
339 points_rate_input: rate.points_per_1k_input_tokens,
340 points_rate_output: rate.points_per_1k_output_tokens,
341 points_rate_cache_hit: rate.points_per_1k_cache_hit_tokens,
342 request_started_at: started,
343 response_completed_at: Some(completed),
344 latency_ms: (completed - started).whole_milliseconds() as i32,
345 client_ip: ctx.client_ip.clone(),
346 user_agent: ctx.user_agent.clone(),
347 }
348 }
349
350 pub async fn count_tokens(&self, request: &ChatRequest) -> Result<usize, LproviderError> {
351 self.manager.snapshot().await.estimate_tokens(request)
352 }
353
354 pub async fn handle_chat(&self, ctx: ChatContext) -> Result<CallOutcomeFull, LproviderError> {
355 let started = OffsetDateTime::now_utc();
356 let lp = self.lp();
357 let snapshot = self.manager.snapshot().await;
358 let prepared = self.prepare(&ctx).await;
359
360 let prepared = match prepared {
361 Ok(p) => p,
362 Err(e) => {
363 let completed = OffsetDateTime::now_utc();
364 let entry = self.build_entry(
365 &ctx,
366 None,
367 None,
368 "",
369 None,
370 Self::worst_rate(&lp.pricing),
371 0,
372 started,
373 completed,
374 Some(error_kind_of(&e)),
375 Some(&crate::audit::sanitize_error(&e.to_string())),
376 &snapshot,
377 );
378 let audit = self.audit.clone();
379 tokio::spawn(async move {
380 if let Err(err) = audit.insert_call_log(&entry).await {
381 warn!(error = %err, "audit insert failed");
382 }
383 });
384 return Err(e);
385 }
386 };
387
388 let outcome =
389 chat_with_rotation(&self.manager, &snapshot, &prepared.ranked, &ctx.request).await;
390
391 let outcome = match outcome {
392 Ok(o) => o,
393 Err(e) => {
394 let completed = OffsetDateTime::now_utc();
395 let entry = self.build_entry(
396 &ctx,
397 None,
398 None,
399 &prepared.strategy,
400 None,
401 Self::worst_rate(&lp.pricing),
402 prepared.input_tokens,
403 started,
404 completed,
405 Some(error_kind_of(&e)),
406 Some(&crate::audit::sanitize_error(&e.to_string())),
407 &snapshot,
408 );
409 let audit = self.audit.clone();
410 tokio::spawn(async move {
411 if let Err(err) = audit.insert_call_log(&entry).await {
412 warn!(error = %err, "audit insert failed");
413 }
414 });
415 return Err(e);
416 }
417 };
418
419 let rate = *rates_for(&outcome.upstream_code, &lp.pricing);
420 let usage = &outcome.response.usage;
421 let bill = compute_bill(
422 usage.input_tokens as i64,
423 usage.output_tokens as i64,
424 0,
425 &rate,
426 );
427
428 let balance_after = if let Some(wallet_id) = ctx.wallet_id {
429 match self
430 .wallets
431 .debit(wallet_id, bill.points_total, ctx.request_id, "consume")
432 .await
433 {
434 Ok(ba) => Some(ba),
435 Err(e) => {
436 let completed = OffsetDateTime::now_utc();
437 let entry = self.build_entry(
438 &ctx,
439 Some(&outcome),
440 None,
441 &prepared.strategy,
442 Some(&bill),
443 rate,
444 prepared.input_tokens,
445 started,
446 completed,
447 Some("wallet_debit_failed"),
448 Some(&crate::audit::sanitize_error(&e.to_string())),
449 &snapshot,
450 );
451 let audit = self.audit.clone();
452 tokio::spawn(async move {
453 if let Err(err) = audit.insert_call_log(&entry).await {
454 warn!(error = %err, "audit insert failed");
455 }
456 });
457 return Err(e);
458 }
459 }
460 } else {
461 None
462 };
463
464 let completed = OffsetDateTime::now_utc();
465 let entry = self.build_entry(
466 &ctx,
467 Some(&outcome),
468 None,
469 &prepared.strategy,
470 Some(&bill),
471 rate,
472 prepared.input_tokens,
473 started,
474 completed,
475 None,
476 None,
477 &snapshot,
478 );
479 let audit = self.audit.clone();
480 tokio::spawn(async move {
481 if let Err(err) = audit.insert_call_log(&entry).await {
482 warn!(error = %err, "audit insert failed");
483 }
484 });
485
486 if let Some(key) = prepared.session_key.as_deref() {
487 self.record_affinity(
488 ctx.tenant_id,
489 &ctx.model_alias,
490 key,
491 &prepared.binding,
492 &outcome,
493 )
494 .await;
495 }
496
497 Ok(CallOutcomeFull {
498 request_id: ctx.request_id,
499 upstream_code: outcome.upstream_code.clone(),
500 key_label: outcome.key_label.clone(),
501 model: outcome.model.clone(),
502 chosen_strategy: prepared.strategy,
503 response: outcome.response,
504 bill,
505 balance_after,
506 })
507 }
508
509 async fn record_affinity(
510 &self,
511 tenant_id: i64,
512 model_alias: &str,
513 session_key: &str,
514 binding: &Option<SessionBinding>,
515 outcome: &CallOutcome,
516 ) {
517 let prompt_tokens = outcome.response.usage.input_tokens as i64;
518 match binding {
519 Some(b) if b.upstream_code == outcome.upstream_code => {
520 self.state.record_hit();
521 let a = self.affinity.clone();
522 let tenant = tenant_id;
523 let alias = model_alias.to_string();
524 let key = session_key.to_string();
525 tokio::spawn(async move {
526 if let Err(e) = a
527 .record_feedback(tenant, &alias, &key, true, prompt_tokens)
528 .await
529 {
530 warn!(error = %e, "affinity feedback failed");
531 }
532 });
533 }
534 Some(_) => {
535 self.state.record_miss();
536 let a = self.affinity.clone();
537 let tenant = tenant_id;
538 let alias = model_alias.to_string();
539 let key = session_key.to_string();
540 tokio::spawn(async move {
541 if let Err(e) = a.record_feedback(tenant, &alias, &key, false, 0).await {
542 warn!(error = %e, "affinity feedback failed");
543 }
544 if let Err(e) = a.invalidate(tenant, &alias, &key).await {
545 warn!(error = %e, "affinity invalidate failed");
546 }
547 });
548 }
549 None => {
550 let binding = SessionBinding {
551 upstream_code: outcome.upstream_code.clone(),
552 key_label: outcome.key_label.clone(),
553 model: outcome.model.clone(),
554 prefix_hash: session_key.to_string(),
555 bound_at: OffsetDateTime::now_utc().unix_timestamp(),
556 last_hit_tokens: 0,
557 hit_count: 0,
558 miss_count: 0,
559 };
560 self.state.record_bind();
561 let a = self.affinity.clone();
562 let tenant = tenant_id;
563 let alias = model_alias.to_string();
564 let key = session_key.to_string();
565 tokio::spawn(async move {
566 if let Err(e) = a.bind(tenant, &alias, &key, &binding).await {
567 warn!(error = %e, "affinity bind failed");
568 }
569 });
570 }
571 }
572 }
573
574 pub async fn handle_stream(&self, ctx: ChatContext) -> Result<BilledStream, LproviderError> {
575 let started = OffsetDateTime::now_utc();
576 let lp = self.lp();
577 let snapshot = self.manager.snapshot().await;
578
579 let prepared = match self.prepare(&ctx).await {
580 Ok(p) => p,
581 Err(e) => {
582 let completed = OffsetDateTime::now_utc();
583 let entry = self.build_entry(
584 &ctx,
585 None,
586 None,
587 "",
588 None,
589 Self::worst_rate(&lp.pricing),
590 0,
591 started,
592 completed,
593 Some(error_kind_of(&e)),
594 Some(&crate::audit::sanitize_error(&e.to_string())),
595 &snapshot,
596 );
597 let audit = self.audit.clone();
598 tokio::spawn(async move {
599 if let Err(err) = audit.insert_call_log(&entry).await {
600 warn!(error = %err, "audit insert failed");
601 }
602 });
603 return Err(e);
604 }
605 };
606
607 let stream =
608 chat_stream_with_rotation(&self.manager, &snapshot, &prepared.ranked, &ctx.request)
609 .await;
610
611 let stream = match stream {
612 Ok(s) => s,
613 Err(e) => {
614 let completed = OffsetDateTime::now_utc();
615 let entry = self.build_entry(
616 &ctx,
617 None,
618 None,
619 &prepared.strategy,
620 None,
621 Self::worst_rate(&lp.pricing),
622 prepared.input_tokens,
623 started,
624 completed,
625 Some(error_kind_of(&e)),
626 Some(&crate::audit::sanitize_error(&e.to_string())),
627 &snapshot,
628 );
629 let audit = self.audit.clone();
630 tokio::spawn(async move {
631 if let Err(err) = audit.insert_call_log(&entry).await {
632 warn!(error = %err, "audit insert failed");
633 }
634 });
635 return Err(e);
636 }
637 };
638
639 let meta = stream
640 .meta
641 .as_ref()
642 .cloned()
643 .ok_or_else(|| LproviderError::Internal("流元数据缺失".to_string()))?;
644 let rate = *rates_for(&meta.upstream_code, &lp.pricing);
645 let entry = self.build_entry(
646 &ctx,
647 None,
648 Some(&meta),
649 &prepared.strategy,
650 None,
651 rate,
652 prepared.input_tokens,
653 started,
654 started,
655 None,
656 None,
657 &snapshot,
658 );
659
660 let plan = StreamFinishPlan {
661 wallet_id: ctx.wallet_id,
662 rate,
663 entry,
664 wallets: self.wallets.clone(),
665 audit_service: self.audit.clone(),
666 };
667
668 Ok(BilledStream {
669 inner: stream,
670 plan: Some(plan),
671 usage: TokenUsageStats::new(0, 0),
672 finished: false,
673 })
674 }
675}
676
677fn max_component(r: &PointsPer1k) -> i64 {
678 r.points_per_1k_input_tokens
679 .max(r.points_per_1k_output_tokens)
680 .max(r.points_per_1k_cache_hit_tokens)
681}
682
683fn error_kind_of(e: &LproviderError) -> &'static str {
684 match e {
685 LproviderError::Wallet(crate::error::WalletError::InsufficientBalance) => {
686 "insufficient_balance"
687 }
688 LproviderError::Wallet(crate::error::WalletError::DebitFailed) => "wallet_debit_failed",
689 LproviderError::QuotaExceeded(_) => "quota_exceeded",
690 LproviderError::Upstream(crate::error::UpstreamError::ContextLengthExceeded) => {
691 "context_length_exceeded"
692 }
693 LproviderError::Upstream(crate::error::UpstreamError::NoUpstreamAvailable) => "no_upstream",
694 LproviderError::Upstream(crate::error::UpstreamError::NoKeyAvailable(_)) => "no_key",
695 LproviderError::Upstream(crate::error::UpstreamError::AllFailed(_)) => "upstream_failed",
696 LproviderError::InvalidRequest(_) => "invalid_request",
697 LproviderError::Auth(_) => "auth",
698 _ => "internal",
699 }
700}
701
702fn messages_have_images(messages: &[ChatMessage]) -> bool {
703 messages.iter().any(|m| match &m.content {
704 MessageContent::Text(_) => false,
705 MessageContent::Parts(parts) => {
706 parts.iter().any(|p| matches!(p, ContentPart::Image { .. }))
707 }
708 })
709}
710
711pub struct BilledStream {
712 inner: RotatingStream,
713 plan: Option<StreamFinishPlan>,
714 usage: TokenUsageStats,
715 finished: bool,
716}
717
718impl Stream for BilledStream {
719 type Item = Result<StreamEvent, llm::error::LlmError>;
720
721 fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
722 if self.finished {
723 return Poll::Ready(None);
724 }
725 match Pin::new(&mut self.inner).poll_next(cx) {
726 Poll::Ready(item) => match item {
727 Some(Ok(StreamEvent::Usage(u))) => {
728 if u.input_tokens > self.usage.input_tokens {
729 self.usage.input_tokens = u.input_tokens;
730 }
731 if u.output_tokens > self.usage.output_tokens {
732 self.usage.output_tokens = u.output_tokens;
733 }
734 if u.reasoning_tokens.is_some() {
735 self.usage.reasoning_tokens = u.reasoning_tokens;
736 }
737 Poll::Ready(Some(Ok(StreamEvent::Usage(u))))
738 }
739 Some(x) => Poll::Ready(Some(x)),
740 None => {
741 self.finished = true;
742 if let Some(plan) = self.plan.take() {
743 let usage = self.usage.clone();
744 tokio::spawn(finalize_stream(plan, usage));
745 }
746 Poll::Ready(None)
747 }
748 },
749 Poll::Pending => Poll::Pending,
750 }
751 }
752}
753
754async fn finalize_stream(plan: StreamFinishPlan, usage: TokenUsageStats) {
755 let bill = compute_bill(
756 usage.input_tokens as i64,
757 usage.output_tokens as i64,
758 0,
759 &plan.rate,
760 );
761 if let Some(wallet_id) = plan.wallet_id
762 && let Err(e) = plan
763 .wallets
764 .debit(
765 wallet_id,
766 bill.points_total,
767 plan.entry.request_id,
768 "consume",
769 )
770 .await
771 {
772 warn!(error = %e, wallet_id, "stream debit failed");
773 }
774 let mut entry = plan.entry;
775 entry.prompt_tokens = usage.input_tokens as i32;
776 entry.completion_tokens = usage.output_tokens as i32;
777 entry.billable_input_tokens = bill.billable_input_tokens as i32;
778 entry.billable_output_tokens = bill.billable_output_tokens as i32;
779 entry.points_input = bill.points_input;
780 entry.points_output = bill.points_output;
781 entry.points_cache_hit = bill.points_cache_hit;
782 entry.points_total = bill.points_total;
783 entry.status = "success".to_string();
784 entry.response_completed_at = Some(OffsetDateTime::now_utc());
785 if let Some(completed) = entry.response_completed_at {
786 entry.latency_ms = (completed - entry.request_started_at).whole_milliseconds() as i32;
787 }
788 if let Err(e) = plan.audit_service.insert_call_log(&entry).await {
789 warn!(error = %e, "stream audit insert failed");
790 }
791}
Added agent/lprovider/selector.rs agent/lprovider/selector.rs +183 −0 View agent/lprovider/selector.rs
@@ -0,0 +1,183 @@
1use config::lprovider::PriorityStrategyOverride;
2
3use crate::error::{LproviderError, UpstreamError};
4use crate::upstream::{UpstreamEntry, UpstreamSnapshot};
5
6pub struct CandidateFilter {
7 pub required_capabilities: Vec<String>,
8 pub input_tokens: usize,
9 pub max_output: usize,
10}
11
12#[derive(Debug)]
13pub struct Selected {
14 pub upstream_idx: usize,
15 pub model: String,
16 pub effective_context: usize,
17}
18
19pub fn resolve_strategy(
20 default_strategy: &str,
21 priority_overrides: &[PriorityStrategyOverride],
22 priority: i32,
23 wallet_balance: i64,
24 request_strategy: Option<&str>,
25 low_balance_force_cost_points: i64,
26) -> String {
27 if let Some(s) = request_strategy
28 && !s.is_empty()
29 {
30 return s.to_string();
31 }
32 for ov in priority_overrides {
33 let within =
34 priority >= ov.min_priority && ov.max_priority.is_none_or(|max| priority < max);
35 if within {
36 return ov.strategy.clone();
37 }
38 }
39 if wallet_balance < low_balance_force_cost_points {
40 return "COST".to_string();
41 }
42 default_strategy.to_string()
43}
44
45#[allow(clippy::too_many_arguments)]
46pub fn rank(
47 snapshot: &UpstreamSnapshot,
48 alias: &str,
49 filter: &CandidateFilter,
50 strategy: &str,
51 allowed: &[String],
52 excluded: &[String],
53 balance_bias: f64,
54 gateway_max_context: usize,
55) -> Result<Vec<Selected>, LproviderError> {
56 let mut pool: Vec<(usize, &UpstreamEntry)> = Vec::new();
57 for (idx, up) in snapshot.upstreams.iter().enumerate() {
58 if alias != "auto" && up.code != alias {
59 continue;
60 }
61 if !allowed.is_empty() && !allowed.iter().any(|a| a == &up.code) {
62 continue;
63 }
64 if excluded.iter().any(|e| e == &up.code) {
65 continue;
66 }
67 if !capabilities_ok(up, filter) {
68 continue;
69 }
70 pool.push((idx, up));
71 }
72 if pool.is_empty() {
73 return Err(LproviderError::Upstream(UpstreamError::NoUpstreamAvailable));
74 }
75
76 pool.retain(|(_, up)| {
77 let effective = gateway_max_context.min(up.context_window);
78 filter.input_tokens + filter.max_output <= effective
79 });
80 if pool.is_empty() {
81 return Err(LproviderError::Upstream(
82 UpstreamError::ContextLengthExceeded,
83 ));
84 }
85
86 sort_pool(&mut pool, strategy, allowed, balance_bias);
87
88 Ok(pool
89 .into_iter()
90 .map(|(upstream_idx, up)| Selected {
91 upstream_idx,
92 model: if alias == "auto" {
93 up.default_model.clone()
94 } else {
95 alias.to_string()
96 },
97 effective_context: gateway_max_context.min(up.context_window),
98 })
99 .collect())
100}
101
102#[allow(clippy::too_many_arguments)]
103pub fn select(
104 snapshot: &UpstreamSnapshot,
105 alias: &str,
106 filter: &CandidateFilter,
107 strategy: &str,
108 allowed: &[String],
109 excluded: &[String],
110 balance_bias: f64,
111 gateway_max_context: usize,
112) -> Result<Selected, LproviderError> {
113 rank(
114 snapshot,
115 alias,
116 filter,
117 strategy,
118 allowed,
119 excluded,
120 balance_bias,
121 gateway_max_context,
122 )?
123 .into_iter()
124 .next()
125 .ok_or(LproviderError::Upstream(UpstreamError::NoUpstreamAvailable))
126}
127
128fn capabilities_ok(up: &UpstreamEntry, filter: &CandidateFilter) -> bool {
129 filter
130 .required_capabilities
131 .iter()
132 .all(|need| up.capabilities.iter().any(|have| have == need))
133}
134
135fn sort_pool(
136 pool: &mut [(usize, &UpstreamEntry)],
137 strategy: &str,
138 allowed: &[String],
139 balance_bias: f64,
140) {
141 match strategy {
142 "COST" => pool.sort_by(|a, b| {
143 a.1.cost_tier
144 .cmp(&b.1.cost_tier)
145 .then(b.1.quality_tier.cmp(&a.1.quality_tier))
146 .then(a.1.latency_tier.cmp(&b.1.latency_tier))
147 }),
148 "QUALITY" => pool.sort_by(|a, b| {
149 b.1.quality_tier
150 .cmp(&a.1.quality_tier)
151 .then(a.1.cost_tier.cmp(&b.1.cost_tier))
152 .then(a.1.latency_tier.cmp(&b.1.latency_tier))
153 }),
154 "LATENCY" => pool.sort_by(|a, b| {
155 a.1.latency_tier
156 .cmp(&b.1.latency_tier)
157 .then(b.1.quality_tier.cmp(&a.1.quality_tier))
158 .then(a.1.cost_tier.cmp(&b.1.cost_tier))
159 }),
160 "BALANCED" => pool.sort_by(|a, b| {
161 let sa = balanced_score(a.1, balance_bias);
162 let sb = balanced_score(b.1, balance_bias);
163 sb.partial_cmp(&sa)
164 .unwrap_or(std::cmp::Ordering::Equal)
165 .then(b.1.quality_tier.cmp(&a.1.quality_tier))
166 .then(a.1.cost_tier.cmp(&b.1.cost_tier))
167 }),
168 "EXPLICIT" if !allowed.is_empty() => pool.sort_by_key(|(_, up)| {
169 allowed
170 .iter()
171 .position(|a| a == &up.code)
172 .unwrap_or(usize::MAX)
173 }),
174 _ => {}
175 }
176}
177
178fn balanced_score(up: &UpstreamEntry, bias: f64) -> f64 {
179 0.5 * f64::from(up.quality_tier)
180 - 0.3 * f64::from(up.cost_tier)
181 - 0.2 * f64::from(up.latency_tier)
182 + bias
183}
Added agent/lprovider/state.rs agent/lprovider/state.rs +108 −0 View agent/lprovider/state.rs
@@ -0,0 +1,108 @@
1use std::sync::Arc;
2use std::sync::atomic::{AtomicI64, Ordering};
3
4use cache::RedisCache;
5use db::conn::Db;
6use sea_orm::{ConnectionTrait, Statement};
7use time::OffsetDateTime;
8use tokio::sync::RwLock;
9
10use crate::config::{CliOverrides, RuntimeConfig};
11
12pub struct AppState {
13 pub runtime: Arc<RwLock<RuntimeConfig>>,
14 pub db: Option<Db>,
15 pub redis: Option<Arc<RedisCache>>,
16 pub instance_id: String,
17 pub cluster_name: String,
18 pub started_at: OffsetDateTime,
19 pub pid: u32,
20 pub affinity: AffinityCounters,
21 pub shutdown: tokio_util::sync::CancellationToken,
22 pub cli: CliOverrides,
23}
24
25#[derive(Default)]
26pub struct AffinityCounters {
27 pub bind_total: AtomicI64,
28 pub hit_count: AtomicI64,
29 pub miss_count: AtomicI64,
30}
31
32impl AppState {
33 pub fn new(
34 runtime: RuntimeConfig,
35 db: Option<Db>,
36 redis: Option<Arc<RedisCache>>,
37 instance_id: String,
38 cluster_name: String,
39 cli: CliOverrides,
40 ) -> Arc<Self> {
41 Arc::new(Self {
42 runtime: Arc::new(RwLock::new(runtime)),
43 db,
44 redis,
45 instance_id,
46 cluster_name,
47 started_at: OffsetDateTime::now_utc(),
48 pid: std::process::id(),
49 affinity: AffinityCounters::default(),
50 shutdown: tokio_util::sync::CancellationToken::new(),
51 cli,
52 })
53 }
54
55 pub async fn redis_ok(&self) -> bool {
56 let Some(redis) = &self.redis else {
57 return false;
58 };
59 let mut conn = redis.connection();
60 redis::cmd("PING")
61 .query_async::<String>(&mut conn)
62 .await
63 .is_ok()
64 }
65
66 pub async fn db_writable(&self) -> bool {
67 let Some(db) = &self.db else {
68 return false;
69 };
70 let backend = db.writer().get_database_backend();
71 db.writer()
72 .query_one_raw(Statement::from_string(backend, "SELECT 1"))
73 .await
74 .is_ok()
75 }
76
77 pub async fn ready(&self) -> Result<(), String> {
78 if !self.redis_ok().await {
79 return Err("redis unreachable".to_string());
80 }
81 if !self.db_writable().await {
82 return Err("database not writable".to_string());
83 }
84 let runtime = self.runtime.read().await;
85 runtime
86 .lp()
87 .is_valid()
88 .map_err(|e| format!("config invalid: {e}"))
89 }
90
91 pub fn snapshot_version(&self) -> String {
92 let runtime = self.runtime.blocking_read();
93 let sha = &runtime.config_sha256;
94 sha.chars().take(12).collect()
95 }
96
97 pub fn record_bind(&self) {
98 self.affinity.bind_total.fetch_add(1, Ordering::Relaxed);
99 }
100
101 pub fn record_hit(&self) {
102 self.affinity.hit_count.fetch_add(1, Ordering::Relaxed);
103 }
104
105 pub fn record_miss(&self) {
106 self.affinity.miss_count.fetch_add(1, Ordering::Relaxed);
107 }
108}
Added agent/lprovider/upstream.rs agent/lprovider/upstream.rs +220 −0 View agent/lprovider/upstream.rs
@@ -0,0 +1,220 @@
1use std::sync::Arc;
2use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
3use std::time::{Duration, Instant};
4
5use cache::RedisCache;
6use config::lprovider::LproviderConfig;
7use llm::provider::LlmProvider;
8use llm::provider::openai::{OpenAIConfig, OpenAIProvider};
9use redis::AsyncCommands;
10use tokio::sync::RwLock;
11
12use crate::credentials::{self, ResolvedCredential};
13use crate::error::LproviderError;
14
15pub struct KeyEntry {
16 pub label: String,
17 pub fingerprint: String,
18 pub qps_limit: i64,
19 pub rpm_limit: i64,
20 pub provider: Arc<dyn LlmProvider>,
21 pub in_flight: Arc<AtomicI64>,
22 pub cooldown_until: RwLock<Option<Instant>>,
23 pub rr: AtomicU64,
24}
25
26pub struct UpstreamEntry {
27 pub code: String,
28 pub default_model: String,
29 pub capabilities: Vec<String>,
30 pub context_window: usize,
31 pub max_output: usize,
32 pub cost_tier: i32,
33 pub quality_tier: i32,
34 pub latency_tier: i32,
35 pub tags: Vec<String>,
36 pub keys: Vec<KeyEntry>,
37}
38
39pub struct UpstreamSnapshot {
40 pub upstreams: Vec<UpstreamEntry>,
41 pub estimator: Arc<dyn LlmProvider>,
42}
43
44impl UpstreamSnapshot {
45 pub fn estimate_tokens(&self, req: &llm::types::ChatRequest) -> Result<usize, LproviderError> {
46 Ok(self.estimator.count_tokens(req))
47 }
48}
49
50pub struct KeyLease {
51 in_flight: Arc<AtomicI64>,
52}
53
54impl Drop for KeyLease {
55 fn drop(&mut self) {
56 self.in_flight.fetch_sub(1, Ordering::SeqCst);
57 }
58}
59
60pub struct UpstreamManager {
61 snapshot: RwLock<Arc<UpstreamSnapshot>>,
62 redis: Option<Arc<RedisCache>>,
63 cooldown_seconds: i64,
64 key_prefix: String,
65}
66
67impl UpstreamManager {
68 pub fn new(
69 cfg: &LproviderConfig,
70 redis: Option<Arc<RedisCache>>,
71 ) -> Result<Self, LproviderError> {
72 let snapshot = Arc::new(build_snapshot(cfg)?);
73 Ok(Self {
74 snapshot: RwLock::new(snapshot),
75 redis,
76 cooldown_seconds: cfg.redis.cooldown_seconds,
77 key_prefix: cfg.redis.key_prefix.clone(),
78 })
79 }
80
81 pub async fn snapshot(&self) -> Arc<UpstreamSnapshot> {
82 self.snapshot.read().await.clone()
83 }
84
85 pub fn reload(&self, cfg: &LproviderConfig) -> Result<(), LproviderError> {
86 let new_snapshot = Arc::new(build_snapshot(cfg)?);
87 *self.snapshot.blocking_write() = new_snapshot;
88 Ok(())
89 }
90
91 pub async fn select_key(&self, upstream: &UpstreamEntry) -> Option<usize> {
92 let mut best: Option<(usize, i64, u64)> = None;
93 for (idx, key) in upstream.keys.iter().enumerate() {
94 if self.is_cooldown(&upstream.code, key).await {
95 continue;
96 }
97 let inflight = key.in_flight.load(Ordering::SeqCst);
98 let rr = key.rr.load(Ordering::SeqCst);
99 let candidate = (idx, inflight, rr);
100 match best {
101 Some(b) if (b.1, b.2) <= (candidate.1, candidate.2) => {}
102 _ => best = Some(candidate),
103 }
104 }
105 let (idx, _, _) = best?;
106 upstream.keys[idx].rr.fetch_add(1, Ordering::SeqCst);
107 Some(idx)
108 }
109
110 pub async fn begin_key(&self, upstream_idx: usize, key_idx: usize) -> Option<KeyLease> {
111 let snapshot = self.snapshot().await;
112 let upstream = snapshot.upstreams.get(upstream_idx)?;
113 let key = upstream.keys.get(key_idx)?;
114 key.in_flight.fetch_add(1, Ordering::SeqCst);
115 Some(KeyLease {
116 in_flight: Arc::clone(&key.in_flight),
117 })
118 }
119
120 pub async fn mark_cooldown(&self, code: &str, label: &str) {
121 let snapshot = self.snapshot().await;
122 if let Some(key) = snapshot
123 .upstreams
124 .iter()
125 .find(|u| u.code == code)
126 .and_then(|u| u.keys.iter().find(|k| k.label == label))
127 {
128 *key.cooldown_until.write().await =
129 Some(Instant::now() + Duration::from_secs(self.cooldown_seconds.max(1) as u64));
130 }
131 if let Some(cache) = &self.redis {
132 let keyname = format!("{}:cooldown:{}:{}", self.key_prefix, code, label);
133 let mut c = cache.connection();
134 let _: redis::RedisResult<()> = c
135 .set_ex(&keyname, "1", self.cooldown_seconds.max(1) as u64)
136 .await;
137 }
138 }
139
140 async fn is_cooldown(&self, code: &str, key: &KeyEntry) -> bool {
141 if let Some(until) = *key.cooldown_until.read().await
142 && until > Instant::now()
143 {
144 return true;
145 }
146 let Some(cache) = &self.redis else {
147 return false;
148 };
149 let keyname = format!("{}:cooldown:{}:{}", self.key_prefix, code, key.label);
150 let mut c = cache.connection();
151 c.exists(&keyname).await.unwrap_or(false)
152 }
153}
154
155fn build_snapshot(cfg: &LproviderConfig) -> Result<UpstreamSnapshot, LproviderError> {
156 let creds = credentials::resolve_all(cfg)
157 .map_err(|e| LproviderError::Config(format!("credential resolution failed: {e}")))?;
158 let request_timeout = Duration::from_millis(cfg.request_timeout_ms);
159 let mut upstreams = Vec::new();
160 for up in &cfg.upstreams {
161 if !up.enabled {
162 continue;
163 }
164 let mut keys = Vec::new();
165 for cred in &up.credentials {
166 let resolved = find_credential(&creds, &up.code, &cred.label)?;
167 let oc = OpenAIConfig::new(resolved.api_key.as_str())
168 .with_model(up.default_model.clone())
169 .with_base_url(up.base_url.clone())
170 .with_request_timeout(request_timeout)
171 .with_max_context_tokens(up.context_window)
172 .with_max_output_tokens(up.max_output);
173 let provider = OpenAIProvider::new(oc).map_err(|e| {
174 LproviderError::Config(format!(
175 "build provider {}/{} failed: {e}",
176 up.code, cred.label
177 ))
178 })?;
179 keys.push(KeyEntry {
180 label: cred.label.clone(),
181 fingerprint: resolved.fingerprint.clone(),
182 qps_limit: cred.qps_limit,
183 rpm_limit: cred.rpm_limit,
184 provider: Arc::new(provider),
185 in_flight: Arc::new(AtomicI64::new(0)),
186 cooldown_until: RwLock::new(None),
187 rr: AtomicU64::new(0),
188 });
189 }
190 upstreams.push(UpstreamEntry {
191 code: up.code.clone(),
192 default_model: up.default_model.clone(),
193 capabilities: up.capabilities.clone(),
194 context_window: up.context_window,
195 max_output: up.max_output,
196 cost_tier: up.cost_tier,
197 quality_tier: up.quality_tier,
198 latency_tier: up.latency_tier,
199 tags: up.tags.clone(),
200 keys,
201 });
202 }
203 let estimator = OpenAIProvider::new(OpenAIConfig::new("estimator"))
204 .map_err(|e| LproviderError::Config(format!("build token estimator failed: {e}")))?;
205 Ok(UpstreamSnapshot {
206 upstreams,
207 estimator: Arc::new(estimator),
208 })
209}
210
211fn find_credential<'a>(
212 creds: &'a [ResolvedCredential],
213 code: &str,
214 label: &str,
215) -> Result<&'a ResolvedCredential, LproviderError> {
216 creds
217 .iter()
218 .find(|c| c.upstream_code == code && c.label == label)
219 .ok_or_else(|| LproviderError::Config(format!("credential {code}:{label} missing")))
220}
Added agent/lprovider/wallet.rs agent/lprovider/wallet.rs +512 −0 View agent/lprovider/wallet.rs
@@ -0,0 +1,512 @@
1use std::sync::Arc;
2
3use db::conn::Db;
4use db::lprovider::{llm_wallet, llm_wallet_freeze_log, llm_wallet_ledger};
5use sea_orm::{
6 ActiveModelTrait, ColumnTrait, Condition, ConnectionTrait, DatabaseBackend, EntityTrait,
7 PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, Statement, TransactionTrait, Value,
8};
9use time::OffsetDateTime;
10use uuid::Uuid;
11
12use crate::error::{LproviderError, WalletError};
13
14#[derive(Clone, Debug)]
15pub struct WalletInfo {
16 pub wallet_id: i64,
17 pub owner_type: String,
18 pub owner_id: i64,
19 pub tenant_id: i64,
20 pub status: String,
21 pub balance_points: i64,
22 pub total_topped_up: i64,
23 pub total_consumed: i64,
24 pub total_refunded: i64,
25 pub created_at: OffsetDateTime,
26 pub updated_at: OffsetDateTime,
27}
28
29#[derive(Clone, Debug)]
30pub struct LedgerEntry {
31 pub wallet_id: i64,
32 pub request_id: Option<Uuid>,
33 pub delta_points: i64,
34 pub balance_after: i64,
35 pub kind: String,
36 pub idempotency_key: Option<String>,
37 pub reason: Option<String>,
38 pub operator: Option<String>,
39 pub created_at: OffsetDateTime,
40}
41
42#[derive(Clone)]
43pub struct WalletService {
44 db: Option<Arc<Db>>,
45}
46
47impl WalletService {
48 pub fn new(db: Option<Arc<Db>>) -> Self {
49 Self { db }
50 }
51
52 fn db(&self) -> Result<&Db, LproviderError> {
53 self.db.as_deref().ok_or_else(|| {
54 LproviderError::Db(db::conn::DbError::InvalidConfig("db not configured".into()))
55 })
56 }
57
58 pub async fn get_wallet(&self, wallet_id: i64) -> Result<WalletInfo, LproviderError> {
59 let db = self.db()?;
60 let model = llm_wallet::Entity::find_by_id(wallet_id)
61 .one(db.primary_connection())
62 .await
63 .map_err(db_err)?
64 .ok_or(LproviderError::Wallet(WalletError::WalletNotFound(
65 wallet_id,
66 )))?;
67 Ok(wallet_from_model(model))
68 }
69
70 pub async fn get_wallet_by_owner(
71 &self,
72 owner_type: &str,
73 owner_id: i64,
74 ) -> Result<WalletInfo, LproviderError> {
75 let db = self.db()?;
76 let model = llm_wallet::Entity::find()
77 .filter(
78 Condition::all()
79 .add(llm_wallet::Column::OwnerType.eq(owner_type))
80 .add(llm_wallet::Column::OwnerId.eq(owner_id)),
81 )
82 .one(db.primary_connection())
83 .await
84 .map_err(db_err)?
85 .ok_or_else(|| {
86 LproviderError::NotFound(format!("wallet for {owner_type}:{owner_id}"))
87 })?;
88 Ok(wallet_from_model(model))
89 }
90
91 pub async fn list_wallets(
92 &self,
93 tenant_id: Option<i64>,
94 owner_type: Option<String>,
95 owner_id: Option<i64>,
96 status: Option<String>,
97 limit: u64,
98 offset: u64,
99 ) -> Result<(Vec<WalletInfo>, u64), LproviderError> {
100 let db = self.db()?;
101 let mut cond = sea_orm::Condition::all();
102 if let Some(t) = tenant_id {
103 cond = cond.add(llm_wallet::Column::TenantId.eq(t));
104 }
105 if let Some(ot) = owner_type {
106 cond = cond.add(llm_wallet::Column::OwnerType.eq(ot));
107 }
108 if let Some(oi) = owner_id {
109 cond = cond.add(llm_wallet::Column::OwnerId.eq(oi));
110 }
111 if let Some(s) = status {
112 cond = cond.add(llm_wallet::Column::Status.eq(s));
113 }
114 let total = llm_wallet::Entity::find()
115 .filter(cond.clone())
116 .count(db.primary_connection())
117 .await
118 .map_err(db_err)?;
119 let models = llm_wallet::Entity::find()
120 .filter(cond)
121 .order_by_asc(llm_wallet::Column::WalletId)
122 .limit(limit)
123 .offset(offset)
124 .all(db.primary_connection())
125 .await
126 .map_err(db_err)?;
127 Ok((models.into_iter().map(wallet_from_model).collect(), total))
128 }
129
130 pub fn verify_active(&self, wallet: &WalletInfo) -> Result<(), LproviderError> {
131 match wallet.status.as_str() {
132 "active" => Ok(()),
133 _ => Err(LproviderError::Wallet(WalletError::WalletNotActive(
134 wallet.wallet_id,
135 ))),
136 }
137 }
138
139 pub fn verify_owner(
140 &self,
141 wallet: &WalletInfo,
142 tenant_id: i64,
143 subject_type: &str,
144 subject_id: i64,
145 ) -> Result<(), LproviderError> {
146 if wallet.tenant_id == tenant_id
147 && wallet.owner_type == subject_type
148 && wallet.owner_id == subject_id
149 {
150 Ok(())
151 } else {
152 Err(LproviderError::Wallet(WalletError::OwnerMismatch))
153 }
154 }
155
156 pub async fn debit(
157 &self,
158 wallet_id: i64,
159 actual_points: i64,
160 request_id: Uuid,
161 kind: &str,
162 ) -> Result<i64, LproviderError> {
163 let db = self.db()?;
164 let writer = db.writer();
165 let txn = writer.begin().await.map_err(db_err)?;
166 let row = txn
167 .query_one_raw(Statement::from_sql_and_values(
168 DatabaseBackend::Postgres,
169 "UPDATE llm_wallet SET balance_points = balance_points - $1, total_consumed = total_consumed + $1, last_consume_at = now(), updated_at = now() WHERE wallet_id = $2 AND status = 'active' AND balance_points >= $1 RETURNING balance_points AS balance_after",
170 [Value::BigInt(Some(actual_points)), Value::BigInt(Some(wallet_id))],
171 ))
172 .await
173 .map_err(db_err)?;
174 let Some(result) = row else {
175 txn.rollback().await.map_err(db_err)?;
176 let wallet = self.get_wallet(wallet_id).await?;
177 if wallet.status != "active" {
178 return Err(LproviderError::Wallet(WalletError::WalletNotActive(
179 wallet_id,
180 )));
181 }
182 return Err(LproviderError::Wallet(WalletError::InsufficientBalance));
183 };
184 let balance_after: i64 = result.try_get("", "balance_after").unwrap_or_default();
185 let ledger: llm_wallet_ledger::ActiveModel = llm_wallet_ledger::ActiveModel::builder()
186 .set_wallet_id(wallet_id)
187 .set_request_id(Some(request_id))
188 .set_delta_points(-actual_points)
189 .set_balance_after(balance_after)
190 .set_kind(kind.to_string())
191 .set_reason(None)
192 .set_operator(None)
193 .set_created_at(OffsetDateTime::now_utc())
194 .into();
195 ledger.insert(&txn).await.map_err(db_err)?;
196 txn.commit().await.map_err(db_err)?;
197 Ok(balance_after)
198 }
199
200 pub async fn topup(
201 &self,
202 wallet_id: i64,
203 points: i64,
204 idempotency_key: &str,
205 reason: Option<String>,
206 operator: Option<String>,
207 ) -> Result<i64, LproviderError> {
208 self.adjust_balance(
209 wallet_id,
210 points,
211 "topup",
212 idempotency_key,
213 reason,
214 operator,
215 )
216 .await
217 }
218
219 pub async fn adjust(
220 &self,
221 wallet_id: i64,
222 delta_points: i64,
223 idempotency_key: &str,
224 reason: Option<String>,
225 operator: Option<String>,
226 ) -> Result<i64, LproviderError> {
227 self.adjust_balance(
228 wallet_id,
229 delta_points,
230 "adjust",
231 idempotency_key,
232 reason,
233 operator,
234 )
235 .await
236 }
237
238 async fn adjust_balance(
239 &self,
240 wallet_id: i64,
241 delta_points: i64,
242 kind: &str,
243 idempotency_key: &str,
244 reason: Option<String>,
245 operator: Option<String>,
246 ) -> Result<i64, LproviderError> {
247 let db = self.db()?;
248 let writer = db.writer();
249 let txn = writer.begin().await.map_err(db_err)?;
250 let exists = llm_wallet_ledger::Entity::find()
251 .filter(
252 Condition::all()
253 .add(llm_wallet_ledger::Column::WalletId.eq(wallet_id))
254 .add(llm_wallet_ledger::Column::Kind.eq(kind.to_string()))
255 .add(
256 llm_wallet_ledger::Column::IdempotencyKey
257 .eq(Some(idempotency_key.to_string())),
258 ),
259 )
260 .one(&txn)
261 .await
262 .map_err(db_err)?;
263 if let Some(existing) = exists {
264 txn.rollback().await.map_err(db_err)?;
265 return Ok(existing.balance_after);
266 }
267 let row = txn
268 .query_one_raw(Statement::from_sql_and_values(
269 DatabaseBackend::Postgres,
270 "UPDATE llm_wallet SET balance_points = balance_points + $1, total_topped_up = total_topped_up + CASE WHEN $2 = 'topup' AND $1 > 0 THEN $1 ELSE 0 END, total_refunded = total_refunded + CASE WHEN $2 = 'refund' AND $1 > 0 THEN $1 ELSE 0 END, last_topup_at = now(), updated_at = now() WHERE wallet_id = $3 RETURNING balance_points AS balance_after",
271 [Value::BigInt(Some(delta_points)), Value::String(Some(kind.to_string())), Value::BigInt(Some(wallet_id))],
272 ))
273 .await
274 .map_err(db_err)?;
275 let Some(result) = row else {
276 txn.rollback().await.map_err(db_err)?;
277 return Err(LproviderError::Wallet(WalletError::WalletNotFound(
278 wallet_id,
279 )));
280 };
281 let balance_after: i64 = result.try_get("", "balance_after").unwrap_or_default();
282 let ledger: llm_wallet_ledger::ActiveModel = llm_wallet_ledger::ActiveModel::builder()
283 .set_wallet_id(wallet_id)
284 .set_request_id(None)
285 .set_delta_points(delta_points)
286 .set_balance_after(balance_after)
287 .set_kind(kind.to_string())
288 .set_idempotency_key(Some(idempotency_key.to_string()))
289 .set_reason(reason)
290 .set_operator(operator)
291 .set_created_at(OffsetDateTime::now_utc())
292 .into();
293 ledger.insert(&txn).await.map_err(db_err)?;
294 txn.commit().await.map_err(db_err)?;
295 Ok(balance_after)
296 }
297
298 pub async fn refund(
299 &self,
300 wallet_id: i64,
301 points: i64,
302 request_id: Uuid,
303 reason: Option<String>,
304 operator: Option<String>,
305 ) -> Result<i64, LproviderError> {
306 let db = self.db()?;
307 let writer = db.writer();
308 let txn = writer.begin().await.map_err(db_err)?;
309 let exists = llm_wallet_ledger::Entity::find()
310 .filter(
311 Condition::all()
312 .add(llm_wallet_ledger::Column::RequestId.eq(Some(request_id)))
313 .add(llm_wallet_ledger::Column::Kind.eq("refund")),
314 )
315 .one(&txn)
316 .await
317 .map_err(db_err)?;
318 if let Some(existing) = exists {
319 txn.rollback().await.map_err(db_err)?;
320 return Ok(existing.balance_after);
321 }
322 let row = txn
323 .query_one_raw(Statement::from_sql_and_values(
324 DatabaseBackend::Postgres,
325 "UPDATE llm_wallet SET balance_points = balance_points + $1, total_refunded = total_refunded + $1, updated_at = now() WHERE wallet_id = $2 RETURNING balance_points AS balance_after",
326 [Value::BigInt(Some(points)), Value::BigInt(Some(wallet_id))],
327 ))
328 .await
329 .map_err(db_err)?;
330 let Some(result) = row else {
331 txn.rollback().await.map_err(db_err)?;
332 return Err(LproviderError::Wallet(WalletError::WalletNotFound(
333 wallet_id,
334 )));
335 };
336 let balance_after: i64 = result.try_get("", "balance_after").unwrap_or_default();
337 let ledger: llm_wallet_ledger::ActiveModel = llm_wallet_ledger::ActiveModel::builder()
338 .set_wallet_id(wallet_id)
339 .set_request_id(Some(request_id))
340 .set_delta_points(points)
341 .set_balance_after(balance_after)
342 .set_kind("refund".to_string())
343 .set_reason(reason)
344 .set_operator(operator)
345 .set_created_at(OffsetDateTime::now_utc())
346 .into();
347 ledger.insert(&txn).await.map_err(db_err)?;
348 txn.commit().await.map_err(db_err)?;
349 Ok(balance_after)
350 }
351
352 pub async fn set_status(
353 &self,
354 wallet_id: i64,
355 status: &str,
356 reason: String,
357 operator: Option<String>,
358 ) -> Result<(), LproviderError> {
359 let db = self.db()?;
360 let writer = db.writer();
361 let txn = writer.begin().await.map_err(db_err)?;
362 let updated = txn
363 .execute_raw(Statement::from_sql_and_values(
364 DatabaseBackend::Postgres,
365 "UPDATE llm_wallet SET status = $1, updated_at = now() WHERE wallet_id = $2",
366 [
367 Value::String(Some(status.to_string())),
368 Value::BigInt(Some(wallet_id)),
369 ],
370 ))
371 .await
372 .map_err(db_err)?;
373 if updated.rows_affected() == 0 {
374 txn.rollback().await.map_err(db_err)?;
375 return Err(LproviderError::Wallet(WalletError::WalletNotFound(
376 wallet_id,
377 )));
378 }
379 let log: llm_wallet_freeze_log::ActiveModel = llm_wallet_freeze_log::ActiveModel::builder()
380 .set_wallet_id(wallet_id)
381 .set_action(status.to_string())
382 .set_reason(reason)
383 .set_operator(operator)
384 .set_created_at(OffsetDateTime::now_utc())
385 .into();
386 log.insert(&txn).await.map_err(db_err)?;
387 txn.commit().await.map_err(db_err)?;
388 Ok(())
389 }
390
391 pub async fn freeze(
392 &self,
393 wallet_id: i64,
394 reason: String,
395 operator: Option<String>,
396 ) -> Result<(), LproviderError> {
397 self.set_status(wallet_id, "frozen", reason, operator).await
398 }
399
400 pub async fn unfreeze(
401 &self,
402 wallet_id: i64,
403 reason: String,
404 operator: Option<String>,
405 ) -> Result<(), LproviderError> {
406 self.set_status(wallet_id, "active", reason, operator).await
407 }
408
409 pub async fn disable(
410 &self,
411 wallet_id: i64,
412 reason: String,
413 operator: Option<String>,
414 ) -> Result<(), LproviderError> {
415 self.set_status(wallet_id, "disabled", reason, operator)
416 .await
417 }
418
419 pub async fn list_ledger(
420 &self,
421 wallet_id: Option<i64>,
422 kind: Option<String>,
423 limit: u64,
424 offset: u64,
425 ) -> Result<(Vec<LedgerEntry>, u64), LproviderError> {
426 let db = self.db()?;
427 let mut cond = sea_orm::Condition::all();
428 if let Some(w) = wallet_id {
429 cond = cond.add(llm_wallet_ledger::Column::WalletId.eq(w));
430 }
431 if let Some(k) = kind {
432 cond = cond.add(llm_wallet_ledger::Column::Kind.eq(k));
433 }
434 let total = llm_wallet_ledger::Entity::find()
435 .filter(cond.clone())
436 .count(db.primary_connection())
437 .await
438 .map_err(db_err)?;
439 let models = llm_wallet_ledger::Entity::find()
440 .filter(cond)
441 .order_by_desc(llm_wallet_ledger::Column::Id)
442 .limit(limit)
443 .offset(offset)
444 .all(db.primary_connection())
445 .await
446 .map_err(db_err)?;
447 Ok((
448 models
449 .into_iter()
450 .map(|m| LedgerEntry {
451 wallet_id: m.wallet_id,
452 request_id: m.request_id,
453 delta_points: m.delta_points,
454 balance_after: m.balance_after,
455 kind: m.kind,
456 idempotency_key: m.idempotency_key,
457 reason: m.reason,
458 operator: m.operator,
459 created_at: m.created_at,
460 })
461 .collect(),
462 total,
463 ))
464 }
465
466 pub async fn list_freeze_log(
467 &self,
468 wallet_id: Option<i64>,
469 limit: u64,
470 offset: u64,
471 ) -> Result<(Vec<llm_wallet_freeze_log::Model>, u64), LproviderError> {
472 let db = self.db()?;
473 let mut cond = sea_orm::Condition::all();
474 if let Some(w) = wallet_id {
475 cond = cond.add(llm_wallet_freeze_log::Column::WalletId.eq(w));
476 }
477 let total = llm_wallet_freeze_log::Entity::find()
478 .filter(cond.clone())
479 .count(db.primary_connection())
480 .await
481 .map_err(db_err)?;
482 let models = llm_wallet_freeze_log::Entity::find()
483 .filter(cond)
484 .order_by_desc(llm_wallet_freeze_log::Column::Id)
485 .limit(limit)
486 .offset(offset)
487 .all(db.primary_connection())
488 .await
489 .map_err(db_err)?;
490 Ok((models, total))
491 }
492}
493
494fn wallet_from_model(m: llm_wallet::Model) -> WalletInfo {
495 WalletInfo {
496 wallet_id: m.wallet_id,
497 owner_type: m.owner_type,
498 owner_id: m.owner_id,
499 tenant_id: m.tenant_id,
500 status: m.status,
501 balance_points: m.balance_points,
502 total_topped_up: m.total_topped_up,
503 total_consumed: m.total_consumed,
504 total_refunded: m.total_refunded,
505 created_at: m.created_at,
506 updated_at: m.updated_at,
507 }
508}
509
510fn db_err(err: sea_orm::DbErr) -> LproviderError {
511 LproviderError::Db(db::conn::DbError::InvalidConfig(format!("sea-orm: {err}")))
512}