gproxy_transform/transform/
routing.rs1use serde_json::Value;
5
6use crate::protocol::{Operation, OperationKey, OperationKind};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum RuleImpl {
11 Passthrough,
12 TransformTo,
13 Local,
14 Unsupported,
15}
16
17#[derive(Debug, Clone)]
19pub struct CompiledRoutingRule {
20 pub operation: Operation,
21 pub kind: OperationKind,
22 pub implementation: RuleImpl,
23 pub dest_operation: Option<Operation>,
24 pub dest_kind: Option<OperationKind>,
25}
26
27pub struct RoutingRuleSpec<'a> {
29 pub id: i64,
30 pub provider_id: i64,
31 pub operation: &'a str,
32 pub kind: &'a str,
33 pub implementation: &'a str,
34 pub dest_operation: Option<&'a str>,
35 pub dest_kind: Option<&'a str>,
36 pub sort_order: i64,
37 pub enabled: bool,
38}
39
40pub fn compile(rows: &[RoutingRuleSpec<'_>]) -> Vec<CompiledRoutingRule> {
43 let mut rows: Vec<&RoutingRuleSpec<'_>> = rows.iter().filter(|r| r.enabled).collect();
44 rows.sort_by_key(|r| r.sort_order);
45 let mut out = Vec::new();
46 for row in rows {
47 match compile_row(row) {
48 Some(rule) => out.push(rule),
49 None => tracing::warn!(
50 rule_id = row.id,
51 provider_id = row.provider_id,
52 "skipping unparsable routing rule"
53 ),
54 }
55 }
56 out
57}
58
59fn compile_row(row: &RoutingRuleSpec<'_>) -> Option<CompiledRoutingRule> {
60 Some(CompiledRoutingRule {
61 operation: parse_str(row.operation)?,
62 kind: parse_str(row.kind)?,
63 implementation: match row.implementation {
64 "passthrough" => RuleImpl::Passthrough,
65 "transform_to" => RuleImpl::TransformTo,
66 "local" => RuleImpl::Local,
67 "unsupported" => RuleImpl::Unsupported,
68 _ => return None,
69 },
70 dest_operation: match row.dest_operation {
71 Some(s) => Some(parse_str(s)?),
72 None => None,
73 },
74 dest_kind: match row.dest_kind {
75 Some(s) => Some(parse_str(s)?),
76 None => None,
77 },
78 })
79}
80
81fn parse_str<T: serde::de::DeserializeOwned>(s: &str) -> Option<T> {
86 serde_json::from_value(Value::String(s.to_owned())).ok()
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum RoutingDecision {
92 Passthrough,
93 TransformTo(OperationKey),
94 Local,
95 Unsupported,
96}
97
98pub fn decide(rules: &[CompiledRoutingRule], source: OperationKey) -> RoutingDecision {
105 if let Some(rule) = rules
106 .iter()
107 .find(|r| r.operation == source.operation && r.kind == source.kind)
108 {
109 return match rule.implementation {
110 RuleImpl::Passthrough => RoutingDecision::Passthrough,
111 RuleImpl::Local => RoutingDecision::Local,
112 RuleImpl::Unsupported => RoutingDecision::Unsupported,
113 RuleImpl::TransformTo => match rule.dest_kind {
114 Some(kind) => RoutingDecision::TransformTo(OperationKey {
115 operation: rule.dest_operation.unwrap_or(source.operation),
116 kind,
117 }),
118 None => RoutingDecision::Unsupported,
119 },
120 };
121 }
122 RoutingDecision::Unsupported
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128 use crate::protocol::ContentGenerationKind;
129
130 fn cg(op: Operation, k: ContentGenerationKind) -> OperationKey {
131 OperationKey::content_generation(op, k)
132 }
133
134 #[test]
135 fn no_rule_is_unsupported() {
136 use ContentGenerationKind as K;
137 let src = cg(Operation::GenerateContent, K::ClaudeMessages);
139 assert_eq!(decide(&[], src), RoutingDecision::Unsupported);
140 }
141
142 #[test]
143 fn explicit_rule_wins() {
144 let rule = CompiledRoutingRule {
145 operation: Operation::GenerateContent,
146 kind: OperationKind::ContentGeneration(ContentGenerationKind::ClaudeMessages),
147 implementation: RuleImpl::Unsupported,
148 dest_operation: None,
149 dest_kind: None,
150 };
151 let src = cg(
152 Operation::GenerateContent,
153 ContentGenerationKind::ClaudeMessages,
154 );
155 assert_eq!(decide(&[rule], src), RoutingDecision::Unsupported);
156 }
157}