mirror of
https://github.com/whit3rabbit/anyllm-proxy.git
synced 2026-09-21 16:00:49 +00:00
Large refactor of files over 400+ of lines of code to make it more managable. Just clean up/splitting of files.
This commit is contained in:
whit3rabbit
@@ -875,6 +875,30 @@
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6",
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-luna",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"output_cost_per_token": 6e-06,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-sol",
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-terra",
|
||||
"input_cost_per_token": 2.5e-06,
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-audio",
|
||||
"input_cost_per_token": 2.5e-06,
|
||||
@@ -929,6 +953,18 @@
|
||||
"output_cost_per_token": 1.6e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2.1",
|
||||
"input_cost_per_token": 4e-06,
|
||||
"output_cost_per_token": 2.4e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2.1-mini",
|
||||
"input_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2025-08-28",
|
||||
"input_cost_per_token": 4e-06,
|
||||
|
||||
@@ -280,350 +280,5 @@ fn decisions_hash(edits: &[(usize, BufferId, EditScript)]) -> u64 {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::budget::HeuristicBudgetCounter;
|
||||
use crate::traits::{CacheModel, Pricing, UniformScorer};
|
||||
use crate::types::{Message, Role};
|
||||
|
||||
struct TestStrategy;
|
||||
impl CacheStrategy for TestStrategy {
|
||||
fn pricing(&self) -> Pricing {
|
||||
Pricing {
|
||||
input: 1.0,
|
||||
cached_read: 0.5,
|
||||
cache_write_mult: 1.0,
|
||||
}
|
||||
}
|
||||
fn model(&self) -> CacheModel {
|
||||
CacheModel::ExplicitBreakpoints
|
||||
}
|
||||
fn breakpoint_at(&self, frontier: usize) -> Option<usize> {
|
||||
Some(frontier)
|
||||
}
|
||||
}
|
||||
|
||||
fn long_user(text: &str) -> Message {
|
||||
Message {
|
||||
role: Role::User,
|
||||
blocks: vec![ContentBlock::Text(text.into())],
|
||||
protection: Protection::Mutable,
|
||||
client_cache_marker: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn convo(n: usize) -> Conversation {
|
||||
let long = "The quick brown fox jumps over the lazy dog again and again across \
|
||||
the wide green field toward the distant blue mountains beyond the river \
|
||||
and the tall dark trees under a bright and cloudless summer sky at noon."
|
||||
.to_string();
|
||||
Conversation::new((0..n).map(|_| long_user(&long)).collect())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_never_renders_but_reports() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Shadow,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.rendered.is_none());
|
||||
assert_eq!(out.report.mode, Mode::Shadow);
|
||||
assert!(out.report.frontier > 0);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_renders_and_is_deterministic() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let a = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
let b = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(a.rendered.is_some());
|
||||
assert_eq!(a.report.decisions_hash, b.report.decisions_hash);
|
||||
assert!(a.report.applied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_delta_is_signed_and_nonzero_for_compression() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.report.applied);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
let expected = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&TestStrategy.pricing(),
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
assert!(expected > 0.0, "fixture should have a positive delta");
|
||||
assert_eq!(out.report.est_cost_delta_usd, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_delta_is_zero_for_noop() {
|
||||
// No message reaches `min_len`, so no edits are produced and dt stays 0 — the
|
||||
// report's delta must be exactly 0.0, not the raw formula's rewrite-cost artifact.
|
||||
let conv = Conversation::new(vec![long_user("too short to compress")]);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(!out.report.applied);
|
||||
assert_eq!(out.report.removed_tokens_est, 0);
|
||||
assert_eq!(out.report.est_cost_delta_usd, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_override_turns_off_one_route_while_another_still_compresses() {
|
||||
use crate::policy::{OptimizationPolicy, RouteOverride};
|
||||
use std::collections::HashMap;
|
||||
|
||||
let conv = convo(20);
|
||||
let mut routes = HashMap::new();
|
||||
routes.insert(
|
||||
"batch".to_string(),
|
||||
RouteOverride {
|
||||
mode: Some(Mode::Off),
|
||||
ratios: None,
|
||||
pricing: None,
|
||||
},
|
||||
);
|
||||
let opt_policy = OptimizationPolicy {
|
||||
mode: Mode::Live,
|
||||
routes,
|
||||
..OptimizationPolicy::default()
|
||||
};
|
||||
|
||||
// Overridden route: rendered stays None, i.e. output ≡ input for that route.
|
||||
let mut ws = Workspace::new();
|
||||
let off = optimize_for_route(
|
||||
&conv,
|
||||
&opt_policy,
|
||||
"batch",
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(off.rendered.is_none());
|
||||
assert!(!off.report.applied);
|
||||
|
||||
// Unlisted route: falls back to the top-level Live default and still compresses.
|
||||
let mut ws2 = Workspace::new();
|
||||
let on = optimize_for_route(
|
||||
&conv,
|
||||
&opt_policy,
|
||||
"interactive",
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws2,
|
||||
);
|
||||
assert!(on.rendered.is_some());
|
||||
assert!(on.report.applied);
|
||||
}
|
||||
|
||||
/// Scores like `UniformScorer` but sleeps on its very first call, so the deadline
|
||||
/// check ahead of every later message is guaranteed to observe it expired
|
||||
/// regardless of scheduler jitter (sleep duration >> deadline budget below).
|
||||
struct SlowFirstCallScorer {
|
||||
calls: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
impl TokenScorer for SlowFirstCallScorer {
|
||||
fn score_words(
|
||||
&self,
|
||||
words: &[&str],
|
||||
_ws: &mut Workspace,
|
||||
) -> Result<Vec<f32>, crate::error::ScoreError> {
|
||||
if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
Ok(vec![0.5; words.len()])
|
||||
}
|
||||
fn artifact_hash(&self) -> u64 {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deadline_expiry_leaves_later_messages_byte_identical() {
|
||||
// 20 messages -> frontier eligible_end=16 (keep_recent=4), all long enough to
|
||||
// clear `min_len`. A 50ms sleep on the first scorer call vs. a 10ms deadline
|
||||
// guarantees message 0 is scored before the deadline and every later message
|
||||
// observes it already expired (elapsed only ever grows).
|
||||
let conv = convo(20);
|
||||
let mut policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
policy.compression.deadline = Duration::from_millis(10);
|
||||
|
||||
let scorer = SlowFirstCallScorer {
|
||||
calls: std::sync::atomic::AtomicUsize::new(0),
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&scorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
|
||||
assert!(
|
||||
out.report.messages_compressed >= 1,
|
||||
"message scored before the deadline should be compressed"
|
||||
);
|
||||
assert!(
|
||||
out.report.messages_skipped_deadline >= 1,
|
||||
"messages after deadline expiry must be counted as skipped"
|
||||
);
|
||||
assert_eq!(
|
||||
out.report.messages_compressed as usize + out.report.messages_skipped_deadline as usize,
|
||||
out.report.frontier,
|
||||
"every eligible message is either compressed or explicitly deadline-skipped, never silently dropped"
|
||||
);
|
||||
|
||||
let rendered = out
|
||||
.rendered
|
||||
.expect("gate should apply given a nonzero saving");
|
||||
let compressed = out.report.messages_compressed as usize;
|
||||
// Deadline-skipped messages must render byte-identical to the source.
|
||||
for i in compressed..out.report.frontier {
|
||||
assert_eq!(
|
||||
rendered.messages[i].blocks[0], conv.messages[i].blocks[0],
|
||||
"deadline-skipped message {i} must stay byte-identical"
|
||||
);
|
||||
}
|
||||
// The message scored before the deadline actually got edited.
|
||||
assert_ne!(
|
||||
rendered.messages[0].blocks[0], conv.messages[0].blocks[0],
|
||||
"the message scored before the deadline should have been edited"
|
||||
);
|
||||
}
|
||||
|
||||
/// M4.3: `Policy::pricing_override`, loaded from a config string via
|
||||
/// `Pricing::from_config_str` (no hardcoded-table constant), must be what the
|
||||
/// orchestrator actually uses for the cost gate/report — not `TestStrategy`'s own
|
||||
/// (hardcoded-in-code) `pricing()`. Proven by asserting the report's dollar delta
|
||||
/// matches `net_cost_delta_usd` computed with the config pricing and differs from
|
||||
/// the value that would result from `TestStrategy::pricing()` alone.
|
||||
#[test]
|
||||
fn pricing_comes_from_config() {
|
||||
use crate::traits::Pricing;
|
||||
|
||||
let conv = convo(20);
|
||||
|
||||
// Deliberately far from `TestStrategy::pricing()` (input:1.0, cached_read:0.5,
|
||||
// cache_write_mult:1.0) and from any of `anyllm_optimize_passes::cost_gate`'s
|
||||
// hardcoded tables — proves the number really came from this config string.
|
||||
let config_pricing =
|
||||
Pricing::from_config_str("input=9.0\ncached_read=0.05\ncache_write_mult=2.0\n")
|
||||
.expect("well-formed config parses");
|
||||
assert_ne!(config_pricing, TestStrategy.pricing());
|
||||
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
pricing_override: Some(config_pricing),
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.report.applied);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
|
||||
let expected_from_config = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&config_pricing,
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
let would_be_from_hardcoded_strategy = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&TestStrategy.pricing(),
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
|
||||
assert_eq!(out.report.est_cost_delta_usd, expected_from_config);
|
||||
assert_ne!(
|
||||
out.report.est_cost_delta_usd, would_be_from_hardcoded_strategy,
|
||||
"orchestrator must use the config-loaded pricing, not the strategy's hardcoded table"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_conversation_is_noop() {
|
||||
let conv = Conversation::default();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&Policy::default(),
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut Workspace::new(),
|
||||
);
|
||||
assert!(out.rendered.is_none());
|
||||
assert_eq!(out.report.frontier, 0);
|
||||
}
|
||||
}
|
||||
#[path = "orchestrator_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
use super::*;
|
||||
use crate::budget::HeuristicBudgetCounter;
|
||||
use crate::traits::{CacheModel, Pricing, UniformScorer};
|
||||
use crate::types::{Message, Role};
|
||||
|
||||
struct TestStrategy;
|
||||
impl CacheStrategy for TestStrategy {
|
||||
fn pricing(&self) -> Pricing {
|
||||
Pricing {
|
||||
input: 1.0,
|
||||
cached_read: 0.5,
|
||||
cache_write_mult: 1.0,
|
||||
}
|
||||
}
|
||||
fn model(&self) -> CacheModel {
|
||||
CacheModel::ExplicitBreakpoints
|
||||
}
|
||||
fn breakpoint_at(&self, frontier: usize) -> Option<usize> {
|
||||
Some(frontier)
|
||||
}
|
||||
}
|
||||
|
||||
fn long_user(text: &str) -> Message {
|
||||
Message {
|
||||
role: Role::User,
|
||||
blocks: vec![ContentBlock::Text(text.into())],
|
||||
protection: Protection::Mutable,
|
||||
client_cache_marker: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn convo(n: usize) -> Conversation {
|
||||
let long = "The quick brown fox jumps over the lazy dog again and again across \
|
||||
the wide green field toward the distant blue mountains beyond the river \
|
||||
and the tall dark trees under a bright and cloudless summer sky at noon."
|
||||
.to_string();
|
||||
Conversation::new((0..n).map(|_| long_user(&long)).collect())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_never_renders_but_reports() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Shadow,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.rendered.is_none());
|
||||
assert_eq!(out.report.mode, Mode::Shadow);
|
||||
assert!(out.report.frontier > 0);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_renders_and_is_deterministic() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let a = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
let b = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(a.rendered.is_some());
|
||||
assert_eq!(a.report.decisions_hash, b.report.decisions_hash);
|
||||
assert!(a.report.applied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_delta_is_signed_and_nonzero_for_compression() {
|
||||
let conv = convo(20);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.report.applied);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
let expected = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&TestStrategy.pricing(),
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
assert!(expected > 0.0, "fixture should have a positive delta");
|
||||
assert_eq!(out.report.est_cost_delta_usd, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_delta_is_zero_for_noop() {
|
||||
// No message reaches `min_len`, so no edits are produced and dt stays 0 — the
|
||||
// report's delta must be exactly 0.0, not the raw formula's rewrite-cost artifact.
|
||||
let conv = Conversation::new(vec![long_user("too short to compress")]);
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(!out.report.applied);
|
||||
assert_eq!(out.report.removed_tokens_est, 0);
|
||||
assert_eq!(out.report.est_cost_delta_usd, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_override_turns_off_one_route_while_another_still_compresses() {
|
||||
use crate::policy::{OptimizationPolicy, RouteOverride};
|
||||
use std::collections::HashMap;
|
||||
|
||||
let conv = convo(20);
|
||||
let mut routes = HashMap::new();
|
||||
routes.insert(
|
||||
"batch".to_string(),
|
||||
RouteOverride {
|
||||
mode: Some(Mode::Off),
|
||||
ratios: None,
|
||||
pricing: None,
|
||||
},
|
||||
);
|
||||
let opt_policy = OptimizationPolicy {
|
||||
mode: Mode::Live,
|
||||
routes,
|
||||
..OptimizationPolicy::default()
|
||||
};
|
||||
|
||||
// Overridden route: rendered stays None, i.e. output ≡ input for that route.
|
||||
let mut ws = Workspace::new();
|
||||
let off = optimize_for_route(
|
||||
&conv,
|
||||
&opt_policy,
|
||||
"batch",
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(off.rendered.is_none());
|
||||
assert!(!off.report.applied);
|
||||
|
||||
// Unlisted route: falls back to the top-level Live default and still compresses.
|
||||
let mut ws2 = Workspace::new();
|
||||
let on = optimize_for_route(
|
||||
&conv,
|
||||
&opt_policy,
|
||||
"interactive",
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws2,
|
||||
);
|
||||
assert!(on.rendered.is_some());
|
||||
assert!(on.report.applied);
|
||||
}
|
||||
|
||||
/// Scores like `UniformScorer` but sleeps on its very first call, so the deadline
|
||||
/// check ahead of every later message is guaranteed to observe it expired
|
||||
/// regardless of scheduler jitter (sleep duration >> deadline budget below).
|
||||
struct SlowFirstCallScorer {
|
||||
calls: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
impl TokenScorer for SlowFirstCallScorer {
|
||||
fn score_words(
|
||||
&self,
|
||||
words: &[&str],
|
||||
_ws: &mut Workspace,
|
||||
) -> Result<Vec<f32>, crate::error::ScoreError> {
|
||||
if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
Ok(vec![0.5; words.len()])
|
||||
}
|
||||
fn artifact_hash(&self) -> u64 {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deadline_expiry_leaves_later_messages_byte_identical() {
|
||||
// 20 messages -> frontier eligible_end=16 (keep_recent=4), all long enough to
|
||||
// clear `min_len`. A 50ms sleep on the first scorer call vs. a 10ms deadline
|
||||
// guarantees message 0 is scored before the deadline and every later message
|
||||
// observes it already expired (elapsed only ever grows).
|
||||
let conv = convo(20);
|
||||
let mut policy = Policy {
|
||||
mode: Mode::Live,
|
||||
..Default::default()
|
||||
};
|
||||
policy.compression.deadline = Duration::from_millis(10);
|
||||
|
||||
let scorer = SlowFirstCallScorer {
|
||||
calls: std::sync::atomic::AtomicUsize::new(0),
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&scorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
|
||||
assert!(
|
||||
out.report.messages_compressed >= 1,
|
||||
"message scored before the deadline should be compressed"
|
||||
);
|
||||
assert!(
|
||||
out.report.messages_skipped_deadline >= 1,
|
||||
"messages after deadline expiry must be counted as skipped"
|
||||
);
|
||||
assert_eq!(
|
||||
out.report.messages_compressed as usize + out.report.messages_skipped_deadline as usize,
|
||||
out.report.frontier,
|
||||
"every eligible message is either compressed or explicitly deadline-skipped, never silently dropped"
|
||||
);
|
||||
|
||||
let rendered = out
|
||||
.rendered
|
||||
.expect("gate should apply given a nonzero saving");
|
||||
let compressed = out.report.messages_compressed as usize;
|
||||
// Deadline-skipped messages must render byte-identical to the source.
|
||||
for i in compressed..out.report.frontier {
|
||||
assert_eq!(
|
||||
rendered.messages[i].blocks[0], conv.messages[i].blocks[0],
|
||||
"deadline-skipped message {i} must stay byte-identical"
|
||||
);
|
||||
}
|
||||
// The message scored before the deadline actually got edited.
|
||||
assert_ne!(
|
||||
rendered.messages[0].blocks[0], conv.messages[0].blocks[0],
|
||||
"the message scored before the deadline should have been edited"
|
||||
);
|
||||
}
|
||||
|
||||
/// M4.3: `Policy::pricing_override`, loaded from a config string via
|
||||
/// `Pricing::from_config_str` (no hardcoded-table constant), must be what the
|
||||
/// orchestrator actually uses for the cost gate/report — not `TestStrategy`'s own
|
||||
/// (hardcoded-in-code) `pricing()`. Proven by asserting the report's dollar delta
|
||||
/// matches `net_cost_delta_usd` computed with the config pricing and differs from
|
||||
/// the value that would result from `TestStrategy::pricing()` alone.
|
||||
#[test]
|
||||
fn pricing_comes_from_config() {
|
||||
use crate::traits::Pricing;
|
||||
|
||||
let conv = convo(20);
|
||||
|
||||
// Deliberately far from `TestStrategy::pricing()` (input:1.0, cached_read:0.5,
|
||||
// cache_write_mult:1.0) and from any of `anyllm_optimize_passes::cost_gate`'s
|
||||
// hardcoded tables — proves the number really came from this config string.
|
||||
let config_pricing =
|
||||
Pricing::from_config_str("input=9.0\ncached_read=0.05\ncache_write_mult=2.0\n")
|
||||
.expect("well-formed config parses");
|
||||
assert_ne!(config_pricing, TestStrategy.pricing());
|
||||
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
pricing_override: Some(config_pricing),
|
||||
..Default::default()
|
||||
};
|
||||
let mut ws = Workspace::new();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&policy,
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
assert!(out.report.applied);
|
||||
assert!(out.report.removed_tokens_est > 0);
|
||||
|
||||
let expected_from_config = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&config_pricing,
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
let would_be_from_hardcoded_strategy = net_cost_delta_usd(
|
||||
out.report.removed_tokens_est,
|
||||
out.report.rewrite_suffix_tokens,
|
||||
policy.horizon,
|
||||
&TestStrategy.pricing(),
|
||||
&TestStrategy.model(),
|
||||
);
|
||||
|
||||
assert_eq!(out.report.est_cost_delta_usd, expected_from_config);
|
||||
assert_ne!(
|
||||
out.report.est_cost_delta_usd, would_be_from_hardcoded_strategy,
|
||||
"orchestrator must use the config-loaded pricing, not the strategy's hardcoded table"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_conversation_is_noop() {
|
||||
let conv = Conversation::default();
|
||||
let out = optimize(
|
||||
&conv,
|
||||
&Policy::default(),
|
||||
&TestStrategy,
|
||||
&UniformScorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&mut Workspace::new(),
|
||||
);
|
||||
assert!(out.rendered.is_none());
|
||||
assert_eq!(out.report.frontier, 0);
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
use super::{Api, Args};
|
||||
use anyhow::{Context, Result};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// One provider response: token usage + the assistant text.
|
||||
pub struct Reply {
|
||||
pub prompt_tokens: u64,
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// Send one request and return usage + assistant text.
|
||||
pub fn send_request(
|
||||
client: &reqwest::blocking::Client,
|
||||
args: &Args,
|
||||
body: &Value,
|
||||
base_url: &str,
|
||||
api: Api,
|
||||
) -> Result<Reply> {
|
||||
let mut body = body.clone();
|
||||
let url = match api {
|
||||
Api::Openai => format!("{}/chat/completions", base_url.trim_end_matches('/')),
|
||||
Api::Anthropic => format!("{}/messages", base_url.trim_end_matches('/')),
|
||||
};
|
||||
if args.stream {
|
||||
body["stream"] = json!(true);
|
||||
if api == Api::Openai {
|
||||
body["stream_options"] = json!({ "include_usage": true });
|
||||
}
|
||||
}
|
||||
|
||||
let key = resolve_key(args);
|
||||
let mut req = client.post(&url).json(&body);
|
||||
req = match api {
|
||||
Api::Openai => {
|
||||
if key.is_empty() {
|
||||
req
|
||||
} else {
|
||||
req.bearer_auth(key)
|
||||
}
|
||||
}
|
||||
Api::Anthropic => req
|
||||
.header("x-api-key", key)
|
||||
.header("anthropic-version", "2023-06-01"),
|
||||
};
|
||||
|
||||
let resp = req.send().context("request send failed")?;
|
||||
let status = resp.status();
|
||||
let text = resp.text().context("reading response body")?;
|
||||
if !status.is_success() {
|
||||
anyhow::bail!(
|
||||
"HTTP {status}: {}",
|
||||
text.chars().take(300).collect::<String>()
|
||||
);
|
||||
}
|
||||
Ok(Reply {
|
||||
prompt_tokens: extract_prompt_tokens(&text, api, args.stream)
|
||||
.context("could not find prompt/input token usage")?,
|
||||
text: extract_response_text(&text, api, args.stream),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_key(args: &Args) -> String {
|
||||
args.api_key
|
||||
.clone()
|
||||
.or_else(|| std::env::var("OPENROUTER_API_KEY").ok())
|
||||
.or_else(|| std::env::var("ANYLLM_API_KEY").ok())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Pull prompt tokens from a JSON body or an SSE stream.
|
||||
pub fn extract_prompt_tokens(text: &str, api: Api, stream: bool) -> Result<u64> {
|
||||
let want = |v: &Value| -> Option<u64> {
|
||||
match api {
|
||||
Api::Openai => v.get("usage")?.get("prompt_tokens")?.as_u64(),
|
||||
Api::Anthropic => v
|
||||
.get("usage")
|
||||
.or_else(|| v.get("message").and_then(|m| m.get("usage")))?
|
||||
.get("input_tokens")?
|
||||
.as_u64(),
|
||||
}
|
||||
};
|
||||
if !stream {
|
||||
let v: Value = serde_json::from_str(text)?;
|
||||
return want(&v).context("usage field missing");
|
||||
}
|
||||
for payload in sse_payloads(text) {
|
||||
if let Ok(v) = serde_json::from_str::<Value>(&payload) {
|
||||
if let Some(n) = want(&v) {
|
||||
return Ok(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::bail!("no usage in stream")
|
||||
}
|
||||
|
||||
/// Pull assistant text from a JSON body or an SSE stream. Best-effort: returns "" if the
|
||||
/// shape is unrecognized (quality just reads as low, never panics).
|
||||
pub fn extract_response_text(text: &str, api: Api, stream: bool) -> String {
|
||||
if !stream {
|
||||
let Ok(v) = serde_json::from_str::<Value>(text) else {
|
||||
return String::new();
|
||||
};
|
||||
return match api {
|
||||
Api::Openai => v
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
Api::Anthropic => v
|
||||
.get("content")
|
||||
.and_then(|c| c.as_array())
|
||||
.map(|blocks| {
|
||||
blocks
|
||||
.iter()
|
||||
.filter_map(|b| b.get("text").and_then(|t| t.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
};
|
||||
}
|
||||
// Streaming: concatenate incremental text deltas.
|
||||
let mut out = String::new();
|
||||
for payload in sse_payloads(text) {
|
||||
let Ok(v) = serde_json::from_str::<Value>(&payload) else {
|
||||
continue;
|
||||
};
|
||||
match api {
|
||||
Api::Openai => {
|
||||
if let Some(s) = v
|
||||
.pointer("/choices/0/delta/content")
|
||||
.and_then(|c| c.as_str())
|
||||
{
|
||||
out.push_str(s);
|
||||
}
|
||||
}
|
||||
Api::Anthropic => {
|
||||
if let Some(s) = v.pointer("/delta/text").and_then(|c| c.as_str()) {
|
||||
out.push_str(s);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Yield the JSON payload of each `data:` SSE line, skipping `[DONE]`.
|
||||
pub fn sse_payloads(text: &str) -> impl Iterator<Item = String> + '_ {
|
||||
text.lines().filter_map(|line| {
|
||||
let payload = line.trim().strip_prefix("data:")?.trim();
|
||||
(payload != "[DONE]").then(|| payload.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
/// Ask an OpenAI-compatible judge model to score 1-5 how well `comp` preserves `raw`.
|
||||
pub fn run_judge(
|
||||
client: &reqwest::blocking::Client,
|
||||
args: &Args,
|
||||
raw: &str,
|
||||
comp: &str,
|
||||
) -> Result<u32> {
|
||||
let model = args.judge_model.as_ref().context("no judge model")?;
|
||||
let base = args.judge_base_url.as_deref().unwrap_or(&args.base_url);
|
||||
let prompt = format!(
|
||||
"You compare two AI assistant responses to the same user request. Response A is \
|
||||
from the full prompt; Response B is from a compressed prompt. Score 1-5 how well B \
|
||||
preserves A's meaning and quality (5 = equivalent, 1 = badly degraded). Reply with \
|
||||
ONLY the integer.\n\n[A]\n{raw}\n\n[B]\n{comp}"
|
||||
);
|
||||
let body = json!({
|
||||
"model": model,
|
||||
"messages": [{ "role": "user", "content": prompt }],
|
||||
"max_tokens": 4,
|
||||
"temperature": 0,
|
||||
});
|
||||
let url = format!("{}/chat/completions", base.trim_end_matches('/'));
|
||||
let key = resolve_key(args);
|
||||
let mut req = client.post(&url).json(&body);
|
||||
if !key.is_empty() {
|
||||
req = req.bearer_auth(key);
|
||||
}
|
||||
let resp = req.send().context("judge send failed")?;
|
||||
let text = resp.text().context("reading judge response")?;
|
||||
let v: Value = serde_json::from_str(&text).context("judge response not JSON")?;
|
||||
let content = v
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(|c| c.as_str())
|
||||
.context("judge response has no content")?;
|
||||
let score = content
|
||||
.chars()
|
||||
.find(|c| ('1'..='5').contains(c))
|
||||
.and_then(|c| c.to_digit(10))
|
||||
.context("no 1-5 score in judge reply")?;
|
||||
Ok(score)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use super::{build_scorer, compress, ensure_model, load_inputs, Api, Args};
|
||||
use anyhow::Result;
|
||||
use anyllm_optimize_core::{net_cost_delta_usd, CacheModel, CacheStrategy};
|
||||
use anyllm_optimize_passes::{AnthropicStrategy, OpenAiStrategy};
|
||||
|
||||
/// EH-0002 bite 2: simulate the frozen-frontier policy over `--input` and print the
|
||||
/// signed net USD delta (compress vs skip) the cost gate computes per row, plus a class
|
||||
/// total. Fully offline (no network) — the ΔT/S the gate uses come straight out of the
|
||||
/// same `optimize()` call the harness already runs, so this is not a re-derivation, it is
|
||||
/// the actual decision the frontier + cost gate would make.
|
||||
pub fn run_net_cost(args: &Args) -> Result<()> {
|
||||
let bodies = load_inputs(args)?;
|
||||
let scorer = build_scorer(args);
|
||||
let (pricing, model) = match args.cost_model.unwrap_or(args.api) {
|
||||
Api::Openai => (
|
||||
OpenAiStrategy::default().pricing(),
|
||||
OpenAiStrategy::default().model(),
|
||||
),
|
||||
Api::Anthropic => (
|
||||
AnthropicStrategy::default().pricing(),
|
||||
AnthropicStrategy::default().model(),
|
||||
),
|
||||
};
|
||||
println!(
|
||||
"{:<4} {:>9} {:>7} {:>7} {:>7} {:>14}",
|
||||
"#", "frontier", "dt", "s", "apply", "net_usd"
|
||||
);
|
||||
let mut total_dt = 0u64;
|
||||
let mut total_net = 0.0f64;
|
||||
for (i, mut raw) in bodies.into_iter().enumerate() {
|
||||
ensure_model(&mut raw, args);
|
||||
// dt/s come from `--api`'s adapter (must match the input's wire shape, or tool/JSON
|
||||
// blocks get mis-segmented — see cost_model doc comment); the apply decision is
|
||||
// then re-derived under `--cost-model`'s pricing/CacheModel, since `report.applied`
|
||||
// reflects `--api`'s own strategy and the two can differ.
|
||||
let (_compressed, report) = compress(&raw, args, scorer.as_ref());
|
||||
let dt = report.removed_tokens_est;
|
||||
let s = report.rewrite_suffix_tokens;
|
||||
let applies = anyllm_optimize_core::should_apply(dt, s, args.horizon, &pricing, &model);
|
||||
// Only realize the delta when the gate actually applies: `applies == false` means
|
||||
// the original is forwarded untouched, so the realized cost change is $0, not the
|
||||
// negative "forced a no-op rewrite" value the raw formula would give for dt=0
|
||||
// (which `should_apply` exists precisely to avoid ever choosing).
|
||||
let net = if applies {
|
||||
net_cost_delta_usd(dt, s, args.horizon, &pricing, &model)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
total_dt += dt;
|
||||
total_net += net;
|
||||
println!(
|
||||
"{:<4} {:>9} {:>7} {:>7} {:>7} {:>14.6}",
|
||||
i, report.frontier, dt, s, applies, net
|
||||
);
|
||||
}
|
||||
let model_name = match model {
|
||||
CacheModel::ImplicitPrefix => "implicit-prefix",
|
||||
CacheModel::ExplicitBreakpoints => "explicit-breakpoints",
|
||||
};
|
||||
println!(
|
||||
"\nNET COST (horizon={}, {model_name}, input=${:.2}/Mtok cached_read=${:.2}/Mtok \
|
||||
write_mult={:.2}x): total ΔT={total_dt} tokens, net_delta=${total_net:.6}",
|
||||
args.horizon, pricing.input, pricing.cached_read, pricing.cache_write_mult,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn opt_u64(v: Option<u64>) -> String {
|
||||
v.map(|x| x.to_string()).unwrap_or_else(|| "-".into())
|
||||
}
|
||||
|
||||
pub fn pct(raw: u64, comp: u64) -> f64 {
|
||||
if raw == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(raw.saturating_sub(comp)) as f64 / raw as f64 * 100.0
|
||||
}
|
||||
}
|
||||
@@ -17,12 +17,9 @@
|
||||
//! Fail-open: a network/API error on one row is reported and skipped; the process exits
|
||||
//! non-zero if any row failed, but never panics.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use anyllm_optimize_core::{
|
||||
net_cost_delta_usd, BudgetCounter, CacheModel, CacheStrategy, HeuristicBudgetCounter, Mode,
|
||||
OptimizationReport, Policy, TokenScorer, UniformScorer, Workspace,
|
||||
Mode, OptimizationReport, Policy, TokenScorer, UniformScorer, Workspace,
|
||||
};
|
||||
use anyllm_optimize_passes::adapter::{anthropic, openai};
|
||||
use anyllm_optimize_passes::{
|
||||
@@ -32,10 +29,22 @@ use clap::{Parser, ValueEnum};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
mod budget_tiktoken;
|
||||
pub(crate) mod budget_tiktoken;
|
||||
|
||||
mod client;
|
||||
mod cost;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
mod token_count;
|
||||
mod utils;
|
||||
|
||||
use client::{run_judge, send_request};
|
||||
use cost::{opt_u64, pct, run_net_cost};
|
||||
use token_count::{build_counter, count_local};
|
||||
use utils::{jaccard, wrap_prompt};
|
||||
|
||||
#[derive(Clone, Copy, Debug, ValueEnum, PartialEq)]
|
||||
enum Api {
|
||||
pub(crate) enum Api {
|
||||
Openai,
|
||||
Anthropic,
|
||||
}
|
||||
@@ -45,71 +54,65 @@ enum Api {
|
||||
name = "optimize-eval",
|
||||
about = "FFEC token-savings + response-quality harness"
|
||||
)]
|
||||
struct Args {
|
||||
pub(crate) struct Args {
|
||||
/// Wire format of the target endpoint.
|
||||
#[arg(long, value_enum, default_value_t = Api::Openai)]
|
||||
api: Api,
|
||||
pub(crate) api: Api,
|
||||
/// Base URL. Ollama: http://localhost:11434/v1 ; OpenRouter: https://openrouter.ai/api/v1
|
||||
#[arg(long, default_value = "http://localhost:11434/v1")]
|
||||
base_url: String,
|
||||
pub(crate) base_url: String,
|
||||
/// Target model id.
|
||||
#[arg(long)]
|
||||
model: String,
|
||||
pub(crate) model: String,
|
||||
/// API key. Falls back to $OPENROUTER_API_KEY / $ANYLLM_API_KEY; empty for local.
|
||||
#[arg(long)]
|
||||
api_key: Option<String>,
|
||||
pub(crate) api_key: Option<String>,
|
||||
/// JSONL file: each line a request body (`{"messages":[...]}`) or a bare prompt string.
|
||||
#[arg(long, conflicts_with = "prompt")]
|
||||
input: Option<std::path::PathBuf>,
|
||||
pub(crate) input: Option<std::path::PathBuf>,
|
||||
/// A single prompt to test instead of --input.
|
||||
#[arg(long)]
|
||||
prompt: Option<String>,
|
||||
pub(crate) prompt: Option<String>,
|
||||
/// Use streaming requests (adds include_usage for OpenAI).
|
||||
#[arg(long)]
|
||||
stream: bool,
|
||||
pub(crate) stream: bool,
|
||||
/// max_tokens for the response (Anthropic requires it; also sent to OpenAI).
|
||||
#[arg(long, default_value_t = 256)]
|
||||
max_tokens: u64,
|
||||
pub(crate) max_tokens: u64,
|
||||
/// Expected remaining turns (cost-gate horizon).
|
||||
#[arg(long, default_value_t = 8)]
|
||||
horizon: u64,
|
||||
pub(crate) horizon: u64,
|
||||
/// Only compute local estimates; do not call the network (no quality signal).
|
||||
#[arg(long)]
|
||||
offline: bool,
|
||||
pub(crate) offline: bool,
|
||||
/// Print both full responses per row (raw vs compressed) for eyeballing quality.
|
||||
#[arg(long)]
|
||||
show_responses: bool,
|
||||
pub(crate) show_responses: bool,
|
||||
/// Optional OpenAI-compatible judge model. If set, scores 1-5 how well the compressed
|
||||
/// response preserves the raw response's meaning/quality (adds one call per row).
|
||||
#[arg(long)]
|
||||
judge_model: Option<String>,
|
||||
pub(crate) judge_model: Option<String>,
|
||||
/// Base URL for the judge (defaults to --base-url). Must be OpenAI-compatible.
|
||||
#[arg(long)]
|
||||
judge_base_url: Option<String>,
|
||||
pub(crate) judge_base_url: Option<String>,
|
||||
/// Simulate the frozen-frontier cost gate over the input and print the signed net
|
||||
/// USD delta (compress vs skip) per row plus a class total, using the CacheModel and
|
||||
/// Pricing table implied by `--api` unless `--cost-model` overrides it (EH-0002 M0.3:
|
||||
/// offline, no network needed).
|
||||
#[arg(long)]
|
||||
net_cost: bool,
|
||||
pub(crate) net_cost: bool,
|
||||
/// Override which provider's CacheModel/Pricing the `--net-cost` dollar math uses,
|
||||
/// independent of `--api` (which must still match the input's actual wire shape — the
|
||||
/// adapter mis-segments tool/content blocks if it doesn't). Defaults to `--api`.
|
||||
#[arg(long, value_enum)]
|
||||
cost_model: Option<Api>,
|
||||
pub(crate) cost_model: Option<Api>,
|
||||
/// M3.6: directory containing `model.onnx` + `tokenizer.json` for the real
|
||||
/// LLMLingua2Pass ML scorer (ROADMAP D8), used in place of `UniformScorer` for every
|
||||
/// message the frontier already selects. Requires building with `--features onnx`;
|
||||
/// without it (or on a load failure) this falls back to `UniformScorer` with a
|
||||
/// warning, per the fail-open invariant — never a hard error.
|
||||
#[arg(long)]
|
||||
llmlingua2_model_dir: Option<std::path::PathBuf>,
|
||||
}
|
||||
|
||||
/// One provider response: token usage + the assistant text.
|
||||
struct Reply {
|
||||
prompt_tokens: u64,
|
||||
text: String,
|
||||
pub(crate) llmlingua2_model_dir: Option<std::path::PathBuf>,
|
||||
}
|
||||
|
||||
fn main() -> std::process::ExitCode {
|
||||
@@ -254,7 +257,11 @@ fn run(args: &Args) -> Result<u32> {
|
||||
/// `UniformScorer` by default, or the real `LLMLingua2Pass` when `--llmlingua2-model-dir`
|
||||
/// resolves (see `build_scorer`) — either way it is only ever invoked by `optimize()` for
|
||||
/// frontier-eligible D8 targets (tool results, old RAG blocks, old assistant messages).
|
||||
fn compress(raw: &Value, args: &Args, scorer: &dyn TokenScorer) -> (Value, OptimizationReport) {
|
||||
pub(crate) fn compress(
|
||||
raw: &Value,
|
||||
args: &Args,
|
||||
scorer: &dyn TokenScorer,
|
||||
) -> (Value, OptimizationReport) {
|
||||
let policy = Policy {
|
||||
mode: Mode::Live,
|
||||
horizon: args.horizon,
|
||||
@@ -271,7 +278,7 @@ fn compress(raw: &Value, args: &Args, scorer: &dyn TokenScorer) -> (Value, Optim
|
||||
&policy,
|
||||
&OpenAiStrategy::default(),
|
||||
scorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&anyllm_optimize_core::HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
if let Some(r) = &res.rendered {
|
||||
@@ -286,7 +293,7 @@ fn compress(raw: &Value, args: &Args, scorer: &dyn TokenScorer) -> (Value, Optim
|
||||
&policy,
|
||||
&AnthropicStrategy::default(),
|
||||
scorer,
|
||||
&HeuristicBudgetCounter::default(),
|
||||
&anyllm_optimize_core::HeuristicBudgetCounter::default(),
|
||||
&mut ws,
|
||||
);
|
||||
if let Some(r) = &res.rendered {
|
||||
@@ -305,7 +312,7 @@ fn compress(raw: &Value, args: &Args, scorer: &dyn TokenScorer) -> (Value, Optim
|
||||
/// that gating already lives in `optimize()`/`compress_message` (see `LLMLingua2Pass`'s
|
||||
/// doc comment). Any failure to build/load falls back to `UniformScorer` with a warning
|
||||
/// rather than aborting the run.
|
||||
fn build_scorer(args: &Args) -> Box<dyn TokenScorer> {
|
||||
pub(crate) fn build_scorer(args: &Args) -> Box<dyn TokenScorer> {
|
||||
let Some(dir) = &args.llmlingua2_model_dir else {
|
||||
return Box::new(UniformScorer);
|
||||
};
|
||||
@@ -364,339 +371,16 @@ fn fnv1a64(bytes: &[u8]) -> u64 {
|
||||
.fold(OFFSET, |h, &b| (h ^ b as u64).wrapping_mul(PRIME))
|
||||
}
|
||||
|
||||
/// EH-0002 bite 2: simulate the frozen-frontier policy over `--input` and print the
|
||||
/// signed net USD delta (compress vs skip) the cost gate computes per row, plus a class
|
||||
/// total. Fully offline (no network) — the ΔT/S the gate uses come straight out of the
|
||||
/// same `optimize()` call the harness already runs, so this is not a re-derivation, it is
|
||||
/// the actual decision the frontier + cost gate would make.
|
||||
fn run_net_cost(args: &Args) -> Result<()> {
|
||||
let bodies = load_inputs(args)?;
|
||||
let scorer = build_scorer(args);
|
||||
let (pricing, model) = match args.cost_model.unwrap_or(args.api) {
|
||||
Api::Openai => (
|
||||
OpenAiStrategy::default().pricing(),
|
||||
OpenAiStrategy::default().model(),
|
||||
),
|
||||
Api::Anthropic => (
|
||||
AnthropicStrategy::default().pricing(),
|
||||
AnthropicStrategy::default().model(),
|
||||
),
|
||||
};
|
||||
println!(
|
||||
"{:<4} {:>9} {:>7} {:>7} {:>7} {:>14}",
|
||||
"#", "frontier", "dt", "s", "apply", "net_usd"
|
||||
);
|
||||
let mut total_dt = 0u64;
|
||||
let mut total_net = 0.0f64;
|
||||
for (i, mut raw) in bodies.into_iter().enumerate() {
|
||||
ensure_model(&mut raw, args);
|
||||
// dt/s come from `--api`'s adapter (must match the input's wire shape, or tool/JSON
|
||||
// blocks get mis-segmented — see cost_model doc comment); the apply decision is
|
||||
// then re-derived under `--cost-model`'s pricing/CacheModel, since `report.applied`
|
||||
// reflects `--api`'s own strategy and the two can differ.
|
||||
let (_compressed, report) = compress(&raw, args, scorer.as_ref());
|
||||
let dt = report.removed_tokens_est;
|
||||
let s = report.rewrite_suffix_tokens;
|
||||
let applies = anyllm_optimize_core::should_apply(dt, s, args.horizon, &pricing, &model);
|
||||
// Only realize the delta when the gate actually applies: `applies == false` means
|
||||
// the original is forwarded untouched, so the realized cost change is $0, not the
|
||||
// negative "forced a no-op rewrite" value the raw formula would give for dt=0
|
||||
// (which `should_apply` exists precisely to avoid ever choosing).
|
||||
let net = if applies {
|
||||
net_cost_delta_usd(dt, s, args.horizon, &pricing, &model)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
total_dt += dt;
|
||||
total_net += net;
|
||||
println!(
|
||||
"{:<4} {:>9} {:>7} {:>7} {:>7} {:>14.6}",
|
||||
i, report.frontier, dt, s, applies, net
|
||||
);
|
||||
}
|
||||
let model_name = match model {
|
||||
CacheModel::ImplicitPrefix => "implicit-prefix",
|
||||
CacheModel::ExplicitBreakpoints => "explicit-breakpoints",
|
||||
};
|
||||
println!(
|
||||
"\nNET COST (horizon={}, {model_name}, input=${:.2}/Mtok cached_read=${:.2}/Mtok \
|
||||
write_mult={:.2}x): total ΔT={total_dt} tokens, net_delta=${total_net:.6}",
|
||||
args.horizon, pricing.input, pricing.cached_read, pricing.cache_write_mult,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_model(body: &mut Value, args: &Args) {
|
||||
pub(crate) fn ensure_model(body: &mut Value, args: &Args) {
|
||||
body["model"] = json!(args.model);
|
||||
if args.api == Api::Anthropic && body.get("max_tokens").is_none() {
|
||||
body["max_tokens"] = json!(args.max_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
/// M3.6: selects the counter `count_local` uses to report est_raw/est_comp. Defaults to
|
||||
/// `HeuristicBudgetCounter` (bytes/3.6, provider-agnostic). When this binary is built with
|
||||
/// `--features tiktoken` AND the target is OpenAI-shaped, uses the exact `o200k_base`
|
||||
/// tokenizer instead — Anthropic's tokenizer is unpublished, so the heuristic remains the
|
||||
/// only option there (ROADMAP risk 3). This only changes what the harness *reports*; the
|
||||
/// planning counter `compress()` passes into `optimize()` is unchanged.
|
||||
#[cfg_attr(not(feature = "tiktoken"), allow(unused_variables))]
|
||||
fn build_counter(args: &Args) -> Box<dyn BudgetCounter> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
if args.api == Api::Openai {
|
||||
return Box::new(budget_tiktoken::TiktokenBudgetCounter);
|
||||
}
|
||||
Box::new(HeuristicBudgetCounter::default())
|
||||
}
|
||||
|
||||
/// Local token estimate: sum the budget counter over all message text content.
|
||||
fn count_local(body: &Value, api: Api, counter: &dyn BudgetCounter) -> u64 {
|
||||
let mut total = 0u64;
|
||||
if api == Api::Anthropic {
|
||||
if let Some(sys) = body.get("system") {
|
||||
total += count_content(sys, counter);
|
||||
}
|
||||
}
|
||||
if let Some(msgs) = body.get("messages").and_then(|m| m.as_array()) {
|
||||
for m in msgs {
|
||||
if let Some(c) = m.get("content") {
|
||||
total += count_content(c, counter);
|
||||
}
|
||||
}
|
||||
}
|
||||
total
|
||||
}
|
||||
|
||||
fn count_content(c: &Value, counter: &dyn BudgetCounter) -> u64 {
|
||||
match c {
|
||||
Value::String(s) => counter.count(s),
|
||||
Value::Array(parts) => parts
|
||||
.iter()
|
||||
.map(|p| match p.get("text").and_then(|t| t.as_str()) {
|
||||
Some(t) => counter.count(t),
|
||||
None => counter.count(&p.to_string()),
|
||||
})
|
||||
.sum(),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send one request and return usage + assistant text.
|
||||
fn send_request(
|
||||
client: &reqwest::blocking::Client,
|
||||
args: &Args,
|
||||
body: &Value,
|
||||
base_url: &str,
|
||||
api: Api,
|
||||
) -> Result<Reply> {
|
||||
let mut body = body.clone();
|
||||
let url = match api {
|
||||
Api::Openai => format!("{}/chat/completions", base_url.trim_end_matches('/')),
|
||||
Api::Anthropic => format!("{}/messages", base_url.trim_end_matches('/')),
|
||||
};
|
||||
if args.stream {
|
||||
body["stream"] = json!(true);
|
||||
if api == Api::Openai {
|
||||
body["stream_options"] = json!({ "include_usage": true });
|
||||
}
|
||||
}
|
||||
|
||||
let key = resolve_key(args);
|
||||
let mut req = client.post(&url).json(&body);
|
||||
req = match api {
|
||||
Api::Openai => {
|
||||
if key.is_empty() {
|
||||
req
|
||||
} else {
|
||||
req.bearer_auth(key)
|
||||
}
|
||||
}
|
||||
Api::Anthropic => req
|
||||
.header("x-api-key", key)
|
||||
.header("anthropic-version", "2023-06-01"),
|
||||
};
|
||||
|
||||
let resp = req.send().context("request send failed")?;
|
||||
let status = resp.status();
|
||||
let text = resp.text().context("reading response body")?;
|
||||
if !status.is_success() {
|
||||
anyhow::bail!(
|
||||
"HTTP {status}: {}",
|
||||
text.chars().take(300).collect::<String>()
|
||||
);
|
||||
}
|
||||
Ok(Reply {
|
||||
prompt_tokens: extract_prompt_tokens(&text, api, args.stream)
|
||||
.context("could not find prompt/input token usage")?,
|
||||
text: extract_response_text(&text, api, args.stream),
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_key(args: &Args) -> String {
|
||||
args.api_key
|
||||
.clone()
|
||||
.or_else(|| std::env::var("OPENROUTER_API_KEY").ok())
|
||||
.or_else(|| std::env::var("ANYLLM_API_KEY").ok())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Pull prompt tokens from a JSON body or an SSE stream.
|
||||
fn extract_prompt_tokens(text: &str, api: Api, stream: bool) -> Result<u64> {
|
||||
let want = |v: &Value| -> Option<u64> {
|
||||
match api {
|
||||
Api::Openai => v.get("usage")?.get("prompt_tokens")?.as_u64(),
|
||||
Api::Anthropic => v
|
||||
.get("usage")
|
||||
.or_else(|| v.get("message").and_then(|m| m.get("usage")))?
|
||||
.get("input_tokens")?
|
||||
.as_u64(),
|
||||
}
|
||||
};
|
||||
if !stream {
|
||||
let v: Value = serde_json::from_str(text)?;
|
||||
return want(&v).context("usage field missing");
|
||||
}
|
||||
for payload in sse_payloads(text) {
|
||||
if let Ok(v) = serde_json::from_str::<Value>(&payload) {
|
||||
if let Some(n) = want(&v) {
|
||||
return Ok(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::bail!("no usage in stream")
|
||||
}
|
||||
|
||||
/// Pull assistant text from a JSON body or an SSE stream. Best-effort: returns "" if the
|
||||
/// shape is unrecognized (quality just reads as low, never panics).
|
||||
fn extract_response_text(text: &str, api: Api, stream: bool) -> String {
|
||||
if !stream {
|
||||
let Ok(v) = serde_json::from_str::<Value>(text) else {
|
||||
return String::new();
|
||||
};
|
||||
return match api {
|
||||
Api::Openai => v
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
Api::Anthropic => v
|
||||
.get("content")
|
||||
.and_then(|c| c.as_array())
|
||||
.map(|blocks| {
|
||||
blocks
|
||||
.iter()
|
||||
.filter_map(|b| b.get("text").and_then(|t| t.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
};
|
||||
}
|
||||
// Streaming: concatenate incremental text deltas.
|
||||
let mut out = String::new();
|
||||
for payload in sse_payloads(text) {
|
||||
let Ok(v) = serde_json::from_str::<Value>(&payload) else {
|
||||
continue;
|
||||
};
|
||||
match api {
|
||||
Api::Openai => {
|
||||
if let Some(s) = v
|
||||
.pointer("/choices/0/delta/content")
|
||||
.and_then(|c| c.as_str())
|
||||
{
|
||||
out.push_str(s);
|
||||
}
|
||||
}
|
||||
Api::Anthropic => {
|
||||
if let Some(s) = v.pointer("/delta/text").and_then(|c| c.as_str()) {
|
||||
out.push_str(s);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Yield the JSON payload of each `data:` SSE line, skipping `[DONE]`.
|
||||
fn sse_payloads(text: &str) -> impl Iterator<Item = String> + '_ {
|
||||
text.lines().filter_map(|line| {
|
||||
let payload = line.trim().strip_prefix("data:")?.trim();
|
||||
(payload != "[DONE]").then(|| payload.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
/// Word-level Jaccard similarity in [0,1]. 1.0 if both empty.
|
||||
fn jaccard(a: &str, b: &str) -> f64 {
|
||||
let sa: HashSet<&str> = a.split_whitespace().collect();
|
||||
let sb: HashSet<&str> = b.split_whitespace().collect();
|
||||
if sa.is_empty() && sb.is_empty() {
|
||||
return 1.0;
|
||||
}
|
||||
let inter = sa.intersection(&sb).count() as f64;
|
||||
let uni = sa.union(&sb).count() as f64;
|
||||
if uni == 0.0 {
|
||||
1.0
|
||||
} else {
|
||||
inter / uni
|
||||
}
|
||||
}
|
||||
|
||||
/// Ask an OpenAI-compatible judge model to score 1-5 how well `comp` preserves `raw`.
|
||||
fn run_judge(
|
||||
client: &reqwest::blocking::Client,
|
||||
args: &Args,
|
||||
raw: &str,
|
||||
comp: &str,
|
||||
) -> Result<u32> {
|
||||
let model = args.judge_model.as_ref().context("no judge model")?;
|
||||
let base = args.judge_base_url.as_deref().unwrap_or(&args.base_url);
|
||||
let prompt = format!(
|
||||
"You compare two AI assistant responses to the same user request. Response A is \
|
||||
from the full prompt; Response B is from a compressed prompt. Score 1-5 how well B \
|
||||
preserves A's meaning and quality (5 = equivalent, 1 = badly degraded). Reply with \
|
||||
ONLY the integer.\n\n[A]\n{raw}\n\n[B]\n{comp}"
|
||||
);
|
||||
let body = json!({
|
||||
"model": model,
|
||||
"messages": [{ "role": "user", "content": prompt }],
|
||||
"max_tokens": 4,
|
||||
"temperature": 0,
|
||||
});
|
||||
let url = format!("{}/chat/completions", base.trim_end_matches('/'));
|
||||
let key = resolve_key(args);
|
||||
let mut req = client.post(&url).json(&body);
|
||||
if !key.is_empty() {
|
||||
req = req.bearer_auth(key);
|
||||
}
|
||||
let resp = req.send().context("judge send failed")?;
|
||||
let text = resp.text().context("reading judge response")?;
|
||||
let v: Value = serde_json::from_str(&text).context("judge response not JSON")?;
|
||||
let content = v
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(|c| c.as_str())
|
||||
.context("judge response has no content")?;
|
||||
let score = content
|
||||
.chars()
|
||||
.find(|c| ('1'..='5').contains(c))
|
||||
.and_then(|c| c.to_digit(10))
|
||||
.context("no 1-5 score in judge reply")?;
|
||||
Ok(score)
|
||||
}
|
||||
|
||||
fn opt_u64(v: Option<u64>) -> String {
|
||||
v.map(|x| x.to_string()).unwrap_or_else(|| "-".into())
|
||||
}
|
||||
|
||||
fn pct(raw: u64, comp: u64) -> f64 {
|
||||
if raw == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(raw.saturating_sub(comp)) as f64 / raw as f64 * 100.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Load input bodies. Each JSONL line is either a full request body (`{"messages":...}`)
|
||||
/// or a bare JSON string (treated as one user message).
|
||||
fn load_inputs(args: &Args) -> Result<Vec<Value>> {
|
||||
pub(crate) fn load_inputs(args: &Args) -> Result<Vec<Value>> {
|
||||
if let Some(p) = &args.prompt {
|
||||
return Ok(vec![wrap_prompt(p)]);
|
||||
}
|
||||
@@ -721,54 +405,3 @@ fn load_inputs(args: &Args) -> Result<Vec<Value>> {
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn wrap_prompt(p: &str) -> Value {
|
||||
json!({ "messages": [ { "role": "user", "content": p } ] })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn jaccard_bounds() {
|
||||
assert_eq!(jaccard("", ""), 1.0);
|
||||
assert_eq!(jaccard("a b c", "a b c"), 1.0);
|
||||
assert!((jaccard("a b", "b c") - 1.0 / 3.0).abs() < 1e-9);
|
||||
assert_eq!(jaccard("a", "z"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_openai_nonstream() {
|
||||
let body =
|
||||
r#"{"choices":[{"message":{"content":"hello there"}}],"usage":{"prompt_tokens":42}}"#;
|
||||
assert_eq!(extract_prompt_tokens(body, Api::Openai, false).unwrap(), 42);
|
||||
assert_eq!(
|
||||
extract_response_text(body, Api::Openai, false),
|
||||
"hello there"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_anthropic_nonstream() {
|
||||
let body = r#"{"content":[{"type":"text","text":"hi"},{"type":"text","text":" world"}],"usage":{"input_tokens":7}}"#;
|
||||
assert_eq!(
|
||||
extract_prompt_tokens(body, Api::Anthropic, false).unwrap(),
|
||||
7
|
||||
);
|
||||
assert_eq!(
|
||||
extract_response_text(body, Api::Anthropic, false),
|
||||
"hi world"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_openai_stream_text_and_usage() {
|
||||
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"foo\"}}]}\n\
|
||||
data: {\"choices\":[{\"delta\":{\"content\":\"bar\"}}]}\n\
|
||||
data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11}}\n\
|
||||
data: [DONE]\n";
|
||||
assert_eq!(extract_prompt_tokens(sse, Api::Openai, true).unwrap(), 11);
|
||||
assert_eq!(extract_response_text(sse, Api::Openai, true), "foobar");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
use super::{client::*, utils::*, Api};
|
||||
|
||||
#[test]
|
||||
fn jaccard_bounds() {
|
||||
assert_eq!(jaccard("", ""), 1.0);
|
||||
assert_eq!(jaccard("a b c", "a b c"), 1.0);
|
||||
assert!((jaccard("a b", "b c") - 1.0 / 3.0).abs() < 1e-9);
|
||||
assert_eq!(jaccard("a", "z"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_openai_nonstream() {
|
||||
let body =
|
||||
r#"{"choices":[{"message":{"content":"hello there"}}],"usage":{"prompt_tokens":42}}"#;
|
||||
assert_eq!(extract_prompt_tokens(body, Api::Openai, false).unwrap(), 42);
|
||||
assert_eq!(
|
||||
extract_response_text(body, Api::Openai, false),
|
||||
"hello there"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_anthropic_nonstream() {
|
||||
let body = r#"{"content":[{"type":"text","text":"hi"},{"type":"text","text":" world"}],"usage":{"input_tokens":7}}"#;
|
||||
assert_eq!(
|
||||
extract_prompt_tokens(body, Api::Anthropic, false).unwrap(),
|
||||
7
|
||||
);
|
||||
assert_eq!(
|
||||
extract_response_text(body, Api::Anthropic, false),
|
||||
"hi world"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_openai_stream_text_and_usage() {
|
||||
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"foo\"}}]}\n\
|
||||
data: {\"choices\":[{\"delta\":{\"content\":\"bar\"}}]}\n\
|
||||
data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11}}\n\
|
||||
data: [DONE]\n";
|
||||
assert_eq!(extract_prompt_tokens(sse, Api::Openai, true).unwrap(), 11);
|
||||
assert_eq!(extract_response_text(sse, Api::Openai, true), "foobar");
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
use super::{Api, Args};
|
||||
use anyllm_optimize_core::{BudgetCounter, HeuristicBudgetCounter};
|
||||
use serde_json::Value;
|
||||
|
||||
/// M3.6: selects the counter `count_local` uses to report est_raw/est_comp. Defaults to
|
||||
/// `HeuristicBudgetCounter` (bytes/3.6, provider-agnostic). When this binary is built with
|
||||
/// `--features tiktoken` AND the target is OpenAI-shaped, uses the exact `o200k_base`
|
||||
/// tokenizer instead — Anthropic's tokenizer is unpublished, so the heuristic remains the
|
||||
/// only option there (ROADMAP risk 3). This only changes what the harness *reports*; the
|
||||
/// planning counter `compress()` passes into `optimize()` is unchanged.
|
||||
#[cfg_attr(not(feature = "tiktoken"), allow(unused_variables))]
|
||||
pub fn build_counter(args: &Args) -> Box<dyn BudgetCounter> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
if args.api == Api::Openai {
|
||||
return Box::new(super::budget_tiktoken::TiktokenBudgetCounter);
|
||||
}
|
||||
Box::new(HeuristicBudgetCounter::default())
|
||||
}
|
||||
|
||||
/// Local token estimate: sum the budget counter over all message text content.
|
||||
pub fn count_local(body: &Value, api: Api, counter: &dyn BudgetCounter) -> u64 {
|
||||
let mut total = 0u64;
|
||||
if api == Api::Anthropic {
|
||||
if let Some(sys) = body.get("system") {
|
||||
total += count_content(sys, counter);
|
||||
}
|
||||
}
|
||||
if let Some(msgs) = body.get("messages").and_then(|m| m.as_array()) {
|
||||
for m in msgs {
|
||||
if let Some(c) = m.get("content") {
|
||||
total += count_content(c, counter);
|
||||
}
|
||||
}
|
||||
}
|
||||
total
|
||||
}
|
||||
|
||||
pub fn count_content(c: &Value, counter: &dyn BudgetCounter) -> u64 {
|
||||
match c {
|
||||
Value::String(s) => counter.count(s),
|
||||
Value::Array(parts) => parts
|
||||
.iter()
|
||||
.map(|p| match p.get("text").and_then(|t| t.as_str()) {
|
||||
Some(t) => counter.count(t),
|
||||
None => counter.count(&p.to_string()),
|
||||
})
|
||||
.sum(),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashSet;
|
||||
|
||||
/// Word-level Jaccard similarity in [0,1]. 1.0 if both empty.
|
||||
pub fn jaccard(a: &str, b: &str) -> f64 {
|
||||
let sa: HashSet<&str> = a.split_whitespace().collect();
|
||||
let sb: HashSet<&str> = b.split_whitespace().collect();
|
||||
if sa.is_empty() && sb.is_empty() {
|
||||
return 1.0;
|
||||
}
|
||||
let inter = sa.intersection(&sb).count() as f64;
|
||||
let uni = sa.union(&sb).count() as f64;
|
||||
if uni == 0.0 {
|
||||
1.0
|
||||
} else {
|
||||
inter / uni
|
||||
}
|
||||
}
|
||||
|
||||
pub fn wrap_prompt(p: &str) -> Value {
|
||||
json!({ "messages": [ { "role": "user", "content": p } ] })
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
# anyllm_providers
|
||||
|
||||
Provider and model catalog for the [anyllm-proxy](https://github.com/whit3rabbit/llm-translate-api) workspace.
|
||||
Provider and model catalog.
|
||||
|
||||
## What this crate is
|
||||
|
||||
|
||||
+8
-8
File diff suppressed because one or more lines are too long
@@ -1,758 +1,17 @@
|
||||
// Mirrors the JSON shapes returned by /admin/api/* endpoints.
|
||||
// Keep in sync with Rust structs in crates/proxy/src/admin/state.rs and routes/.
|
||||
|
||||
/** Represents the status and configuration details of the proxy server. */
|
||||
export interface ProxyStatus {
|
||||
/** True if the proxy backend has been configured. */
|
||||
configured: boolean
|
||||
/** The port number the proxy listens on for client requests. */
|
||||
proxy_port: number
|
||||
/** Whether the proxy is currently running and accepting connections. */
|
||||
proxy_running: boolean
|
||||
}
|
||||
|
||||
/** Represents real-time usage and performance metrics for the proxy. */
|
||||
export interface Metrics {
|
||||
/** The total number of requests received by the proxy. */
|
||||
total_requests: number
|
||||
/** The number of successful requests. */
|
||||
successful_requests: number
|
||||
/** The number of failed requests. */
|
||||
failed_requests: number
|
||||
/** The current request rate per minute. */
|
||||
requests_per_minute: number
|
||||
/** The p50 latency in milliseconds, if available. */
|
||||
p50_latency_ms: number | null
|
||||
/** The p95 latency in milliseconds, if available. */
|
||||
p95_latency_ms: number | null
|
||||
/** The percentage rate of request failures. */
|
||||
error_rate: number
|
||||
/** The number of streaming connection requests started. */
|
||||
streams_started: number
|
||||
/** The number of streaming connections completed. */
|
||||
streams_completed: number
|
||||
/** The number of streaming connections that failed. */
|
||||
streams_failed: number
|
||||
/** The number of streaming connections disconnected by the client. */
|
||||
streams_client_disconnected: number
|
||||
/** Requests where pxpipe text-to-image compression fired. */
|
||||
pxpipe_compressed_total: number
|
||||
/** Total PNG image blocks pxpipe emitted. */
|
||||
pxpipe_images_total: number
|
||||
/** Total source chars pxpipe replaced with images. */
|
||||
pxpipe_imaged_chars_total: number
|
||||
/** Requests where RTK tool-output compression fired. */
|
||||
rtk_compressed_total: number
|
||||
/** Total tool-result payloads RTK rewrote. */
|
||||
rtk_blocks_total: number
|
||||
/** Total source chars RTK removed from tool output. */
|
||||
rtk_saved_chars_total: number
|
||||
}
|
||||
|
||||
/** A single request transaction log entry. */
|
||||
export interface RequestLogEntry {
|
||||
/** Unique transaction ID. */
|
||||
request_id: string
|
||||
/** ISO 8601 timestamp when the request was made. */
|
||||
timestamp: string
|
||||
/** The backend model provider targeted. */
|
||||
backend: string
|
||||
/** The model identifier requested by the client. */
|
||||
model_requested: string | null
|
||||
/** The actual model identifier routed to on the backend. */
|
||||
model_mapped: string | null
|
||||
/** HTTP response status code. */
|
||||
status_code: number
|
||||
/** Request duration in milliseconds. */
|
||||
latency_ms: number
|
||||
/** Number of input tokens processed. */
|
||||
input_tokens: number | null
|
||||
/** Number of output tokens generated. */
|
||||
output_tokens: number | null
|
||||
/** Whether the request was streamed. */
|
||||
is_streaming: boolean
|
||||
/** The raw error message returned by the backend, if any. */
|
||||
error_message: string | null
|
||||
/** Normalized error class/category. */
|
||||
error_kind: string | null
|
||||
/** The virtual key ID used for authorization, if any. */
|
||||
key_id: number | null
|
||||
/** Calculated transaction cost in USD. */
|
||||
cost_usd: number | null
|
||||
}
|
||||
|
||||
/** Paginated list response for request log queries. */
|
||||
export interface RequestsResponse {
|
||||
/** The page of request log entries. */
|
||||
requests: RequestLogEntry[]
|
||||
/** Maximum number of items returned. */
|
||||
limit: number
|
||||
/** Pagination offset. */
|
||||
offset: number
|
||||
/** True if more records are available. */
|
||||
has_more: boolean
|
||||
}
|
||||
|
||||
/** Represents a virtual API key configuration and its associated limits. */
|
||||
export interface VirtualKey {
|
||||
/** Unique database ID. */
|
||||
id: number
|
||||
/** The prefix of the key shown to users. */
|
||||
key_prefix: string
|
||||
/** Optional descriptive note. */
|
||||
description: string | null
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 expiration timestamp, if set. */
|
||||
expires_at: string | null
|
||||
/** ISO 8601 revocation timestamp, if revoked. */
|
||||
revoked_at: string | null
|
||||
/** Spend limit in USD. */
|
||||
spend_limit: number | null
|
||||
/** Monthly budget limit in USD. */
|
||||
max_budget_usd: number | null
|
||||
/** Duration of the budget period (e.g. 'monthly'). */
|
||||
budget_duration: string | null
|
||||
/** Requests-per-minute limit. */
|
||||
rpm_limit: number | null
|
||||
/** Tokens-per-minute limit. */
|
||||
tpm_limit: number | null
|
||||
/** Total spend in USD across the key lifetime. */
|
||||
total_spend: number
|
||||
/** Total count of requests made with this key. */
|
||||
total_requests: number
|
||||
/** Total count of tokens processed. */
|
||||
total_tokens: number
|
||||
/** ISO 8601 reset timestamp for the current budget period. */
|
||||
period_reset_at: string | null
|
||||
/** List of model names this key is restricted to, if any. */
|
||||
allowed_models: string[] | null
|
||||
/** List of route names this key is restricted to, if any. */
|
||||
allowed_routes: string[] | null
|
||||
/** Active status of the key. */
|
||||
status: 'active' | 'revoked' | 'expired' | 'override'
|
||||
/** Spend in USD during the current budget period. */
|
||||
period_spend_usd: number
|
||||
}
|
||||
|
||||
/** Spent token and requests details for a virtual key. */
|
||||
export interface KeySpend {
|
||||
/** Unique key database ID. */
|
||||
id: number
|
||||
/** Total spend in USD. */
|
||||
total_spend: number
|
||||
/** Total request count. */
|
||||
total_requests: number
|
||||
/** Total token count. */
|
||||
total_tokens: number
|
||||
}
|
||||
|
||||
/** Represents a backend endpoint model and its metrics. */
|
||||
export interface Backend {
|
||||
/** Name of the backend. */
|
||||
name: string
|
||||
/** Model mapped for heavy workloads. */
|
||||
big_model: string
|
||||
/** Model mapped for light workloads. */
|
||||
small_model: string
|
||||
/** Request outcome counters. */
|
||||
metrics: {
|
||||
requests_total: number
|
||||
requests_success: number
|
||||
requests_error: number
|
||||
}
|
||||
}
|
||||
|
||||
/** Represents a single configuration entry key-value pair. */
|
||||
export interface ConfigEntry {
|
||||
/** Configuration key. */
|
||||
key: string
|
||||
/** Configuration value. */
|
||||
value: string
|
||||
/** Timestamp when updated. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Represents the full system configuration and environment variables. */
|
||||
export interface ConfigResponse {
|
||||
/** List of config database overrides. */
|
||||
entries: ConfigEntry[]
|
||||
/** Server environment variables map. */
|
||||
env: Record<string, string>
|
||||
/** System logging level. */
|
||||
log_level: string
|
||||
/** Whether requests and responses are logged. */
|
||||
log_bodies: boolean
|
||||
/** Whether credentials are redacted from logs. */
|
||||
redact_secrets: boolean
|
||||
/** Whether to attempt to repair broken thinking blocks in Anthropic streams. */
|
||||
anthropic_thinking_repair: boolean
|
||||
/** True if pxpipe image compression is enabled. */
|
||||
pxpipe_compress: boolean
|
||||
/** CSV of model bases in pxpipe compression scope. */
|
||||
pxpipe_models: string
|
||||
/** Vision-capable Claude models offered as per-model scope toggles. */
|
||||
pxpipe_available_models: string[]
|
||||
/** Whether RTK tool-output compression is active. */
|
||||
rtk_compress: boolean
|
||||
/** CSV of model bases in RTK compression scope (empty = all models). */
|
||||
rtk_models: string
|
||||
/** Whether to pass client authorization headers to backend. */
|
||||
forward_client_auth: boolean
|
||||
/** Tool guardrail mode configured. */
|
||||
tool_guardrail_mode: string
|
||||
/** Prompt-compression optimizer mode: 'off' | 'shadow' | 'live'. */
|
||||
optimizer_mode: string
|
||||
/** Mappings of backend names to their configured big/small models. */
|
||||
backends: Record<string, { big_model: string; small_model: string }>
|
||||
/** List of keys whose overrides are active. */
|
||||
overridden_keys: string[]
|
||||
}
|
||||
|
||||
/** Status of the optional LLMLingua-2 ONNX model artifact (optimizer scorer tier). */
|
||||
export interface OptimizerModelStatus {
|
||||
/** Proxy built with the `optimizer-onnx` feature. When false the tier is inert. */
|
||||
compiled_in: boolean
|
||||
/** Verified model artifact is present on disk. */
|
||||
present: boolean
|
||||
/** A download+verify is currently in flight. */
|
||||
downloading: boolean
|
||||
/** Last download error, if any. */
|
||||
error: string | null
|
||||
/** Pinned sha256 the artifact is verified against. */
|
||||
sha256: string
|
||||
/** Expected download size in bytes. */
|
||||
size_bytes: number
|
||||
}
|
||||
|
||||
/** A time-series data point for observability charts. */
|
||||
export interface ObservabilityPoint {
|
||||
/** UNIX timestamp representing the bucket start time. */
|
||||
bucket_start: number
|
||||
/** Total requests in this bucket. */
|
||||
requests: number
|
||||
/** Total errors in this bucket. */
|
||||
errors: number
|
||||
/** Total input tokens processed in this bucket. */
|
||||
input_tokens: number
|
||||
/** Total output tokens generated in this bucket. */
|
||||
output_tokens: number
|
||||
/** Calculated cost in USD in this bucket. */
|
||||
cost_usd: number
|
||||
}
|
||||
|
||||
/** Summary of error occurrences on a backend. */
|
||||
export interface ObservabilityFailure {
|
||||
/** Class or classification of error. */
|
||||
error_kind: string
|
||||
/** Occurrence count. */
|
||||
count: number
|
||||
/** Last occurrence timestamp. */
|
||||
last_seen: string
|
||||
/** Last error message received. */
|
||||
last_message: string
|
||||
}
|
||||
|
||||
/** Represents a single trace event in the observability timeline. */
|
||||
export interface ObservabilityTimeline {
|
||||
/** Transaction request ID. */
|
||||
request_id: string
|
||||
/** ISO 8601 request timestamp. */
|
||||
timestamp: string
|
||||
/** Backend provider routed to. */
|
||||
backend: string
|
||||
/** Model mapped to. */
|
||||
model: string
|
||||
/** Latency in milliseconds. */
|
||||
latency_ms: number
|
||||
/** Request outcome status. */
|
||||
status: string
|
||||
}
|
||||
|
||||
/** Observability stats summary response. */
|
||||
export interface ObservabilityResponse {
|
||||
/** The time window in hours for metrics. */
|
||||
window_hours: number
|
||||
/** The name of the backend. */
|
||||
backend: string
|
||||
/** Total requests within the window. */
|
||||
total_requests: number
|
||||
/** Total errors within the window. */
|
||||
total_errors: number
|
||||
/** Total input tokens processed. */
|
||||
total_input_tokens: number
|
||||
/** Total output tokens generated. */
|
||||
total_output_tokens: number
|
||||
/** Total cost in USD. */
|
||||
total_cost_usd: number
|
||||
/** Historical metrics series points. */
|
||||
series: ObservabilityPoint[]
|
||||
/** Error breakdown summary. */
|
||||
failures: ObservabilityFailure[]
|
||||
/** Timeline events. */
|
||||
timeline: ObservabilityTimeline[]
|
||||
}
|
||||
|
||||
/** Represents a single model configuration entry. */
|
||||
export interface ModelEntry {
|
||||
/** The unique name/identifier of the model. */
|
||||
model_name: string
|
||||
/** The number of active deployments for this model. */
|
||||
deployments: number
|
||||
}
|
||||
|
||||
/** Response shape for models list requests. */
|
||||
export interface ModelsResponse {
|
||||
/** List of model configurations. */
|
||||
models: ModelEntry[]
|
||||
/** Router strategy (e.g. priority, failover). */
|
||||
strategy: string | null
|
||||
/** Optional descriptive note. */
|
||||
note?: string
|
||||
}
|
||||
|
||||
/** Represents a model discovered from a backend API. */
|
||||
export interface DiscoveredModel {
|
||||
/** The model identifier. */
|
||||
id: string
|
||||
/** Optional human-readable name. */
|
||||
name: string | null
|
||||
}
|
||||
|
||||
/** Response containing discovered models. */
|
||||
export interface DiscoverResponse {
|
||||
/** List of discovered models. */
|
||||
models: DiscoveredModel[]
|
||||
/** Source or method used for discovery. */
|
||||
source: string
|
||||
/** True if authorization was used. */
|
||||
auth_used: boolean
|
||||
}
|
||||
|
||||
/** A single audit log entry tracking administrative actions. */
|
||||
export interface AuditEntry {
|
||||
/** Unique database ID. */
|
||||
id: number
|
||||
/** ISO 8601 action timestamp. */
|
||||
timestamp: string
|
||||
/** The action performed. */
|
||||
action: string
|
||||
/** Type of target resource affected. */
|
||||
target_type: string
|
||||
/** ID of target resource affected. */
|
||||
target_id: string | null
|
||||
/** Detailed changes/payload. */
|
||||
detail: string | null
|
||||
/** Source IP of the requester. */
|
||||
source_ip: string | null
|
||||
}
|
||||
|
||||
/** Paginated list response for audit log queries. */
|
||||
export interface AuditResponse {
|
||||
/** List of audit log entries. */
|
||||
entries: AuditEntry[]
|
||||
/** Maximum number of items returned. */
|
||||
limit: number
|
||||
/** Pagination offset. */
|
||||
offset: number
|
||||
/** True if more records are available. */
|
||||
has_more: boolean
|
||||
}
|
||||
|
||||
// --- Traffic tab (new) ---
|
||||
|
||||
/** Real-time metrics for a specific API route. */
|
||||
export interface RouteMetrics {
|
||||
/** The request path/route. */
|
||||
path: string
|
||||
/** Number of requests per minute. */
|
||||
requests_per_min: number
|
||||
/** Percentage rate of failures. */
|
||||
error_rate: number
|
||||
/** Average latency in milliseconds. */
|
||||
avg_latency_ms: number
|
||||
/** The p95 latency in milliseconds. */
|
||||
p95_latency_ms: number
|
||||
/** Total number of requests. */
|
||||
total_requests: number
|
||||
}
|
||||
|
||||
/** Time-series requests count point for a route. */
|
||||
export interface TrafficSeriesPoint {
|
||||
/** UNIX timestamp representing the bucket start time. */
|
||||
bucket_start: number
|
||||
/** API path. */
|
||||
path: string
|
||||
/** Request count. */
|
||||
requests: number
|
||||
}
|
||||
|
||||
/** Response containing traffic analytics. */
|
||||
export interface TrafficResponse {
|
||||
/** Time window in hours. */
|
||||
window_hours: number
|
||||
/** Metrics per route. */
|
||||
routes: RouteMetrics[]
|
||||
/** Time-series data points. */
|
||||
series: TrafficSeriesPoint[]
|
||||
}
|
||||
|
||||
// --- Uptime tab (new) ---
|
||||
|
||||
/** A single day's uptime status. */
|
||||
export interface HistoryDay {
|
||||
/** Date in YYYY-MM-DD format. */
|
||||
date: string
|
||||
/** Availability status. */
|
||||
status: 'up' | 'down' | 'degraded' | 'no-data'
|
||||
}
|
||||
|
||||
/** Uptime info for the proxy itself. */
|
||||
export interface ProxyUptimeInfo {
|
||||
/** UNIX timestamp when the proxy started. */
|
||||
started_at: number
|
||||
/** Uptime percentage over the last 30 days. */
|
||||
uptime_pct_30d: number
|
||||
/** Daily uptime history. */
|
||||
history: HistoryDay[]
|
||||
}
|
||||
|
||||
/** Uptime info for a backend endpoint. */
|
||||
export interface BackendUptimeInfo {
|
||||
/** Name of the backend. */
|
||||
name: string
|
||||
/** Current connection status. */
|
||||
status: 'up' | 'down' | 'unknown'
|
||||
/** UNIX timestamp of the last health check. */
|
||||
last_checked_at: number | null
|
||||
/** Last checked latency in milliseconds. */
|
||||
last_latency_ms: number | null
|
||||
/** 30-day uptime percentage. */
|
||||
uptime_pct_30d: number
|
||||
/** Daily uptime history. */
|
||||
history: HistoryDay[]
|
||||
}
|
||||
|
||||
/** Full uptime summary response. */
|
||||
export interface UptimeResponse {
|
||||
/** Proxy server uptime details. */
|
||||
proxy: ProxyUptimeInfo
|
||||
/** Backend endpoints uptime details. */
|
||||
backends: BackendUptimeInfo[]
|
||||
}
|
||||
|
||||
// --- Env file import / export ---
|
||||
|
||||
/** Warning message generated during env import. */
|
||||
export interface EnvWarning {
|
||||
/** Affected line number. */
|
||||
line: number | null
|
||||
/** Affected environment variable key. */
|
||||
key: string | null
|
||||
/** Warning message. */
|
||||
message: string
|
||||
}
|
||||
|
||||
/** Response detailing the result of importing an env file. */
|
||||
export interface EnvImportResponse {
|
||||
/** Count of imported variables applied. */
|
||||
applied: number
|
||||
/** Non-fatal warnings generated. */
|
||||
warnings: EnvWarning[]
|
||||
}
|
||||
|
||||
/** Error payload for env file imports. */
|
||||
export interface EnvImportError {
|
||||
/** Hard/blocking validation errors. */
|
||||
hard_errors: string[]
|
||||
/** Non-fatal warnings. */
|
||||
warnings: EnvWarning[]
|
||||
}
|
||||
|
||||
// --- Provider catalog ---
|
||||
|
||||
/** Details of a provider available in the LiteLLM catalog. */
|
||||
export interface CatalogProvider {
|
||||
/** Unique provider identifier. */
|
||||
id: string
|
||||
/** User-facing display name. */
|
||||
display_name: string
|
||||
/** Keep for backwards compatibility. */
|
||||
name?: string
|
||||
/** Communication protocol. */
|
||||
protocol: string
|
||||
/** Auth mechanism. */
|
||||
auth: string
|
||||
/** Proxy integration status. */
|
||||
status: 'implemented' | 'wired' | 'stub'
|
||||
/** Default base URL. */
|
||||
default_base_url: string
|
||||
/** Expected environment variables for API keys. */
|
||||
env_vars: string[]
|
||||
/** LiteLLM prefix string. */
|
||||
litellm_prefix: string
|
||||
/** Supported capabilities. */
|
||||
capabilities: {
|
||||
chat_completions: boolean
|
||||
streaming: boolean
|
||||
tool_use: boolean
|
||||
embeddings: boolean
|
||||
vision: boolean
|
||||
batch: boolean
|
||||
}
|
||||
/** Total count of models. */
|
||||
model_count: number
|
||||
/** Count of models currently cached. */
|
||||
cached_model_count: number
|
||||
/** UNIX timestamp of the last cache refresh. */
|
||||
last_refreshed: number | null
|
||||
}
|
||||
|
||||
// --- Managed backends ---
|
||||
|
||||
/** Represents a backend credentials deployment managed by the admin. */
|
||||
export interface ManagedBackend {
|
||||
/** Unique database ID. */
|
||||
id: string
|
||||
/** Unique backend name. */
|
||||
name: string
|
||||
/** Provider ID. */
|
||||
provider_id: string
|
||||
/** True if the API key is configured. */
|
||||
api_key_set: boolean
|
||||
/** True if AWS credentials are set (for AWS Bedrock, etc.). */
|
||||
aws_creds_set: boolean
|
||||
/** Base URL of the API endpoint. */
|
||||
api_base: string | null
|
||||
/** Optional deployment name (e.g. Azure deployment name). */
|
||||
deployment: string | null
|
||||
/** Optional API version. */
|
||||
api_version: string | null
|
||||
/** Optional cloud project ID. */
|
||||
project: string | null
|
||||
/** Optional cloud region (e.g. AWS/Azure region). */
|
||||
region: string | null
|
||||
/** Rate limit: requests per minute override. */
|
||||
rpm: number | null
|
||||
/** Rate limit: tokens per minute override. */
|
||||
tpm: number | null
|
||||
/** True if this backend is enabled. */
|
||||
enabled: boolean
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 last update timestamp. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Response containing managed backends. */
|
||||
export interface ManagedBackendsResponse {
|
||||
/** List of managed backends. */
|
||||
backends: ManagedBackend[]
|
||||
}
|
||||
|
||||
/** Request payload for creating a new managed backend. */
|
||||
export interface CreateManagedBackendRequest {
|
||||
/** Unique name for the backend. */
|
||||
name: string
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** API key credential string. */
|
||||
api_key?: string
|
||||
/** API base URL. */
|
||||
api_base?: string
|
||||
/** API deployment name. */
|
||||
deployment?: string
|
||||
/** API version. */
|
||||
api_version?: string
|
||||
/** Cloud project ID. */
|
||||
project?: string
|
||||
/** Cloud region. */
|
||||
region?: string
|
||||
/** AWS access key ID. */
|
||||
aws_access_key_id?: string
|
||||
/** AWS secret access key. */
|
||||
aws_secret_access_key?: string
|
||||
/** AWS session token. */
|
||||
aws_session_token?: string
|
||||
/** Requests per minute limit. */
|
||||
rpm?: number
|
||||
/** Tokens per minute limit. */
|
||||
tpm?: number
|
||||
/** Active status toggle. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload for updating an existing managed backend. */
|
||||
export type UpdateManagedBackendRequest = Partial<Omit<CreateManagedBackendRequest, 'name'>>
|
||||
|
||||
// --- Routes ---
|
||||
|
||||
/** Represents a configured API routing definition. */
|
||||
export interface Route {
|
||||
/** Unique route database ID. */
|
||||
id: string
|
||||
/** Route name. */
|
||||
name: string
|
||||
/** Optional description. */
|
||||
description: string | null
|
||||
/** Load balancing strategy. */
|
||||
strategy: string
|
||||
/** Requests per minute limit. */
|
||||
rpm: number | null
|
||||
/** Tokens per minute limit. */
|
||||
tpm: number | null
|
||||
/** Cost budget in USD. */
|
||||
budget_usd: number | null
|
||||
/** Active status toggle. */
|
||||
enabled: boolean
|
||||
/** Tool guardrail mode override. */
|
||||
guardrail_mode: string | null
|
||||
/** pxpipe image compression toggle override. */
|
||||
pxpipe_compress: boolean | null
|
||||
/** pxpipe models CSV override. */
|
||||
pxpipe_models: string | null
|
||||
/** Secret redaction toggle override. */
|
||||
redact_secrets: boolean | null
|
||||
/** Position order in the router. */
|
||||
position: number
|
||||
/** Count of providers assigned to this route. */
|
||||
provider_count: number
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 last update timestamp. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Response containing routing configs. */
|
||||
export interface RoutesResponse {
|
||||
/** List of routes. */
|
||||
routes: Route[]
|
||||
}
|
||||
|
||||
/** Request payload to create a new route. */
|
||||
export interface CreateRouteRequest {
|
||||
/** Unique route name. */
|
||||
name: string
|
||||
/** Route description. */
|
||||
description?: string
|
||||
/** Load balancing strategy. */
|
||||
strategy?: string
|
||||
/** Requests per minute limit. */
|
||||
rpm?: number
|
||||
/** Tokens per minute limit. */
|
||||
tpm?: number
|
||||
/** Cost budget in USD. */
|
||||
budget_usd?: number
|
||||
/** Active status toggle. */
|
||||
enabled?: boolean
|
||||
/** Tool guardrail mode override. */
|
||||
guardrail_mode?: string | null
|
||||
/** pxpipe image compression toggle override. */
|
||||
pxpipe_compress?: boolean | null
|
||||
/** pxpipe models CSV override. */
|
||||
pxpipe_models?: string | null
|
||||
/** Secret redaction toggle override. */
|
||||
redact_secrets?: boolean | null
|
||||
/** Position order. */
|
||||
position?: number
|
||||
}
|
||||
|
||||
/** Request payload to update an existing route. */
|
||||
export type UpdateRouteRequest = Partial<CreateRouteRequest>
|
||||
|
||||
/** Represents a provider mapped to a route. */
|
||||
export interface RouteProvider {
|
||||
/** Unique assignment ID. */
|
||||
id: string
|
||||
/** Target route ID. */
|
||||
route_id: string
|
||||
/** Backend ID. */
|
||||
backend_id: string
|
||||
/** Name of the backend. */
|
||||
backend_name: string
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** Supported models list. */
|
||||
models: string[]
|
||||
/** Evaluation priority. */
|
||||
priority: number
|
||||
/** Active status. */
|
||||
enabled: boolean
|
||||
}
|
||||
|
||||
/** Response containing route provider mappings. */
|
||||
export interface RouteProvidersResponse {
|
||||
/** List of mappings. */
|
||||
providers: RouteProvider[]
|
||||
}
|
||||
|
||||
/** Request payload to add a provider assignment to a route. */
|
||||
export interface AddRouteProviderRequest {
|
||||
/** Target backend ID. */
|
||||
backend_id: string
|
||||
/** Supported models list. */
|
||||
models?: string[]
|
||||
/** Priority position. */
|
||||
priority?: number
|
||||
/** Active status. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload to update a provider assignment configuration. */
|
||||
export interface UpdateRouteProviderRequest {
|
||||
/** Supported models list. */
|
||||
models?: string[]
|
||||
/** Priority position. */
|
||||
priority?: number
|
||||
/** Active status. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload to reorder provider mappings. */
|
||||
export interface ReorderRouteProvidersRequest {
|
||||
/** Ordered list of provider assignment IDs. */
|
||||
provider_ids: string[]
|
||||
}
|
||||
|
||||
// --- WebSocket events ---
|
||||
|
||||
/** Types of events sent over the admin WebSocket channel. */
|
||||
export type WSEvent =
|
||||
| { type: 'request_completed'; data: RequestLogEntry }
|
||||
| { type: 'metrics_snapshot'; data: Metrics }
|
||||
| { type: 'config_changed'; data: { key: string; value: string } }
|
||||
| { type: 'backend_health_changed'; data: { backend: string; status: 'up' | 'down'; latency_ms: number | null } }
|
||||
|
||||
/** Details of a model available in the LiteLLM catalog. */
|
||||
export interface CatalogModel {
|
||||
/** Unique model identifier. */
|
||||
id: string
|
||||
/** Maximum context window size in tokens. */
|
||||
context_window: number
|
||||
/** Maximum output tokens. */
|
||||
max_output_tokens: number
|
||||
/** Availability status of the model. */
|
||||
status: 'available' | 'deprecated' | 'stub'
|
||||
/** Capabilities flags. */
|
||||
capabilities: {
|
||||
streaming: boolean
|
||||
tool_use: boolean
|
||||
vision: boolean
|
||||
extended_thinking: boolean
|
||||
}
|
||||
/** Pricing per million tokens. */
|
||||
pricing: {
|
||||
input_per_million_tokens: number
|
||||
output_per_million_tokens: number
|
||||
} | null
|
||||
}
|
||||
|
||||
/** Response containing model lists. */
|
||||
export interface CatalogModelsResponse {
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** True if this provider has models. */
|
||||
has_models: boolean
|
||||
/** List of models. */
|
||||
models: CatalogModel[]
|
||||
}
|
||||
export * from './types/proxy'
|
||||
export * from './types/metrics'
|
||||
export * from './types/requests'
|
||||
export * from './types/keys'
|
||||
export * from './types/backends'
|
||||
export * from './types/config'
|
||||
export * from './types/optimizer'
|
||||
export * from './types/models'
|
||||
export * from './types/audit'
|
||||
export * from './types/uptime'
|
||||
export * from './types/env'
|
||||
export * from './types/catalog'
|
||||
export * from './types/routes'
|
||||
export * from './types/ws'
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
/** A single audit log entry tracking administrative actions. */
|
||||
export interface AuditEntry {
|
||||
/** Unique database ID. */
|
||||
id: number
|
||||
/** ISO 8601 action timestamp. */
|
||||
timestamp: string
|
||||
/** The action performed. */
|
||||
action: string
|
||||
/** Type of target resource affected. */
|
||||
target_type: string
|
||||
/** ID of target resource affected. */
|
||||
target_id: string | null
|
||||
/** Detailed changes/payload. */
|
||||
detail: string | null
|
||||
/** Source IP of the requester. */
|
||||
source_ip: string | null
|
||||
}
|
||||
|
||||
/** Paginated list response for audit log queries. */
|
||||
export interface AuditResponse {
|
||||
/** List of audit log entries. */
|
||||
entries: AuditEntry[]
|
||||
/** Maximum number of items returned. */
|
||||
limit: number
|
||||
/** Pagination offset. */
|
||||
offset: number
|
||||
/** True if more records are available. */
|
||||
has_more: boolean
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
/** Represents a backend endpoint model and its metrics. */
|
||||
export interface Backend {
|
||||
/** Name of the backend. */
|
||||
name: string
|
||||
/** Model mapped for heavy workloads. */
|
||||
big_model: string
|
||||
/** Model mapped for light workloads. */
|
||||
small_model: string
|
||||
/** Request outcome counters. */
|
||||
metrics: {
|
||||
requests_total: number
|
||||
requests_success: number
|
||||
requests_error: number
|
||||
}
|
||||
}
|
||||
|
||||
/** Represents a backend credentials deployment managed by the admin. */
|
||||
export interface ManagedBackend {
|
||||
/** Unique database ID. */
|
||||
id: string
|
||||
/** Unique backend name. */
|
||||
name: string
|
||||
/** Provider ID. */
|
||||
provider_id: string
|
||||
/** True if the API key is configured. */
|
||||
api_key_set: boolean
|
||||
/** True if AWS credentials are set (for AWS Bedrock, etc.). */
|
||||
aws_creds_set: boolean
|
||||
/** Base URL of the API endpoint. */
|
||||
api_base: string | null
|
||||
/** Optional deployment name (e.g. Azure deployment name). */
|
||||
deployment: string | null
|
||||
/** Optional API version. */
|
||||
api_version: string | null
|
||||
/** Optional cloud project ID. */
|
||||
project: string | null
|
||||
/** Optional cloud region (e.g. AWS/Azure region). */
|
||||
region: string | null
|
||||
/** Rate limit: requests per minute override. */
|
||||
rpm: number | null
|
||||
/** Rate limit: tokens per minute override. */
|
||||
tpm: number | null
|
||||
/** True if this backend is enabled. */
|
||||
enabled: boolean
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 last update timestamp. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Response containing managed backends. */
|
||||
export interface ManagedBackendsResponse {
|
||||
/** List of managed backends. */
|
||||
backends: ManagedBackend[]
|
||||
}
|
||||
|
||||
/** Request payload for creating a new managed backend. */
|
||||
export interface CreateManagedBackendRequest {
|
||||
/** Unique name for the backend. */
|
||||
name: string
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** API key credential string. */
|
||||
api_key?: string
|
||||
/** API base URL. */
|
||||
api_base?: string
|
||||
/** API deployment name. */
|
||||
deployment?: string
|
||||
/** API version. */
|
||||
api_version?: string
|
||||
/** Cloud project ID. */
|
||||
project?: string
|
||||
/** Cloud region. */
|
||||
region?: string
|
||||
/** AWS access key ID. */
|
||||
aws_access_key_id?: string
|
||||
/** AWS secret access key. */
|
||||
aws_secret_access_key?: string
|
||||
/** AWS session token. */
|
||||
aws_session_token?: string
|
||||
/** Requests per minute limit. */
|
||||
rpm?: number
|
||||
/** Tokens per minute limit. */
|
||||
tpm?: number
|
||||
/** Active status toggle. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload for updating an existing managed backend. */
|
||||
export type UpdateManagedBackendRequest = Partial<Omit<CreateManagedBackendRequest, 'name'>>
|
||||
@@ -0,0 +1,36 @@
|
||||
/** Details of a provider available in the LiteLLM catalog. */
|
||||
export interface CatalogProvider {
|
||||
/** Unique provider identifier. */
|
||||
id: string
|
||||
/** User-facing display name. */
|
||||
display_name: string
|
||||
/** Keep for backwards compatibility. */
|
||||
name?: string
|
||||
/** Communication protocol. */
|
||||
protocol: string
|
||||
/** Auth mechanism. */
|
||||
auth: string
|
||||
/** Proxy integration status. */
|
||||
status: 'implemented' | 'wired' | 'stub'
|
||||
/** Default base URL. */
|
||||
default_base_url: string
|
||||
/** Expected environment variables for API keys. */
|
||||
env_vars: string[]
|
||||
/** LiteLLM prefix string. */
|
||||
litellm_prefix: string
|
||||
/** Supported capabilities. */
|
||||
capabilities: {
|
||||
chat_completions: boolean
|
||||
streaming: boolean
|
||||
tool_use: boolean
|
||||
embeddings: boolean
|
||||
vision: boolean
|
||||
batch: boolean
|
||||
}
|
||||
/** Total count of models. */
|
||||
model_count: number
|
||||
/** Count of models currently cached. */
|
||||
cached_model_count: number
|
||||
/** UNIX timestamp of the last cache refresh. */
|
||||
last_refreshed: number | null
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
/** Represents a single configuration entry key-value pair. */
|
||||
export interface ConfigEntry {
|
||||
/** Configuration key. */
|
||||
key: string
|
||||
/** Configuration value. */
|
||||
value: string
|
||||
/** Timestamp when updated. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Represents the full system configuration and environment variables. */
|
||||
export interface ConfigResponse {
|
||||
/** List of config database overrides. */
|
||||
entries: ConfigEntry[]
|
||||
/** Server environment variables map. */
|
||||
env: Record<string, string>
|
||||
/** System logging level. */
|
||||
log_level: string
|
||||
/** Whether requests and responses are logged. */
|
||||
log_bodies: boolean
|
||||
/** Whether credentials are redacted from logs. */
|
||||
redact_secrets: boolean
|
||||
/** Whether to attempt to repair broken thinking blocks in Anthropic streams. */
|
||||
anthropic_thinking_repair: boolean
|
||||
/** True if pxpipe image compression is enabled. */
|
||||
pxpipe_compress: boolean
|
||||
/** CSV of model bases in pxpipe compression scope. */
|
||||
pxpipe_models: string
|
||||
/** Vision-capable Claude models offered as per-model scope toggles. */
|
||||
pxpipe_available_models: string[]
|
||||
/** Whether RTK tool-output compression is active. */
|
||||
rtk_compress: boolean
|
||||
/** CSV of model bases in RTK compression scope (empty = all models). */
|
||||
rtk_models: string
|
||||
/** Whether to pass client authorization headers to backend. */
|
||||
forward_client_auth: boolean
|
||||
/** Tool guardrail mode configured. */
|
||||
tool_guardrail_mode: string
|
||||
/** Prompt-compression optimizer mode: 'off' | 'shadow' | 'live'. */
|
||||
optimizer_mode: string
|
||||
/** Mappings of backend names to their configured big/small models. */
|
||||
backends: Record<string, { big_model: string; small_model: string }>
|
||||
/** List of keys whose overrides are active. */
|
||||
overridden_keys: string[]
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
/** Warning message generated during env import. */
|
||||
export interface EnvWarning {
|
||||
/** Affected line number. */
|
||||
line: number | null
|
||||
/** Affected environment variable key. */
|
||||
key: string | null
|
||||
/** Warning message. */
|
||||
message: string
|
||||
}
|
||||
|
||||
/** Response detailing the result of importing an env file. */
|
||||
export interface EnvImportResponse {
|
||||
/** Count of imported variables applied. */
|
||||
applied: number
|
||||
/** Non-fatal warnings generated. */
|
||||
warnings: EnvWarning[]
|
||||
}
|
||||
|
||||
/** Error payload for env file imports. */
|
||||
export interface EnvImportError {
|
||||
/** Hard/blocking validation errors. */
|
||||
hard_errors: string[]
|
||||
/** Non-fatal warnings. */
|
||||
warnings: EnvWarning[]
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
/** Represents a virtual API key configuration and its associated limits. */
|
||||
export interface VirtualKey {
|
||||
/** Unique database ID. */
|
||||
id: number
|
||||
/** The prefix of the key shown to users. */
|
||||
key_prefix: string
|
||||
/** Optional descriptive note. */
|
||||
description: string | null
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 expiration timestamp, if set. */
|
||||
expires_at: string | null
|
||||
/** ISO 8601 revocation timestamp, if revoked. */
|
||||
revoked_at: string | null
|
||||
/** Spend limit in USD. */
|
||||
spend_limit: number | null
|
||||
/** Monthly budget limit in USD. */
|
||||
max_budget_usd: number | null
|
||||
/** Duration of the budget period (e.g. 'monthly'). */
|
||||
budget_duration: string | null
|
||||
/** Requests-per-minute limit. */
|
||||
rpm_limit: number | null
|
||||
/** Tokens-per-minute limit. */
|
||||
tpm_limit: number | null
|
||||
/** Total spend in USD across the key lifetime. */
|
||||
total_spend: number
|
||||
/** Total count of requests made with this key. */
|
||||
total_requests: number
|
||||
/** Total count of tokens processed. */
|
||||
total_tokens: number
|
||||
/** ISO 8601 reset timestamp for the current budget period. */
|
||||
period_reset_at: string | null
|
||||
/** List of model names this key is restricted to, if any. */
|
||||
allowed_models: string[] | null
|
||||
/** List of route names this key is restricted to, if any. */
|
||||
allowed_routes: string[] | null
|
||||
/** Active status of the key. */
|
||||
status: 'active' | 'revoked' | 'expired' | 'override'
|
||||
/** Spend in USD during the current budget period. */
|
||||
period_spend_usd: number
|
||||
}
|
||||
|
||||
/** Spent token and requests details for a virtual key. */
|
||||
export interface KeySpend {
|
||||
/** Unique key database ID. */
|
||||
id: number
|
||||
/** Total spend in USD. */
|
||||
total_spend: number
|
||||
/** Total request count. */
|
||||
total_requests: number
|
||||
/** Total token count. */
|
||||
total_tokens: number
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
/** Represents real-time usage and performance metrics for the proxy. */
|
||||
export interface Metrics {
|
||||
/** The total number of requests received by the proxy. */
|
||||
total_requests: number
|
||||
/** The number of successful requests. */
|
||||
successful_requests: number
|
||||
/** The number of failed requests. */
|
||||
failed_requests: number
|
||||
/** The current request rate per minute. */
|
||||
requests_per_minute: number
|
||||
/** The p50 latency in milliseconds, if available. */
|
||||
p50_latency_ms: number | null
|
||||
/** The p95 latency in milliseconds, if available. */
|
||||
p95_latency_ms: number | null
|
||||
/** The percentage rate of request failures. */
|
||||
error_rate: number
|
||||
/** The number of streaming connection requests started. */
|
||||
streams_started: number
|
||||
/** The number of streaming connections completed. */
|
||||
streams_completed: number
|
||||
/** The number of streaming connections that failed. */
|
||||
streams_failed: number
|
||||
/** The number of streaming connections disconnected by the client. */
|
||||
streams_client_disconnected: number
|
||||
/** Requests where pxpipe text-to-image compression fired. */
|
||||
pxpipe_compressed_total: number
|
||||
/** Total PNG image blocks pxpipe emitted. */
|
||||
pxpipe_images_total: number
|
||||
/** Total source chars pxpipe replaced with images. */
|
||||
pxpipe_imaged_chars_total: number
|
||||
/** Requests where RTK tool-output compression fired. */
|
||||
rtk_compressed_total: number
|
||||
/** Total tool-result payloads RTK rewrote. */
|
||||
rtk_blocks_total: number
|
||||
/** Total source chars RTK removed from tool output. */
|
||||
rtk_saved_chars_total: number
|
||||
}
|
||||
|
||||
/** A time-series data point for observability charts. */
|
||||
export interface ObservabilityPoint {
|
||||
/** UNIX timestamp representing the bucket start time. */
|
||||
bucket_start: number
|
||||
/** Total requests in this bucket. */
|
||||
requests: number
|
||||
/** Total errors in this bucket. */
|
||||
errors: number
|
||||
/** Total input tokens processed in this bucket. */
|
||||
input_tokens: number
|
||||
/** Total output tokens generated in this bucket. */
|
||||
output_tokens: number
|
||||
/** Calculated cost in USD in this bucket. */
|
||||
cost_usd: number
|
||||
}
|
||||
|
||||
/** Summary of error occurrences on a backend. */
|
||||
export interface ObservabilityFailure {
|
||||
/** Class or classification of error. */
|
||||
error_kind: string
|
||||
/** Occurrence count. */
|
||||
count: number
|
||||
/** Last occurrence timestamp. */
|
||||
last_seen: string
|
||||
/** Last error message received. */
|
||||
last_message: string
|
||||
}
|
||||
|
||||
/** Represents a single trace event in the observability timeline. */
|
||||
export interface ObservabilityTimeline {
|
||||
/** Transaction request ID. */
|
||||
request_id: string
|
||||
/** ISO 8601 request timestamp. */
|
||||
timestamp: string
|
||||
/** Backend provider routed to. */
|
||||
backend: string
|
||||
/** Model mapped to. */
|
||||
model: string
|
||||
/** Latency in milliseconds. */
|
||||
latency_ms: number
|
||||
/** Request outcome status. */
|
||||
status: string
|
||||
}
|
||||
|
||||
/** Observability stats summary response. */
|
||||
export interface ObservabilityResponse {
|
||||
/** The time window in hours for metrics. */
|
||||
window_hours: number
|
||||
/** The name of the backend. */
|
||||
backend: string
|
||||
/** Total requests within the window. */
|
||||
total_requests: number
|
||||
/** Total errors within the window. */
|
||||
total_errors: number
|
||||
/** Total input tokens processed. */
|
||||
total_input_tokens: number
|
||||
/** Total output tokens generated. */
|
||||
total_output_tokens: number
|
||||
/** Total cost in USD. */
|
||||
total_cost_usd: number
|
||||
/** Historical metrics series points. */
|
||||
series: ObservabilityPoint[]
|
||||
/** Error breakdown summary. */
|
||||
failures: ObservabilityFailure[]
|
||||
/** Timeline events. */
|
||||
timeline: ObservabilityTimeline[]
|
||||
}
|
||||
|
||||
/** Real-time metrics for a specific API route. */
|
||||
export interface RouteMetrics {
|
||||
/** The request path/route. */
|
||||
path: string
|
||||
/** Number of requests per minute. */
|
||||
requests_per_min: number
|
||||
/** Percentage rate of failures. */
|
||||
error_rate: number
|
||||
/** Average latency in milliseconds. */
|
||||
avg_latency_ms: number
|
||||
/** The p95 latency in milliseconds. */
|
||||
p95_latency_ms: number
|
||||
/** Total number of requests. */
|
||||
total_requests: number
|
||||
}
|
||||
|
||||
/** Time-series requests count point for a route. */
|
||||
export interface TrafficSeriesPoint {
|
||||
/** UNIX timestamp representing the bucket start time. */
|
||||
bucket_start: number
|
||||
/** API path. */
|
||||
path: string
|
||||
/** Request count. */
|
||||
requests: number
|
||||
}
|
||||
|
||||
/** Response containing traffic analytics. */
|
||||
export interface TrafficResponse {
|
||||
/** Time window in hours. */
|
||||
window_hours: number
|
||||
/** Metrics per route. */
|
||||
routes: RouteMetrics[]
|
||||
/** Time-series data points. */
|
||||
series: TrafficSeriesPoint[]
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
/** Represents a single model configuration entry. */
|
||||
export interface ModelEntry {
|
||||
/** The unique name/identifier of the model. */
|
||||
model_name: string
|
||||
/** The number of active deployments for this model. */
|
||||
deployments: number
|
||||
}
|
||||
|
||||
/** Response shape for models list requests. */
|
||||
export interface ModelsResponse {
|
||||
/** List of model configurations. */
|
||||
models: ModelEntry[]
|
||||
/** Router strategy (e.g. priority, failover). */
|
||||
strategy: string | null
|
||||
/** Optional descriptive note. */
|
||||
note?: string
|
||||
}
|
||||
|
||||
/** Represents a model discovered from a backend API. */
|
||||
export interface DiscoveredModel {
|
||||
/** The model identifier. */
|
||||
id: string
|
||||
/** Optional human-readable name. */
|
||||
name: string | null
|
||||
}
|
||||
|
||||
/** Response containing discovered models. */
|
||||
export interface DiscoverResponse {
|
||||
/** List of discovered models. */
|
||||
models: DiscoveredModel[]
|
||||
/** Source or method used for discovery. */
|
||||
source: string
|
||||
/** True if authorization was used. */
|
||||
auth_used: boolean
|
||||
}
|
||||
|
||||
/** Details of a model available in the LiteLLM catalog. */
|
||||
export interface CatalogModel {
|
||||
/** Unique model identifier. */
|
||||
id: string
|
||||
/** Maximum context window size in tokens. */
|
||||
context_window: number
|
||||
/** Maximum output tokens. */
|
||||
max_output_tokens: number
|
||||
/** Availability status of the model. */
|
||||
status: 'available' | 'deprecated' | 'stub'
|
||||
/** Capabilities flags. */
|
||||
capabilities: {
|
||||
streaming: boolean
|
||||
tool_use: boolean
|
||||
vision: boolean
|
||||
extended_thinking: boolean
|
||||
}
|
||||
/** Pricing per million tokens. */
|
||||
pricing: {
|
||||
input_per_million_tokens: number
|
||||
output_per_million_tokens: number
|
||||
} | null
|
||||
}
|
||||
|
||||
/** Response containing model lists. */
|
||||
export interface CatalogModelsResponse {
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** True if this provider has models. */
|
||||
has_models: boolean
|
||||
/** List of models. */
|
||||
models: CatalogModel[]
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
/** Status of the optional LLMLingua-2 ONNX model artifact (optimizer scorer tier). */
|
||||
export interface OptimizerModelStatus {
|
||||
/** Proxy built with the `optimizer-onnx` feature. When false the tier is inert. */
|
||||
compiled_in: boolean
|
||||
/** Verified model artifact is present on disk. */
|
||||
present: boolean
|
||||
/** A download+verify is currently in flight. */
|
||||
downloading: boolean
|
||||
/** Last download error, if any. */
|
||||
error: string | null
|
||||
/** Pinned sha256 the artifact is verified against. */
|
||||
sha256: string
|
||||
/** Expected download size in bytes. */
|
||||
size_bytes: number
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
/** Represents the status and configuration details of the proxy server. */
|
||||
export interface ProxyStatus {
|
||||
/** True if the proxy backend has been configured. */
|
||||
configured: boolean
|
||||
/** The port number the proxy listens on for client requests. */
|
||||
proxy_port: number
|
||||
/** Whether the proxy is currently running and accepting connections. */
|
||||
proxy_running: boolean
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
/** A single request transaction log entry. */
|
||||
export interface RequestLogEntry {
|
||||
/** Unique transaction ID. */
|
||||
request_id: string
|
||||
/** ISO 8601 timestamp when the request was made. */
|
||||
timestamp: string
|
||||
/** The backend model provider targeted. */
|
||||
backend: string
|
||||
/** The model identifier requested by the client. */
|
||||
model_requested: string | null
|
||||
/** The actual model identifier routed to on the backend. */
|
||||
model_mapped: string | null
|
||||
/** HTTP response status code. */
|
||||
status_code: number
|
||||
/** Request duration in milliseconds. */
|
||||
latency_ms: number
|
||||
/** Number of input tokens processed. */
|
||||
input_tokens: number | null
|
||||
/** Number of output tokens generated. */
|
||||
output_tokens: number | null
|
||||
/** Whether the request was streamed. */
|
||||
is_streaming: boolean
|
||||
/** The raw error message returned by the backend, if any. */
|
||||
error_message: string | null
|
||||
/** Normalized error class/category. */
|
||||
error_kind: string | null
|
||||
/** The virtual key ID used for authorization, if any. */
|
||||
key_id: number | null
|
||||
/** Calculated transaction cost in USD. */
|
||||
cost_usd: number | null
|
||||
}
|
||||
|
||||
/** Paginated list response for request log queries. */
|
||||
export interface RequestsResponse {
|
||||
/** The page of request log entries. */
|
||||
requests: RequestLogEntry[]
|
||||
/** Maximum number of items returned. */
|
||||
limit: number
|
||||
/** Pagination offset. */
|
||||
offset: number
|
||||
/** True if more records are available. */
|
||||
has_more: boolean
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
/** Represents a configured API routing definition. */
|
||||
export interface Route {
|
||||
/** Unique route database ID. */
|
||||
id: string
|
||||
/** Route name. */
|
||||
name: string
|
||||
/** Optional description. */
|
||||
description: string | null
|
||||
/** Load balancing strategy. */
|
||||
strategy: string
|
||||
/** Requests per minute limit. */
|
||||
rpm: number | null
|
||||
/** Tokens per minute limit. */
|
||||
tpm: number | null
|
||||
/** Cost budget in USD. */
|
||||
budget_usd: number | null
|
||||
/** Active status toggle. */
|
||||
enabled: boolean
|
||||
/** Tool guardrail mode override. */
|
||||
guardrail_mode: string | null
|
||||
/** pxpipe image compression toggle override. */
|
||||
pxpipe_compress: boolean | null
|
||||
/** pxpipe models CSV override. */
|
||||
pxpipe_models: string | null
|
||||
/** Secret redaction toggle override. */
|
||||
redact_secrets: boolean | null
|
||||
/** Position order in the router. */
|
||||
position: number
|
||||
/** Count of providers assigned to this route. */
|
||||
provider_count: number
|
||||
/** ISO 8601 creation timestamp. */
|
||||
created_at: string
|
||||
/** ISO 8601 last update timestamp. */
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
/** Response containing routing configs. */
|
||||
export interface RoutesResponse {
|
||||
/** List of routes. */
|
||||
routes: Route[]
|
||||
}
|
||||
|
||||
/** Request payload to create a new route. */
|
||||
export interface CreateRouteRequest {
|
||||
/** Unique route name. */
|
||||
name: string
|
||||
/** Route description. */
|
||||
description?: string
|
||||
/** Load balancing strategy. */
|
||||
strategy?: string
|
||||
/** Requests per minute limit. */
|
||||
rpm?: number
|
||||
/** Tokens per minute limit. */
|
||||
tpm?: number
|
||||
/** Cost budget in USD. */
|
||||
budget_usd?: number
|
||||
/** Active status toggle. */
|
||||
enabled?: boolean
|
||||
/** Tool guardrail mode override. */
|
||||
guardrail_mode?: string | null
|
||||
/** pxpipe image compression toggle override. */
|
||||
pxpipe_compress?: boolean | null
|
||||
/** pxpipe models CSV override. */
|
||||
pxpipe_models?: string | null
|
||||
/** Secret redaction toggle override. */
|
||||
redact_secrets?: boolean | null
|
||||
/** Position order. */
|
||||
position?: number
|
||||
}
|
||||
|
||||
/** Request payload to update an existing route. */
|
||||
export type UpdateRouteRequest = Partial<CreateRouteRequest>
|
||||
|
||||
/** Represents a provider mapped to a route. */
|
||||
export interface RouteProvider {
|
||||
/** Unique assignment ID. */
|
||||
id: string
|
||||
/** Target route ID. */
|
||||
route_id: string
|
||||
/** Backend ID. */
|
||||
backend_id: string
|
||||
/** Name of the backend. */
|
||||
backend_name: string
|
||||
/** Catalog provider ID. */
|
||||
provider_id: string
|
||||
/** Supported models list. */
|
||||
models: string[]
|
||||
/** Evaluation priority. */
|
||||
priority: number
|
||||
/** Active status. */
|
||||
enabled: boolean
|
||||
}
|
||||
|
||||
/** Response containing route provider mappings. */
|
||||
export interface RouteProvidersResponse {
|
||||
/** List of mappings. */
|
||||
providers: RouteProvider[]
|
||||
}
|
||||
|
||||
/** Request payload to add a provider assignment to a route. */
|
||||
export interface AddRouteProviderRequest {
|
||||
/** Target backend ID. */
|
||||
backend_id: string
|
||||
/** Supported models list. */
|
||||
models?: string[]
|
||||
/** Priority position. */
|
||||
priority?: number
|
||||
/** Active status. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload to update a provider assignment configuration. */
|
||||
export interface UpdateRouteProviderRequest {
|
||||
/** Supported models list. */
|
||||
models?: string[]
|
||||
/** Priority position. */
|
||||
priority?: number
|
||||
/** Active status. */
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
/** Request payload to reorder provider mappings. */
|
||||
export interface ReorderRouteProvidersRequest {
|
||||
/** Ordered list of provider assignment IDs. */
|
||||
provider_ids: string[]
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
/** A single day's uptime status. */
|
||||
export interface HistoryDay {
|
||||
/** Date in YYYY-MM-DD format. */
|
||||
date: string
|
||||
/** Availability status. */
|
||||
status: 'up' | 'down' | 'degraded' | 'no-data'
|
||||
}
|
||||
|
||||
/** Uptime info for the proxy itself. */
|
||||
export interface ProxyUptimeInfo {
|
||||
/** UNIX timestamp when the proxy started. */
|
||||
started_at: number
|
||||
/** Uptime percentage over the last 30 days. */
|
||||
uptime_pct_30d: number
|
||||
/** Daily uptime history. */
|
||||
history: HistoryDay[]
|
||||
}
|
||||
|
||||
/** Uptime info for a backend endpoint. */
|
||||
export interface BackendUptimeInfo {
|
||||
/** Name of the backend. */
|
||||
name: string
|
||||
/** Current connection status. */
|
||||
status: 'up' | 'down' | 'unknown'
|
||||
/** UNIX timestamp of the last health check. */
|
||||
last_checked_at: number | null
|
||||
/** Last checked latency in milliseconds. */
|
||||
last_latency_ms: number | null
|
||||
/** 30-day uptime percentage. */
|
||||
uptime_pct_30d: number
|
||||
/** Daily uptime history. */
|
||||
history: HistoryDay[]
|
||||
}
|
||||
|
||||
/** Full uptime summary response. */
|
||||
export interface UptimeResponse {
|
||||
/** Proxy server uptime details. */
|
||||
proxy: ProxyUptimeInfo
|
||||
/** Backend endpoints uptime details. */
|
||||
backends: BackendUptimeInfo[]
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
import type { RequestLogEntry } from './requests'
|
||||
import type { Metrics } from './metrics'
|
||||
|
||||
/** Types of events sent over the admin WebSocket channel. */
|
||||
export type WSEvent =
|
||||
| { type: 'request_completed'; data: RequestLogEntry }
|
||||
| { type: 'metrics_snapshot'; data: Metrics }
|
||||
| { type: 'config_changed'; data: { key: string; value: string } }
|
||||
| { type: 'backend_health_changed'; data: { backend: string; status: 'up' | 'down'; latency_ms: number | null } }
|
||||
@@ -0,0 +1,149 @@
|
||||
import { useRef, useState } from 'react'
|
||||
import { useImportEnv, downloadEnvExport } from '../../api/queries'
|
||||
import { AdminButton, AdminSurface } from '../../components/shared/Performative'
|
||||
import type { EnvImportResponse, EnvImportError } from '../../api/types'
|
||||
|
||||
const RESTART_KEY = 'env_import_pending_restart'
|
||||
|
||||
function restartPending() {
|
||||
return sessionStorage.getItem(RESTART_KEY) === '1'
|
||||
}
|
||||
|
||||
export default function EnvFileSection() {
|
||||
const importEnv = useImportEnv()
|
||||
const fileRef = useRef<HTMLInputElement>(null)
|
||||
|
||||
const [importResult, setImportResult] = useState<EnvImportResponse | null>(null)
|
||||
const [importError, setImportError] = useState<EnvImportError | null>(null)
|
||||
const [exportError, setExportError] = useState<string | null>(null)
|
||||
const [showRestartBanner, setShowRestartBanner] = useState(restartPending)
|
||||
|
||||
function handleFileChange(e: React.ChangeEvent<HTMLInputElement>) {
|
||||
const file = e.target.files?.[0]
|
||||
if (!file) return
|
||||
setImportResult(null)
|
||||
setImportError(null)
|
||||
|
||||
importEnv.mutate(file, {
|
||||
onSuccess(data) {
|
||||
setImportResult(data)
|
||||
sessionStorage.setItem(RESTART_KEY, '1')
|
||||
setShowRestartBanner(true)
|
||||
},
|
||||
onError(err) {
|
||||
// Try to parse hard_errors from the response body
|
||||
try {
|
||||
const parsed = JSON.parse(err.message) as EnvImportError
|
||||
if (parsed.hard_errors) {
|
||||
setImportError(parsed)
|
||||
return
|
||||
}
|
||||
} catch {
|
||||
// fall through to generic error
|
||||
}
|
||||
setImportError({ hard_errors: [err.message], warnings: [] })
|
||||
},
|
||||
})
|
||||
|
||||
// Reset file input so the same file can be re-selected after fixing issues
|
||||
if (fileRef.current) fileRef.current.value = ''
|
||||
}
|
||||
|
||||
async function handleExport() {
|
||||
setExportError(null)
|
||||
try {
|
||||
await downloadEnvExport()
|
||||
} catch (err) {
|
||||
setExportError(err instanceof Error ? err.message : String(err))
|
||||
}
|
||||
}
|
||||
|
||||
function dismissRestartBanner() {
|
||||
sessionStorage.removeItem(RESTART_KEY)
|
||||
setShowRestartBanner(false)
|
||||
}
|
||||
|
||||
return (
|
||||
<div style={{ marginBottom: 24 }}>
|
||||
{/* Restart-required banner — shown after a successful import */}
|
||||
{showRestartBanner && (
|
||||
<AdminSurface className="settings-restart-banner" style={{ marginBottom: 16 }}>
|
||||
<span>Restart the proxy for imported env vars to take effect.</span>
|
||||
<AdminButton size="sm" onClick={dismissRestartBanner}>Dismiss</AdminButton>
|
||||
</AdminSurface>
|
||||
)}
|
||||
|
||||
<div className="section-label" style={{ marginBottom: 8 }}>Env File</div>
|
||||
<div style={{ display: 'flex', gap: 8, alignItems: 'center' }}>
|
||||
<input
|
||||
ref={fileRef}
|
||||
type="file"
|
||||
accept=".env,.anyllm.env,text/plain"
|
||||
style={{ display: 'none' }}
|
||||
onChange={handleFileChange}
|
||||
/>
|
||||
<AdminButton
|
||||
size="sm"
|
||||
onClick={() => fileRef.current?.click()}
|
||||
disabled={importEnv.isPending}
|
||||
loading={importEnv.isPending}
|
||||
>
|
||||
Import .anyllm.env
|
||||
</AdminButton>
|
||||
<AdminButton size="sm" onClick={handleExport}>
|
||||
Export .anyllm.env
|
||||
</AdminButton>
|
||||
</div>
|
||||
|
||||
{/* Import success */}
|
||||
{importResult && (
|
||||
<div style={{ marginTop: 10 }}>
|
||||
<div className="dim" style={{ marginBottom: 4 }}>
|
||||
{importResult.applied} variable{importResult.applied !== 1 ? 's' : ''} imported.
|
||||
{importResult.warnings.length === 0 && ' No issues.'}
|
||||
</div>
|
||||
{importResult.warnings.length > 0 && (
|
||||
<div style={{ marginTop: 8, padding: '8px 12px', background: 'var(--warn-dim)', borderLeft: '3px solid var(--warn)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4 }}>Warnings</div>
|
||||
{importResult.warnings.map((w, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>
|
||||
{w.line != null && <span className="dim">[line {w.line}] </span>}
|
||||
{w.key && <span>{w.key}: </span>}
|
||||
{w.message}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Import hard error */}
|
||||
{importError && (
|
||||
<div style={{ marginTop: 10, padding: '8px 12px', background: 'var(--err-dim)', borderLeft: '3px solid var(--err)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4 }}>Import rejected</div>
|
||||
{importError.hard_errors.map((e, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>{e}</div>
|
||||
))}
|
||||
{importError.warnings.length > 0 && (
|
||||
<>
|
||||
<div style={{ fontWeight: 600, marginTop: 8, marginBottom: 4 }}>Warnings (from partial parse)</div>
|
||||
{importError.warnings.map((w, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>
|
||||
{w.line != null && <span className="dim">[line {w.line}] </span>}
|
||||
{w.message}
|
||||
</div>
|
||||
))}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Export error */}
|
||||
{exportError && (
|
||||
<div style={{ marginTop: 10, padding: '8px 12px', background: 'var(--err-dim)', borderLeft: '3px solid var(--err)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
Export failed: {exportError}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
import { Fragment } from 'react'
|
||||
|
||||
interface EnvVariablesSectionProps {
|
||||
envData: Record<string, string> | undefined
|
||||
}
|
||||
|
||||
export default function EnvVariablesSection({ envData }: EnvVariablesSectionProps) {
|
||||
if (!envData) return null
|
||||
|
||||
return (
|
||||
<div className="readonly-section" style={{ marginTop: 16 }}>
|
||||
<div className="section-label">Environment</div>
|
||||
<div style={{ display: 'grid', gridTemplateColumns: '220px 1fr', gap: '4px 12px', marginTop: 8, fontSize: 12 }}>
|
||||
{Object.entries(envData).map(([k, v]) => (
|
||||
<Fragment key={k}>
|
||||
<span className="dim">{k}</span>
|
||||
<span className="mono">{v}</span>
|
||||
</Fragment>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
interface GettingStartedNoticeProps {
|
||||
configured: boolean
|
||||
}
|
||||
|
||||
export default function GettingStartedNotice({ configured }: GettingStartedNoticeProps) {
|
||||
if (configured) return null
|
||||
|
||||
return (
|
||||
<div style={{ marginBottom: 20, padding: '12px 16px', border: '1px solid var(--border)', borderLeft: '3px solid var(--warn)', borderRadius: 'var(--r)', fontSize: 13 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 8 }}>No backend configured — nothing to forward requests to.</div>
|
||||
<div style={{ marginBottom: 10 }}>
|
||||
Add a backend on the <span className="mono">Backends</span> tab, or configure one via env.
|
||||
The proxy needs a backend endpoint (where to forward) and a listen port (where to accept).
|
||||
LISTEN_PORT defaults to 3000. Create a <span className="mono">.anyllm.env</span> and import it below,
|
||||
or pass it at startup: <span className="mono">anyllm-proxy --webui --env-file .anyllm.env</span>
|
||||
</div>
|
||||
<div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr 1fr', gap: 10 }}>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>OpenAI</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_API_KEY=sk-...
|
||||
PROXY_API_KEYS=my-key`}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>Ollama / local LLM</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_BASE_URL=http://localhost:11434/v1
|
||||
PROXY_OPEN_RELAY=true`}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>OpenRouter / custom</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_BASE_URL=https://openrouter.ai/api/v1
|
||||
OPENAI_API_KEY=sk-or-...
|
||||
PROXY_API_KEYS=my-key`}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,376 @@
|
||||
import { useState } from 'react'
|
||||
import {
|
||||
useSaveConfig, useDeleteConfigOverride,
|
||||
useOptimizerModel, useDownloadOptimizerModel,
|
||||
} from '../../api/queries'
|
||||
import ConfirmDialog from '../../components/shared/ConfirmDialog'
|
||||
import { AdminButton } from '../../components/shared/Performative'
|
||||
import type { ConfigResponse } from '../../api/types'
|
||||
|
||||
function fmtMB(bytes: number): string {
|
||||
return `${Math.round(bytes / 1_000_000)} MB`
|
||||
}
|
||||
|
||||
interface RuntimeSettingsSectionProps {
|
||||
cfg: ConfigResponse
|
||||
}
|
||||
|
||||
export default function RuntimeSettingsSection({ cfg }: RuntimeSettingsSectionProps) {
|
||||
const save = useSaveConfig()
|
||||
const del = useDeleteConfigOverride()
|
||||
const { data: model } = useOptimizerModel()
|
||||
const downloadModel = useDownloadOptimizerModel()
|
||||
|
||||
const [form, setForm] = useState<Record<string, string>>({})
|
||||
const [pendingReset, setPendingReset] = useState<string | null>(null)
|
||||
|
||||
/** Resets a config override key back to its default value. */
|
||||
function doReset() {
|
||||
if (!pendingReset) return Promise.resolve()
|
||||
const key = pendingReset
|
||||
return del.mutateAsync(key).then(() => undefined)
|
||||
}
|
||||
|
||||
/** Saves a text configuration setting value. */
|
||||
function handleSave(key: string, currentValue: string) {
|
||||
save.mutate({ [key]: form[key] ?? currentValue })
|
||||
}
|
||||
|
||||
/** Saves a boolean configuration setting value. */
|
||||
function handleBooleanSave(key: string, value: boolean) {
|
||||
save.mutate({ [key]: value })
|
||||
}
|
||||
|
||||
// pxpipe model scope is a CSV of model bases; a model is "in scope" when any
|
||||
// base is a substring of its id (mirrors the backend's model_in_scope).
|
||||
function pxpipeScope(): string[] {
|
||||
return (cfg.pxpipe_models ?? '').split(',').map((s) => s.trim()).filter(Boolean)
|
||||
}
|
||||
function pxpipeModelChecked(model: string): boolean {
|
||||
const m = model.toLowerCase()
|
||||
return pxpipeScope().some((base) => m.includes(base.toLowerCase()))
|
||||
}
|
||||
function togglePxpipeModel(model: string, on: boolean) {
|
||||
const cur = pxpipeScope()
|
||||
const next = on
|
||||
? (pxpipeModelChecked(model) ? cur : [...cur, model])
|
||||
: cur.filter((base) => !model.toLowerCase().includes(base.toLowerCase()))
|
||||
save.mutate({ pxpipe_models: next.join(',') })
|
||||
}
|
||||
|
||||
return (
|
||||
<div>
|
||||
<div className="section-label" style={{ marginBottom: 8 }}>Runtime</div>
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-redact-secrets" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-redact-secrets"
|
||||
type="checkbox"
|
||||
checked={cfg.redact_secrets}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('redact_secrets', e.target.checked)}
|
||||
/>
|
||||
Redact secrets
|
||||
</label>
|
||||
{cfg.overridden_keys.includes('redact_secrets') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('redact_secrets')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-log-bodies" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-log-bodies"
|
||||
type="checkbox"
|
||||
checked={cfg.log_bodies}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('log_bodies', e.target.checked)}
|
||||
/>
|
||||
Log bodies
|
||||
</label>
|
||||
{cfg.overridden_keys.includes('log_bodies') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('log_bodies')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-thinking-repair" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-thinking-repair"
|
||||
type="checkbox"
|
||||
checked={cfg.anthropic_thinking_repair}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('anthropic_thinking_repair', e.target.checked)}
|
||||
/>
|
||||
Anthropic thinking-block repair
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Repairs corrupted thinking/redacted_thinking blocks in Anthropic passthrough
|
||||
requests (applies to any backend running in BACKEND=anthropic passthrough mode,
|
||||
including a named backend in a multi-backend config). Off by default.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('anthropic_thinking_repair') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('anthropic_thinking_repair')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-pxpipe-compress" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-pxpipe-compress"
|
||||
type="checkbox"
|
||||
checked={cfg.pxpipe_compress}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('pxpipe_compress', e.target.checked)}
|
||||
/>
|
||||
Image context compression (pxpipe)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Renders the stable system + tool-definition slab of Anthropic passthrough requests to a
|
||||
PNG image block to save input tokens on vision models. Off by default. Enable per-model
|
||||
below — only models that read imaged text reliably are offered.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('pxpipe_compress') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('pxpipe_compress')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
{cfg.pxpipe_compress && (
|
||||
<div style={{ marginTop: 8 }}>
|
||||
<div className="form-label" style={{ fontSize: 13 }}>Models in scope (vision-capable)</div>
|
||||
{cfg.pxpipe_available_models.length === 0 ? (
|
||||
<div className="dim" style={{ fontSize: 12 }}>No vision-capable models in the catalog.</div>
|
||||
) : (
|
||||
<div style={{ display: 'flex', flexWrap: 'wrap', gap: '4px 16px' }}>
|
||||
{cfg.pxpipe_available_models.map((model) => (
|
||||
<label key={model} style={{ display: 'flex', alignItems: 'center', gap: 6, fontSize: 12 }}>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={pxpipeModelChecked(model)}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => togglePxpipeModel(model, e.target.checked)}
|
||||
/>
|
||||
{model}
|
||||
</label>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
{cfg.overridden_keys.includes('pxpipe_models') && (
|
||||
<div className="form-row" style={{ marginTop: 6 }}>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('pxpipe_models')}>
|
||||
Reset scope
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-rtk-compress" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-rtk-compress"
|
||||
type="checkbox"
|
||||
checked={cfg.rtk_compress}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('rtk_compress', e.target.checked)}
|
||||
/>
|
||||
Tool-output compression (RTK)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Command-aware filtering of tool-result text (test/build/git/log output) using the RTK
|
||||
filter catalog. Shrinks noisy machine output before it reaches the backend; deterministic
|
||||
and cache-safe. Off by default. Applies to Anthropic passthrough and translate paths.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('rtk_compress') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('rtk_compress')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
{cfg.rtk_compress && (
|
||||
<div style={{ marginTop: 8 }}>
|
||||
<div className="form-label" style={{ fontSize: 13 }}>Models in scope (CSV, empty = all)</div>
|
||||
<input
|
||||
type="text"
|
||||
className="form-input"
|
||||
key={cfg.rtk_models}
|
||||
defaultValue={cfg.rtk_models}
|
||||
placeholder="empty = all models; e.g. claude, gpt-5"
|
||||
disabled={save.isPending}
|
||||
onBlur={(e) => {
|
||||
const v = e.target.value.trim()
|
||||
if (v === (cfg.rtk_models ?? '')) return
|
||||
save.mutate({ rtk_models: v })
|
||||
}}
|
||||
/>
|
||||
{cfg.overridden_keys.includes('rtk_models') && (
|
||||
<div className="form-row" style={{ marginTop: 6 }}>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('rtk_models')}>
|
||||
Reset scope
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-forward-client-auth" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-forward-client-auth"
|
||||
type="checkbox"
|
||||
checked={cfg.forward_client_auth}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('forward_client_auth', e.target.checked)}
|
||||
/>
|
||||
Forward client credential (Anthropic passthrough)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Forwards the client's own x-api-key/Authorization header upstream instead of the
|
||||
operator's configured credential (BACKEND=anthropic passthrough only, single-key/BYOK
|
||||
deployments). The proxy refuses to enable this with 2+ PROXY_API_KEYS entries and no
|
||||
PROXY_OPEN_RELAY. Off by default.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('forward_client_auth') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('forward_client_auth')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-tool-guardrail-mode">Tool guardrail mode</label>
|
||||
<div className="form-row">
|
||||
<select
|
||||
id="cfg-tool-guardrail-mode"
|
||||
value={cfg.tool_guardrail_mode}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => save.mutate({ tool_guardrail_mode: e.target.value })}
|
||||
>
|
||||
<option value="disabled">Disabled</option>
|
||||
<option value="standard">Standard</option>
|
||||
</select>
|
||||
{cfg.overridden_keys.includes('tool_guardrail_mode') && (
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('tool_guardrail_mode')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
)}
|
||||
</div>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Applies advisory guardrails to tool calls the proxy auto-executes. Disabled by default.
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-optimizer-mode">Prompt compression (optimizer)</label>
|
||||
<div className="form-row">
|
||||
<select
|
||||
id="cfg-optimizer-mode"
|
||||
value={cfg.optimizer_mode}
|
||||
disabled={save.isPending || (!!model?.compiled_in && !model?.present)}
|
||||
onChange={(e) => save.mutate({ optimizer_mode: e.target.value })}
|
||||
>
|
||||
<option value="off">Off</option>
|
||||
<option value="shadow">Shadow (report only)</option>
|
||||
<option value="live">Live (compress)</option>
|
||||
</select>
|
||||
{cfg.overridden_keys.includes('optimizer_mode') && (
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('optimizer_mode')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
)}
|
||||
</div>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Frozen-Frontier compression of long conversation history (latest turn untouched). Off by default.
|
||||
</div>
|
||||
|
||||
{model && !model.compiled_in && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
Heuristic scorer only. Rebuild the proxy with <code>--features optimizer-onnx</code> to enable the LLMLingua-2 ONNX scorer.
|
||||
</div>
|
||||
)}
|
||||
{model?.compiled_in && !model.present && !model.downloading && (
|
||||
<div className="form-row" style={{ marginTop: 8 }}>
|
||||
<AdminButton
|
||||
size="sm"
|
||||
disabled={downloadModel.isPending}
|
||||
onClick={() => downloadModel.mutate()}
|
||||
>
|
||||
Download model ({fmtMB(model.size_bytes)})
|
||||
</AdminButton>
|
||||
<span className="dim" style={{ fontSize: 12 }}>
|
||||
Required before enabling. Verified against a pinned sha256.
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{model?.downloading && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
Downloading and verifying model ({fmtMB(model.size_bytes)})…
|
||||
</div>
|
||||
)}
|
||||
{model?.error && !model.downloading && (
|
||||
<div style={{ fontSize: 12, marginTop: 8, color: 'var(--danger, #c0392b)' }}>
|
||||
Download failed: {model.error}
|
||||
</div>
|
||||
)}
|
||||
{model?.compiled_in && model.present && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
ONNX scorer ready — live mode uses LLMLingua-2 (loaded on the next request).
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{cfg.entries.filter((entry) => !['redact_secrets', 'log_bodies', 'anthropic_thinking_repair', 'pxpipe_compress', 'pxpipe_models', 'rtk_compress', 'rtk_models', 'forward_client_auth', 'tool_guardrail_mode', 'optimizer_mode'].includes(entry.key)).map((entry) => {
|
||||
const inputId = `cfg-${entry.key}`
|
||||
return (
|
||||
<div className="form-group" key={entry.key}>
|
||||
<label className="form-label" htmlFor={inputId}>{entry.key}</label>
|
||||
<div className="form-row">
|
||||
<input
|
||||
id={inputId}
|
||||
name={entry.key}
|
||||
value={form[entry.key] ?? entry.value}
|
||||
onChange={(e) => setForm((f) => ({ ...f, [entry.key]: e.target.value }))}
|
||||
/>
|
||||
<AdminButton tone="primary" size="sm" onClick={() => handleSave(entry.key, entry.value)}>Save</AdminButton>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset(entry.key)}>Reset</AdminButton>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
})}
|
||||
|
||||
<ConfirmDialog
|
||||
open={pendingReset !== null}
|
||||
onClose={() => setPendingReset(null)}
|
||||
onConfirm={doReset}
|
||||
title="Reset override?"
|
||||
message={
|
||||
<>
|
||||
Reset override for <span className="mono">{pendingReset}</span>? The runtime value will revert
|
||||
to the env-file or default. Active connections are not affected.
|
||||
</>
|
||||
}
|
||||
confirmLabel="Reset"
|
||||
variant="primary"
|
||||
/>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -1,126 +1,19 @@
|
||||
import { Fragment, useRef, useState } from 'react'
|
||||
import {
|
||||
useConfig, useSaveConfig, useDeleteConfigOverride, useEnv,
|
||||
useImportEnv, downloadEnvExport, useStatus,
|
||||
useOptimizerModel, useDownloadOptimizerModel,
|
||||
} from '../../api/queries'
|
||||
import { useConfig, useEnv, useStatus } from '../../api/queries'
|
||||
import EmptyState from '../../components/shared/EmptyState'
|
||||
import ConfirmDialog from '../../components/shared/ConfirmDialog'
|
||||
import { AdminButton, AdminSurface } from '../../components/shared/Performative'
|
||||
import type { EnvImportResponse, EnvImportError } from '../../api/types'
|
||||
|
||||
const RESTART_KEY = 'env_import_pending_restart'
|
||||
|
||||
function restartPending() {
|
||||
return sessionStorage.getItem(RESTART_KEY) === '1'
|
||||
}
|
||||
|
||||
function fmtMB(bytes: number): string {
|
||||
return `${Math.round(bytes / 1_000_000)} MB`
|
||||
}
|
||||
import GettingStartedNotice from './GettingStartedNotice'
|
||||
import EnvFileSection from './EnvFileSection'
|
||||
import RuntimeSettingsSection from './RuntimeSettingsSection'
|
||||
import EnvVariablesSection from './EnvVariablesSection'
|
||||
|
||||
/**
|
||||
* Settings Component.
|
||||
* Provides controls for configuring system settings, environment variables export/import,
|
||||
* and proxy properties.
|
||||
* Coordinates and provides controls for configuring system settings,
|
||||
* environment variables export/import, and proxy properties.
|
||||
*/
|
||||
export default function Settings({ configured = true }: { configured?: boolean }) {
|
||||
const { data: cfg, isLoading, error } = useConfig()
|
||||
const { data: model } = useOptimizerModel()
|
||||
const downloadModel = useDownloadOptimizerModel()
|
||||
const { data: envData } = useEnv()
|
||||
const { data: status } = useStatus()
|
||||
const save = useSaveConfig()
|
||||
const del = useDeleteConfigOverride()
|
||||
const importEnv = useImportEnv()
|
||||
const fileRef = useRef<HTMLInputElement>(null)
|
||||
|
||||
const [form, setForm] = useState<Record<string, string>>({})
|
||||
const [importResult, setImportResult] = useState<EnvImportResponse | null>(null)
|
||||
const [importError, setImportError] = useState<EnvImportError | null>(null)
|
||||
const [exportError, setExportError] = useState<string | null>(null)
|
||||
const [showRestartBanner, setShowRestartBanner] = useState(restartPending)
|
||||
const [pendingReset, setPendingReset] = useState<string | null>(null)
|
||||
|
||||
/** Resets a config override key back to its default value. */
|
||||
function doReset() {
|
||||
if (!pendingReset) return Promise.resolve()
|
||||
const key = pendingReset
|
||||
return del.mutateAsync(key).then(() => undefined)
|
||||
}
|
||||
|
||||
/** Saves a text configuration setting value. */
|
||||
function handleSave(key: string, currentValue: string) {
|
||||
save.mutate({ [key]: form[key] ?? currentValue })
|
||||
}
|
||||
|
||||
/** Saves a boolean configuration setting value. */
|
||||
function handleBooleanSave(key: string, value: boolean) {
|
||||
save.mutate({ [key]: value })
|
||||
}
|
||||
|
||||
// pxpipe model scope is a CSV of model bases; a model is "in scope" when any
|
||||
// base is a substring of its id (mirrors the backend's model_in_scope).
|
||||
function pxpipeScope(): string[] {
|
||||
return (cfg?.pxpipe_models ?? '').split(',').map((s) => s.trim()).filter(Boolean)
|
||||
}
|
||||
function pxpipeModelChecked(model: string): boolean {
|
||||
const m = model.toLowerCase()
|
||||
return pxpipeScope().some((base) => m.includes(base.toLowerCase()))
|
||||
}
|
||||
function togglePxpipeModel(model: string, on: boolean) {
|
||||
const cur = pxpipeScope()
|
||||
const next = on
|
||||
? (pxpipeModelChecked(model) ? cur : [...cur, model])
|
||||
: cur.filter((base) => !model.toLowerCase().includes(base.toLowerCase()))
|
||||
save.mutate({ pxpipe_models: next.join(',') })
|
||||
}
|
||||
|
||||
function handleFileChange(e: React.ChangeEvent<HTMLInputElement>) {
|
||||
const file = e.target.files?.[0]
|
||||
if (!file) return
|
||||
setImportResult(null)
|
||||
setImportError(null)
|
||||
|
||||
importEnv.mutate(file, {
|
||||
onSuccess(data) {
|
||||
setImportResult(data)
|
||||
sessionStorage.setItem(RESTART_KEY, '1')
|
||||
setShowRestartBanner(true)
|
||||
},
|
||||
onError(err) {
|
||||
// Try to parse hard_errors from the response body
|
||||
try {
|
||||
const parsed = JSON.parse(err.message) as EnvImportError
|
||||
if (parsed.hard_errors) {
|
||||
setImportError(parsed)
|
||||
return
|
||||
}
|
||||
} catch {
|
||||
// fall through to generic error
|
||||
}
|
||||
setImportError({ hard_errors: [err.message], warnings: [] })
|
||||
},
|
||||
})
|
||||
|
||||
// Reset file input so the same file can be re-selected after fixing issues
|
||||
if (fileRef.current) fileRef.current.value = ''
|
||||
}
|
||||
|
||||
async function handleExport() {
|
||||
setExportError(null)
|
||||
try {
|
||||
await downloadEnvExport()
|
||||
} catch (err) {
|
||||
setExportError(err instanceof Error ? err.message : String(err))
|
||||
}
|
||||
}
|
||||
|
||||
/** Dismisses the restart banner reminding the operator that imports require a restart. */
|
||||
function dismissRestartBanner() {
|
||||
sessionStorage.removeItem(RESTART_KEY)
|
||||
setShowRestartBanner(false)
|
||||
}
|
||||
|
||||
const proxyUrl = status ? `http://${window.location.hostname}:${status.proxy_port}` : ''
|
||||
|
||||
@@ -142,458 +35,19 @@ export default function Settings({ configured = true }: { configured?: boolean }
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Getting-started notice — shown when no backend is configured (no env/config-file
|
||||
signal and no managed backend). Manage backends on the Backends tab. */}
|
||||
{!configured && (
|
||||
<div style={{ marginBottom: 20, padding: '12px 16px', border: '1px solid var(--border)', borderLeft: '3px solid var(--warn)', borderRadius: 'var(--r)', fontSize: 13 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 8 }}>No backend configured — nothing to forward requests to.</div>
|
||||
<div style={{ marginBottom: 10 }}>
|
||||
Add a backend on the <span className="mono">Backends</span> tab, or configure one via env.
|
||||
The proxy needs a backend endpoint (where to forward) and a listen port (where to accept).
|
||||
LISTEN_PORT defaults to 3000. Create a <span className="mono">.anyllm.env</span> and import it below,
|
||||
or pass it at startup: <span className="mono">anyllm-proxy --webui --env-file .anyllm.env</span>
|
||||
</div>
|
||||
<div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr 1fr', gap: 10 }}>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>OpenAI</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_API_KEY=sk-...
|
||||
PROXY_API_KEYS=my-key`}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>Ollama / local LLM</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_BASE_URL=http://localhost:11434/v1
|
||||
PROXY_OPEN_RELAY=true`}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, fontSize: 12 }}>OpenRouter / custom</div>
|
||||
<pre style={{ margin: 0, padding: '6px 10px', background: 'var(--surface-2)', borderRadius: 'var(--r)', fontSize: 11, overflowX: 'auto' }}>
|
||||
{`OPENAI_BASE_URL=https://openrouter.ai/api/v1
|
||||
OPENAI_API_KEY=sk-or-...
|
||||
PROXY_API_KEYS=my-key`}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Restart-required banner — shown after a successful import */}
|
||||
{showRestartBanner && (
|
||||
<AdminSurface className="settings-restart-banner">
|
||||
<span>Restart the proxy for imported env vars to take effect.</span>
|
||||
<AdminButton size="sm" onClick={dismissRestartBanner}>Dismiss</AdminButton>
|
||||
</AdminSurface>
|
||||
)}
|
||||
{/* Getting-started notice — shown when no backend is configured */}
|
||||
<GettingStartedNotice configured={configured} />
|
||||
|
||||
{/* Env file import / export */}
|
||||
<div style={{ marginBottom: 24 }}>
|
||||
<div className="section-label" style={{ marginBottom: 8 }}>Env File</div>
|
||||
<div style={{ display: 'flex', gap: 8, alignItems: 'center' }}>
|
||||
<input
|
||||
ref={fileRef}
|
||||
type="file"
|
||||
accept=".env,.anyllm.env,text/plain"
|
||||
style={{ display: 'none' }}
|
||||
onChange={handleFileChange}
|
||||
/>
|
||||
<AdminButton
|
||||
size="sm"
|
||||
onClick={() => fileRef.current?.click()}
|
||||
disabled={importEnv.isPending}
|
||||
loading={importEnv.isPending}
|
||||
>
|
||||
Import .anyllm.env
|
||||
</AdminButton>
|
||||
<AdminButton size="sm" onClick={handleExport}>
|
||||
Export .anyllm.env
|
||||
</AdminButton>
|
||||
</div>
|
||||
|
||||
{/* Import success */}
|
||||
{importResult && (
|
||||
<div style={{ marginTop: 10 }}>
|
||||
<div className="dim" style={{ marginBottom: 4 }}>
|
||||
{importResult.applied} variable{importResult.applied !== 1 ? 's' : ''} imported.
|
||||
{importResult.warnings.length === 0 && ' No issues.'}
|
||||
</div>
|
||||
{importResult.warnings.length > 0 && (
|
||||
<div style={{ marginTop: 8, padding: '8px 12px', background: 'var(--warn-dim)', borderLeft: '3px solid var(--warn)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4 }}>Warnings</div>
|
||||
{importResult.warnings.map((w, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>
|
||||
{w.line != null && <span className="dim">[line {w.line}] </span>}
|
||||
{w.key && <span>{w.key}: </span>}
|
||||
{w.message}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Import hard error */}
|
||||
{importError && (
|
||||
<div style={{ marginTop: 10, padding: '8px 12px', background: 'var(--err-dim)', borderLeft: '3px solid var(--err)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4 }}>Import rejected</div>
|
||||
{importError.hard_errors.map((e, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>{e}</div>
|
||||
))}
|
||||
{importError.warnings.length > 0 && (
|
||||
<>
|
||||
<div style={{ fontWeight: 600, marginTop: 8, marginBottom: 4 }}>Warnings (from partial parse)</div>
|
||||
{importError.warnings.map((w, i) => (
|
||||
<div key={i} className="mono" style={{ fontSize: 12 }}>
|
||||
{w.line != null && <span className="dim">[line {w.line}] </span>}
|
||||
{w.message}
|
||||
</div>
|
||||
))}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Export error */}
|
||||
{exportError && (
|
||||
<div style={{ marginTop: 10, padding: '8px 12px', background: 'var(--err-dim)', borderLeft: '3px solid var(--err)', borderRadius: 'var(--r)', fontSize: 12 }}>
|
||||
Export failed: {exportError}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
<EnvFileSection />
|
||||
|
||||
<EmptyState loading={isLoading} error={error?.message} />
|
||||
{cfg && (
|
||||
<div>
|
||||
<div className="section-label" style={{ marginBottom: 8 }}>Runtime</div>
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-redact-secrets" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-redact-secrets"
|
||||
type="checkbox"
|
||||
checked={cfg.redact_secrets}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('redact_secrets', e.target.checked)}
|
||||
/>
|
||||
Redact secrets
|
||||
</label>
|
||||
{cfg.overridden_keys.includes('redact_secrets') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('redact_secrets')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-log-bodies" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-log-bodies"
|
||||
type="checkbox"
|
||||
checked={cfg.log_bodies}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('log_bodies', e.target.checked)}
|
||||
/>
|
||||
Log bodies
|
||||
</label>
|
||||
{cfg.overridden_keys.includes('log_bodies') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('log_bodies')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
{/* Runtime settings override form */}
|
||||
{cfg && <RuntimeSettingsSection cfg={cfg} />}
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-thinking-repair" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-thinking-repair"
|
||||
type="checkbox"
|
||||
checked={cfg.anthropic_thinking_repair}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('anthropic_thinking_repair', e.target.checked)}
|
||||
/>
|
||||
Anthropic thinking-block repair
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Repairs corrupted thinking/redacted_thinking blocks in Anthropic passthrough
|
||||
requests (applies to any backend running in BACKEND=anthropic passthrough mode,
|
||||
including a named backend in a multi-backend config). Off by default.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('anthropic_thinking_repair') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('anthropic_thinking_repair')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-pxpipe-compress" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-pxpipe-compress"
|
||||
type="checkbox"
|
||||
checked={cfg.pxpipe_compress}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('pxpipe_compress', e.target.checked)}
|
||||
/>
|
||||
Image context compression (pxpipe)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Renders the stable system + tool-definition slab of Anthropic passthrough requests to a
|
||||
PNG image block to save input tokens on vision models. Off by default. Enable per-model
|
||||
below — only models that read imaged text reliably are offered.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('pxpipe_compress') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('pxpipe_compress')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
{cfg.pxpipe_compress && (
|
||||
<div style={{ marginTop: 8 }}>
|
||||
<div className="form-label" style={{ fontSize: 13 }}>Models in scope (vision-capable)</div>
|
||||
{cfg.pxpipe_available_models.length === 0 ? (
|
||||
<div className="dim" style={{ fontSize: 12 }}>No vision-capable models in the catalog.</div>
|
||||
) : (
|
||||
<div style={{ display: 'flex', flexWrap: 'wrap', gap: '4px 16px' }}>
|
||||
{cfg.pxpipe_available_models.map((model) => (
|
||||
<label key={model} style={{ display: 'flex', alignItems: 'center', gap: 6, fontSize: 12 }}>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={pxpipeModelChecked(model)}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => togglePxpipeModel(model, e.target.checked)}
|
||||
/>
|
||||
{model}
|
||||
</label>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
{cfg.overridden_keys.includes('pxpipe_models') && (
|
||||
<div className="form-row" style={{ marginTop: 6 }}>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('pxpipe_models')}>
|
||||
Reset scope
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-rtk-compress" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-rtk-compress"
|
||||
type="checkbox"
|
||||
checked={cfg.rtk_compress}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('rtk_compress', e.target.checked)}
|
||||
/>
|
||||
Tool-output compression (RTK)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Command-aware filtering of tool-result text (test/build/git/log output) using the RTK
|
||||
filter catalog. Shrinks noisy machine output before it reaches the backend; deterministic
|
||||
and cache-safe. Off by default. Applies to Anthropic passthrough and translate paths.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('rtk_compress') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('rtk_compress')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
{cfg.rtk_compress && (
|
||||
<div style={{ marginTop: 8 }}>
|
||||
<div className="form-label" style={{ fontSize: 13 }}>Models in scope (CSV, empty = all)</div>
|
||||
<input
|
||||
type="text"
|
||||
className="form-input"
|
||||
key={cfg.rtk_models}
|
||||
defaultValue={cfg.rtk_models}
|
||||
placeholder="empty = all models; e.g. claude, gpt-5"
|
||||
disabled={save.isPending}
|
||||
onBlur={(e) => {
|
||||
const v = e.target.value.trim()
|
||||
if (v === (cfg.rtk_models ?? '')) return
|
||||
save.mutate({ rtk_models: v })
|
||||
}}
|
||||
/>
|
||||
{cfg.overridden_keys.includes('rtk_models') && (
|
||||
<div className="form-row" style={{ marginTop: 6 }}>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('rtk_models')}>
|
||||
Reset scope
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-forward-client-auth" style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
|
||||
<input
|
||||
id="cfg-forward-client-auth"
|
||||
type="checkbox"
|
||||
checked={cfg.forward_client_auth}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => handleBooleanSave('forward_client_auth', e.target.checked)}
|
||||
/>
|
||||
Forward client credential (Anthropic passthrough)
|
||||
</label>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Forwards the client's own x-api-key/Authorization header upstream instead of the
|
||||
operator's configured credential (BACKEND=anthropic passthrough only, single-key/BYOK
|
||||
deployments). The proxy refuses to enable this with 2+ PROXY_API_KEYS entries and no
|
||||
PROXY_OPEN_RELAY. Off by default.
|
||||
</div>
|
||||
{cfg.overridden_keys.includes('forward_client_auth') && (
|
||||
<div className="form-row">
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('forward_client_auth')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-tool-guardrail-mode">Tool guardrail mode</label>
|
||||
<div className="form-row">
|
||||
<select
|
||||
id="cfg-tool-guardrail-mode"
|
||||
value={cfg.tool_guardrail_mode}
|
||||
disabled={save.isPending}
|
||||
onChange={(e) => save.mutate({ tool_guardrail_mode: e.target.value })}
|
||||
>
|
||||
<option value="disabled">Disabled</option>
|
||||
<option value="standard">Standard</option>
|
||||
</select>
|
||||
{cfg.overridden_keys.includes('tool_guardrail_mode') && (
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('tool_guardrail_mode')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
)}
|
||||
</div>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Applies advisory guardrails to tool calls the proxy auto-executes. Disabled by default.
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="form-group">
|
||||
<label className="form-label" htmlFor="cfg-optimizer-mode">Prompt compression (optimizer)</label>
|
||||
<div className="form-row">
|
||||
<select
|
||||
id="cfg-optimizer-mode"
|
||||
value={cfg.optimizer_mode}
|
||||
// Gate enabling on the ONNX model when the tier is compiled in but the
|
||||
// model isn't downloaded yet (point 2: "not toggle if not detected").
|
||||
disabled={save.isPending || (!!model?.compiled_in && !model?.present)}
|
||||
onChange={(e) => save.mutate({ optimizer_mode: e.target.value })}
|
||||
>
|
||||
<option value="off">Off</option>
|
||||
<option value="shadow">Shadow (report only)</option>
|
||||
<option value="live">Live (compress)</option>
|
||||
</select>
|
||||
{cfg.overridden_keys.includes('optimizer_mode') && (
|
||||
<AdminButton size="sm" onClick={() => setPendingReset('optimizer_mode')}>
|
||||
Reset
|
||||
</AdminButton>
|
||||
)}
|
||||
</div>
|
||||
<div className="dim" style={{ fontSize: 12 }}>
|
||||
Frozen-Frontier compression of long conversation history (latest turn untouched). Off by default.
|
||||
</div>
|
||||
|
||||
{model && !model.compiled_in && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
Heuristic scorer only. Rebuild the proxy with <code>--features optimizer-onnx</code> to enable the LLMLingua-2 ONNX scorer.
|
||||
</div>
|
||||
)}
|
||||
{model?.compiled_in && !model.present && !model.downloading && (
|
||||
<div className="form-row" style={{ marginTop: 8 }}>
|
||||
<AdminButton
|
||||
size="sm"
|
||||
disabled={downloadModel.isPending}
|
||||
onClick={() => downloadModel.mutate()}
|
||||
>
|
||||
Download model ({fmtMB(model.size_bytes)})
|
||||
</AdminButton>
|
||||
<span className="dim" style={{ fontSize: 12 }}>
|
||||
Required before enabling. Verified against a pinned sha256.
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{model?.downloading && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
Downloading and verifying model ({fmtMB(model.size_bytes)})…
|
||||
</div>
|
||||
)}
|
||||
{model?.error && !model.downloading && (
|
||||
<div style={{ fontSize: 12, marginTop: 8, color: 'var(--danger, #c0392b)' }}>
|
||||
Download failed: {model.error}
|
||||
</div>
|
||||
)}
|
||||
{model?.compiled_in && model.present && (
|
||||
<div className="dim" style={{ fontSize: 12, marginTop: 8 }}>
|
||||
ONNX scorer ready — live mode uses LLMLingua-2 (loaded on the next request).
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{cfg.entries.filter((entry) => !['redact_secrets', 'log_bodies', 'anthropic_thinking_repair', 'pxpipe_compress', 'pxpipe_models', 'rtk_compress', 'rtk_models', 'forward_client_auth', 'tool_guardrail_mode', 'optimizer_mode'].includes(entry.key)).map((entry) => {
|
||||
const inputId = `cfg-${entry.key}`
|
||||
return (
|
||||
<div className="form-group" key={entry.key}>
|
||||
<label className="form-label" htmlFor={inputId}>{entry.key}</label>
|
||||
<div className="form-row">
|
||||
<input
|
||||
id={inputId}
|
||||
name={entry.key}
|
||||
value={form[entry.key] ?? entry.value}
|
||||
onChange={(e) => setForm((f) => ({ ...f, [entry.key]: e.target.value }))}
|
||||
/>
|
||||
<AdminButton tone="primary" size="sm" onClick={() => handleSave(entry.key, entry.value)}>Save</AdminButton>
|
||||
<AdminButton size="sm" onClick={() => setPendingReset(entry.key)}>Reset</AdminButton>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
{envData && (
|
||||
<div className="readonly-section" style={{ marginTop: 16 }}>
|
||||
<div className="section-label">Environment</div>
|
||||
<div style={{ display: 'grid', gridTemplateColumns: '220px 1fr', gap: '4px 12px', marginTop: 8, fontSize: 12 }}>
|
||||
{Object.entries(envData).map(([k, v]) => (
|
||||
<Fragment key={k}>
|
||||
<span className="dim">{k}</span>
|
||||
<span className="mono">{v}</span>
|
||||
</Fragment>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<ConfirmDialog
|
||||
open={pendingReset !== null}
|
||||
onClose={() => setPendingReset(null)}
|
||||
onConfirm={doReset}
|
||||
title="Reset override?"
|
||||
message={
|
||||
<>
|
||||
Reset override for <span className="mono">{pendingReset}</span>? The runtime value will revert
|
||||
to the env-file or default. Active connections are not affected.
|
||||
</>
|
||||
}
|
||||
confirmLabel="Reset"
|
||||
variant="primary"
|
||||
/>
|
||||
{/* Environment grid */}
|
||||
<EnvVariablesSection envData={envData} />
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -875,6 +875,30 @@
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6",
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-luna",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"output_cost_per_token": 6e-06,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-sol",
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-5.6-terra",
|
||||
"input_cost_per_token": 2.5e-06,
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-audio",
|
||||
"input_cost_per_token": 2.5e-06,
|
||||
@@ -929,6 +953,18 @@
|
||||
"output_cost_per_token": 1.6e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2.1",
|
||||
"input_cost_per_token": 4e-06,
|
||||
"output_cost_per_token": 2.4e-05,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2.1-mini",
|
||||
"input_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"model_pattern": "gpt-realtime-2025-08-28",
|
||||
"input_cost_per_token": 4e-06,
|
||||
|
||||
@@ -458,212 +458,5 @@ pub fn reorder_route_providers(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::admin::db::backends::ManagedBackendRow;
|
||||
|
||||
fn in_memory_db() -> Connection {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
super::super::init_db(&conn).unwrap();
|
||||
conn
|
||||
}
|
||||
|
||||
fn test_row(name: &str) -> ManagedBackendRow {
|
||||
ManagedBackendRow {
|
||||
id: format!("id-{name}"),
|
||||
name: name.to_string(),
|
||||
provider_id: "openai".to_string(),
|
||||
api_key: Some("sk-test".to_string()),
|
||||
api_base: None,
|
||||
deployment: None,
|
||||
api_version: None,
|
||||
project: None,
|
||||
region: None,
|
||||
aws_access_key_id: None,
|
||||
aws_secret_access_key: None,
|
||||
aws_session_token: None,
|
||||
rpm: Some(100),
|
||||
tpm: Some(10_000),
|
||||
enabled: true,
|
||||
created_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
updated_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn seed_route_with_providers(conn: &Connection, count: usize) -> (String, Vec<String>) {
|
||||
let route = RouteRow {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
name: "r".into(),
|
||||
description: None,
|
||||
strategy: "failover".into(),
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: true,
|
||||
guardrail_mode: None,
|
||||
pxpipe_compress: None,
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: 0,
|
||||
created_at: now_iso8601(),
|
||||
updated_at: now_iso8601(),
|
||||
};
|
||||
insert_route(conn, &route).unwrap();
|
||||
|
||||
for i in 0..count {
|
||||
let mut b = test_row(&format!("b{i}"));
|
||||
b.id = format!("backend-{i}");
|
||||
crate::admin::db::backends::insert_managed_backend(conn, &b).unwrap();
|
||||
add_route_provider(conn, &route.id, &b.id, &["*".to_string()], i as i32, true).unwrap();
|
||||
}
|
||||
|
||||
let ids: Vec<String> = list_route_providers(conn, &route.id)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|p| p.id)
|
||||
.collect();
|
||||
(route.id, ids)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_route_sets_and_clears_option_override() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, _) = seed_route_with_providers(&conn, 1);
|
||||
|
||||
// Set an override.
|
||||
let set = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: Some(false),
|
||||
guardrail_mode: Some(Some("standard".into())),
|
||||
pxpipe_compress: Some(Some(true)),
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &set).unwrap());
|
||||
let r = get_route(&conn, &route_id).unwrap().unwrap();
|
||||
assert!(!r.enabled);
|
||||
assert_eq!(r.guardrail_mode.as_deref(), Some("standard"));
|
||||
assert_eq!(r.pxpipe_compress, Some(true));
|
||||
|
||||
// Clear the override back to NULL (inherit).
|
||||
let clear = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: None,
|
||||
guardrail_mode: Some(None),
|
||||
pxpipe_compress: Some(None),
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &clear).unwrap());
|
||||
let r = get_route(&conn, &route_id).unwrap().unwrap();
|
||||
assert_eq!(r.guardrail_mode, None);
|
||||
assert_eq!(r.pxpipe_compress, None);
|
||||
// enabled was left unchanged (None) — still false.
|
||||
assert!(!r.enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_route_excluded_from_enabled_route_ids() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, _) = seed_route_with_providers(&conn, 1);
|
||||
// The seeded backend is "b0"; while enabled the route id is returned.
|
||||
assert_eq!(
|
||||
enabled_route_ids_for_backend_name(&conn, "b0").unwrap(),
|
||||
vec![route_id.clone()]
|
||||
);
|
||||
|
||||
// Disable the route -> it drops out of the virtual-key scope query.
|
||||
let patch = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: Some(false),
|
||||
guardrail_mode: None,
|
||||
pxpipe_compress: None,
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &patch).unwrap());
|
||||
assert!(enabled_route_ids_for_backend_name(&conn, "b0")
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_rewrites_priorities() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Reverse the order.
|
||||
let reversed: Vec<String> = ids.iter().rev().cloned().collect();
|
||||
let outcome = reorder_route_providers(&conn, &route_id, &reversed).unwrap();
|
||||
|
||||
match outcome {
|
||||
ReorderOutcome::Ok(rows) => {
|
||||
let new_order: Vec<String> = rows.iter().map(|p| p.id.clone()).collect();
|
||||
assert_eq!(new_order, reversed);
|
||||
for (i, row) in rows.iter().enumerate() {
|
||||
assert_eq!(row.priority, i as i32);
|
||||
}
|
||||
}
|
||||
ReorderOutcome::Mismatch => panic!("expected Ok"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_mismatch_rolls_back() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Submit a subset — should be rejected as Mismatch, priorities unchanged.
|
||||
let partial: Vec<String> = ids.iter().take(2).cloned().collect();
|
||||
let outcome = reorder_route_providers(&conn, &route_id, &partial).unwrap();
|
||||
assert!(matches!(outcome, ReorderOutcome::Mismatch));
|
||||
|
||||
// Priorities must be untouched.
|
||||
let rows = list_route_providers(&conn, &route_id).unwrap();
|
||||
for (i, row) in rows.iter().enumerate() {
|
||||
assert_eq!(
|
||||
row.priority, i as i32,
|
||||
"priorities must be unchanged after Mismatch"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_rejects_duplicates_and_extras() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Duplicate id.
|
||||
let dup = vec![ids[0].clone(), ids[0].clone(), ids[1].clone()];
|
||||
assert!(matches!(
|
||||
reorder_route_providers(&conn, &route_id, &dup).unwrap(),
|
||||
ReorderOutcome::Mismatch
|
||||
));
|
||||
|
||||
// Extra id that isn't a provider on this route.
|
||||
let mut extra = ids.clone();
|
||||
extra.push("bogus-id".into());
|
||||
assert!(matches!(
|
||||
reorder_route_providers(&conn, &route_id, &extra).unwrap(),
|
||||
ReorderOutcome::Mismatch
|
||||
));
|
||||
}
|
||||
}
|
||||
#[path = "routes/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
use super::*;
|
||||
use crate::admin::db::backends::ManagedBackendRow;
|
||||
|
||||
fn in_memory_db() -> Connection {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
super::super::init_db(&conn).unwrap();
|
||||
conn
|
||||
}
|
||||
|
||||
fn test_row(name: &str) -> ManagedBackendRow {
|
||||
ManagedBackendRow {
|
||||
id: format!("id-{name}"),
|
||||
name: name.to_string(),
|
||||
provider_id: "openai".to_string(),
|
||||
api_key: Some("sk-test".to_string()),
|
||||
api_base: None,
|
||||
deployment: None,
|
||||
api_version: None,
|
||||
project: None,
|
||||
region: None,
|
||||
aws_access_key_id: None,
|
||||
aws_secret_access_key: None,
|
||||
aws_session_token: None,
|
||||
rpm: Some(100),
|
||||
tpm: Some(10_000),
|
||||
enabled: true,
|
||||
created_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
updated_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn seed_route_with_providers(conn: &Connection, count: usize) -> (String, Vec<String>) {
|
||||
let route = RouteRow {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
name: "r".into(),
|
||||
description: None,
|
||||
strategy: "failover".into(),
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: true,
|
||||
guardrail_mode: None,
|
||||
pxpipe_compress: None,
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: 0,
|
||||
created_at: now_iso8601(),
|
||||
updated_at: now_iso8601(),
|
||||
};
|
||||
insert_route(conn, &route).unwrap();
|
||||
|
||||
for i in 0..count {
|
||||
let mut b = test_row(&format!("b{i}"));
|
||||
b.id = format!("backend-{i}");
|
||||
crate::admin::db::backends::insert_managed_backend(conn, &b).unwrap();
|
||||
add_route_provider(conn, &route.id, &b.id, &["*".to_string()], i as i32, true).unwrap();
|
||||
}
|
||||
|
||||
let ids: Vec<String> = list_route_providers(conn, &route.id)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|p| p.id)
|
||||
.collect();
|
||||
(route.id, ids)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_route_sets_and_clears_option_override() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, _) = seed_route_with_providers(&conn, 1);
|
||||
|
||||
// Set an override.
|
||||
let set = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: Some(false),
|
||||
guardrail_mode: Some(Some("standard".into())),
|
||||
pxpipe_compress: Some(Some(true)),
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &set).unwrap());
|
||||
let r = get_route(&conn, &route_id).unwrap().unwrap();
|
||||
assert!(!r.enabled);
|
||||
assert_eq!(r.guardrail_mode.as_deref(), Some("standard"));
|
||||
assert_eq!(r.pxpipe_compress, Some(true));
|
||||
|
||||
// Clear the override back to NULL (inherit).
|
||||
let clear = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: None,
|
||||
guardrail_mode: Some(None),
|
||||
pxpipe_compress: Some(None),
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &clear).unwrap());
|
||||
let r = get_route(&conn, &route_id).unwrap().unwrap();
|
||||
assert_eq!(r.guardrail_mode, None);
|
||||
assert_eq!(r.pxpipe_compress, None);
|
||||
// enabled was left unchanged (None) — still false.
|
||||
assert!(!r.enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_route_excluded_from_enabled_route_ids() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, _) = seed_route_with_providers(&conn, 1);
|
||||
// The seeded backend is "b0"; while enabled the route id is returned.
|
||||
assert_eq!(
|
||||
enabled_route_ids_for_backend_name(&conn, "b0").unwrap(),
|
||||
vec![route_id.clone()]
|
||||
);
|
||||
|
||||
// Disable the route -> it drops out of the virtual-key scope query.
|
||||
let patch = RoutePatch {
|
||||
name: None,
|
||||
description: None,
|
||||
strategy: None,
|
||||
rpm: None,
|
||||
tpm: None,
|
||||
budget_usd: None,
|
||||
enabled: Some(false),
|
||||
guardrail_mode: None,
|
||||
pxpipe_compress: None,
|
||||
pxpipe_models: None,
|
||||
redact_secrets: None,
|
||||
position: None,
|
||||
};
|
||||
assert!(update_route(&conn, &route_id, &patch).unwrap());
|
||||
assert!(enabled_route_ids_for_backend_name(&conn, "b0")
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_rewrites_priorities() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Reverse the order.
|
||||
let reversed: Vec<String> = ids.iter().rev().cloned().collect();
|
||||
let outcome = reorder_route_providers(&conn, &route_id, &reversed).unwrap();
|
||||
|
||||
match outcome {
|
||||
ReorderOutcome::Ok(rows) => {
|
||||
let new_order: Vec<String> = rows.iter().map(|p| p.id.clone()).collect();
|
||||
assert_eq!(new_order, reversed);
|
||||
for (i, row) in rows.iter().enumerate() {
|
||||
assert_eq!(row.priority, i as i32);
|
||||
}
|
||||
}
|
||||
ReorderOutcome::Mismatch => panic!("expected Ok"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_mismatch_rolls_back() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Submit a subset — should be rejected as Mismatch, priorities unchanged.
|
||||
let partial: Vec<String> = ids.iter().take(2).cloned().collect();
|
||||
let outcome = reorder_route_providers(&conn, &route_id, &partial).unwrap();
|
||||
assert!(matches!(outcome, ReorderOutcome::Mismatch));
|
||||
|
||||
// Priorities must be untouched.
|
||||
let rows = list_route_providers(&conn, &route_id).unwrap();
|
||||
for (i, row) in rows.iter().enumerate() {
|
||||
assert_eq!(
|
||||
row.priority, i as i32,
|
||||
"priorities must be unchanged after Mismatch"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reorder_route_providers_rejects_duplicates_and_extras() {
|
||||
let conn = in_memory_db();
|
||||
let (route_id, ids) = seed_route_with_providers(&conn, 3);
|
||||
|
||||
// Duplicate id.
|
||||
let dup = vec![ids[0].clone(), ids[0].clone(), ids[1].clone()];
|
||||
assert!(matches!(
|
||||
reorder_route_providers(&conn, &route_id, &dup).unwrap(),
|
||||
ReorderOutcome::Mismatch
|
||||
));
|
||||
|
||||
// Extra id that isn't a provider on this route.
|
||||
let mut extra = ids.clone();
|
||||
extra.push("bogus-id".into());
|
||||
assert!(matches!(
|
||||
reorder_route_providers(&conn, &route_id, &extra).unwrap(),
|
||||
ReorderOutcome::Mismatch
|
||||
));
|
||||
}
|
||||
@@ -7,75 +7,6 @@ use axum::{
|
||||
};
|
||||
use std::net::SocketAddr;
|
||||
|
||||
/// GET /admin/api/env -- effective environment variable values.
|
||||
/// Secrets (API keys, tokens) are masked; plain config values are shown as-is.
|
||||
pub(super) async fn get_env() -> Json<serde_json::Value> {
|
||||
fn plain(key: &str) -> serde_json::Value {
|
||||
match std::env::var(key) {
|
||||
Ok(v) if !v.is_empty() => serde_json::Value::String(v),
|
||||
_ => serde_json::Value::Null,
|
||||
}
|
||||
}
|
||||
fn secret(key: &str) -> serde_json::Value {
|
||||
match std::env::var(key) {
|
||||
Ok(v) if !v.is_empty() => {
|
||||
serde_json::Value::String(anyllm_translate::util::redact::redact_secret(&v))
|
||||
}
|
||||
_ => serde_json::Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
Json(serde_json::json!({
|
||||
// Core proxy config
|
||||
"BACKEND": plain("BACKEND"),
|
||||
"LISTEN_PORT": plain("LISTEN_PORT"),
|
||||
"BIG_MODEL": plain("BIG_MODEL"),
|
||||
"SMALL_MODEL": plain("SMALL_MODEL"),
|
||||
"RUST_LOG": plain("RUST_LOG"),
|
||||
"LOG_BODIES": plain("LOG_BODIES"),
|
||||
"REDACT_SECRETS": plain("REDACT_SECRETS"),
|
||||
"PROXY_CONFIG": plain("PROXY_CONFIG"),
|
||||
// OpenAI / compatible
|
||||
"OPENAI_BASE_URL": plain("OPENAI_BASE_URL"),
|
||||
"OPENAI_API_FORMAT": plain("OPENAI_API_FORMAT"),
|
||||
"OPENAI_API_KEY": secret("OPENAI_API_KEY"),
|
||||
// Vertex AI
|
||||
"VERTEX_PROJECT": plain("VERTEX_PROJECT"),
|
||||
"VERTEX_REGION": plain("VERTEX_REGION"),
|
||||
"VERTEX_API_KEY": secret("VERTEX_API_KEY"),
|
||||
// Gemini
|
||||
"GEMINI_BASE_URL": plain("GEMINI_BASE_URL"),
|
||||
"GEMINI_API_KEY": secret("GEMINI_API_KEY"),
|
||||
// Azure OpenAI
|
||||
"AZURE_OPENAI_ENDPOINT": plain("AZURE_OPENAI_ENDPOINT"),
|
||||
"AZURE_OPENAI_DEPLOYMENT": plain("AZURE_OPENAI_DEPLOYMENT"),
|
||||
"AZURE_OPENAI_API_KEY": secret("AZURE_OPENAI_API_KEY"),
|
||||
"AZURE_OPENAI_API_VERSION": plain("AZURE_OPENAI_API_VERSION"),
|
||||
// AWS Bedrock
|
||||
"AWS_REGION": plain("AWS_REGION"),
|
||||
"AWS_ACCESS_KEY_ID": secret("AWS_ACCESS_KEY_ID"),
|
||||
"AWS_SECRET_ACCESS_KEY": secret("AWS_SECRET_ACCESS_KEY"),
|
||||
"AWS_SESSION_TOKEN": secret("AWS_SESSION_TOKEN"),
|
||||
// Google OAuth bearer token (full token — treat as secret)
|
||||
"GOOGLE_ACCESS_TOKEN": secret("GOOGLE_ACCESS_TOKEN"),
|
||||
// Auth
|
||||
"PROXY_API_KEYS": secret("PROXY_API_KEYS"),
|
||||
"PROXY_OPEN_RELAY": plain("PROXY_OPEN_RELAY"),
|
||||
// TLS
|
||||
"TLS_CLIENT_CERT_P12": plain("TLS_CLIENT_CERT_P12"),
|
||||
"TLS_CA_CERT": plain("TLS_CA_CERT"),
|
||||
// Network / security
|
||||
"IP_ALLOWLIST": plain("IP_ALLOWLIST"),
|
||||
"TRUST_PROXY_HEADERS": plain("TRUST_PROXY_HEADERS"),
|
||||
"WEBHOOK_URLS": plain("WEBHOOK_URLS"),
|
||||
"RATE_LIMIT_FAIL_POLICY": plain("RATE_LIMIT_FAIL_POLICY"),
|
||||
// Admin
|
||||
"ADMIN_PORT": plain("ADMIN_PORT"),
|
||||
"ADMIN_DB_PATH": plain("ADMIN_DB_PATH"),
|
||||
"ADMIN_LOG_RETENTION_DAYS": plain("ADMIN_LOG_RETENTION_DAYS"),
|
||||
}))
|
||||
}
|
||||
|
||||
/// GET /admin/api/config -- effective config (env defaults + overrides).
|
||||
pub(super) async fn get_config(State(shared): State<SharedState>) -> Json<serde_json::Value> {
|
||||
// Clone config snapshot and drop the read guard before any .await points.
|
||||
|
||||
@@ -17,6 +17,75 @@ use axum::{
|
||||
use serde::Serialize;
|
||||
use std::net::SocketAddr;
|
||||
|
||||
/// GET /admin/api/env -- effective environment variable values.
|
||||
/// Secrets (API keys, tokens) are masked; plain config values are shown as-is.
|
||||
pub(super) async fn get_env() -> Json<serde_json::Value> {
|
||||
fn plain(key: &str) -> serde_json::Value {
|
||||
match std::env::var(key) {
|
||||
Ok(v) if !v.is_empty() => serde_json::Value::String(v),
|
||||
_ => serde_json::Value::Null,
|
||||
}
|
||||
}
|
||||
fn secret(key: &str) -> serde_json::Value {
|
||||
match std::env::var(key) {
|
||||
Ok(v) if !v.is_empty() => {
|
||||
serde_json::Value::String(anyllm_translate::util::redact::redact_secret(&v))
|
||||
}
|
||||
_ => serde_json::Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
Json(serde_json::json!({
|
||||
// Core proxy config
|
||||
"BACKEND": plain("BACKEND"),
|
||||
"LISTEN_PORT": plain("LISTEN_PORT"),
|
||||
"BIG_MODEL": plain("BIG_MODEL"),
|
||||
"SMALL_MODEL": plain("SMALL_MODEL"),
|
||||
"RUST_LOG": plain("RUST_LOG"),
|
||||
"LOG_BODIES": plain("LOG_BODIES"),
|
||||
"REDACT_SECRETS": plain("REDACT_SECRETS"),
|
||||
"PROXY_CONFIG": plain("PROXY_CONFIG"),
|
||||
// OpenAI / compatible
|
||||
"OPENAI_BASE_URL": plain("OPENAI_BASE_URL"),
|
||||
"OPENAI_API_FORMAT": plain("OPENAI_API_FORMAT"),
|
||||
"OPENAI_API_KEY": secret("OPENAI_API_KEY"),
|
||||
// Vertex AI
|
||||
"VERTEX_PROJECT": plain("VERTEX_PROJECT"),
|
||||
"VERTEX_REGION": plain("VERTEX_REGION"),
|
||||
"VERTEX_API_KEY": secret("VERTEX_API_KEY"),
|
||||
// Gemini
|
||||
"GEMINI_BASE_URL": plain("GEMINI_BASE_URL"),
|
||||
"GEMINI_API_KEY": secret("GEMINI_API_KEY"),
|
||||
// Azure OpenAI
|
||||
"AZURE_OPENAI_ENDPOINT": plain("AZURE_OPENAI_ENDPOINT"),
|
||||
"AZURE_OPENAI_DEPLOYMENT": plain("AZURE_OPENAI_DEPLOYMENT"),
|
||||
"AZURE_OPENAI_API_KEY": secret("AZURE_OPENAI_API_KEY"),
|
||||
"AZURE_OPENAI_API_VERSION": plain("AZURE_OPENAI_API_VERSION"),
|
||||
// AWS Bedrock
|
||||
"AWS_REGION": plain("AWS_REGION"),
|
||||
"AWS_ACCESS_KEY_ID": secret("AWS_ACCESS_KEY_ID"),
|
||||
"AWS_SECRET_ACCESS_KEY": secret("AWS_SECRET_ACCESS_KEY"),
|
||||
"AWS_SESSION_TOKEN": secret("AWS_SESSION_TOKEN"),
|
||||
// Google OAuth bearer token (full token — treat as secret)
|
||||
"GOOGLE_ACCESS_TOKEN": secret("GOOGLE_ACCESS_TOKEN"),
|
||||
// Auth
|
||||
"PROXY_API_KEYS": secret("PROXY_API_KEYS"),
|
||||
"PROXY_OPEN_RELAY": plain("PROXY_OPEN_RELAY"),
|
||||
// TLS
|
||||
"TLS_CLIENT_CERT_P12": plain("TLS_CLIENT_CERT_P12"),
|
||||
"TLS_CA_CERT": plain("TLS_CA_CERT"),
|
||||
// Network / security
|
||||
"IP_ALLOWLIST": plain("IP_ALLOWLIST"),
|
||||
"TRUST_PROXY_HEADERS": plain("TRUST_PROXY_HEADERS"),
|
||||
"WEBHOOK_URLS": plain("WEBHOOK_URLS"),
|
||||
"RATE_LIMIT_FAIL_POLICY": plain("RATE_LIMIT_FAIL_POLICY"),
|
||||
// Admin
|
||||
"ADMIN_PORT": plain("ADMIN_PORT"),
|
||||
"ADMIN_DB_PATH": plain("ADMIN_DB_PATH"),
|
||||
"ADMIN_LOG_RETENTION_DAYS": plain("ADMIN_LOG_RETENTION_DAYS"),
|
||||
}))
|
||||
}
|
||||
|
||||
/// Maximum accepted upload size for an env file. Env files are tiny;
|
||||
/// anything larger is almost certainly the wrong file type.
|
||||
/// The router already enforces a 1 MB body limit; this is a tighter application-level check.
|
||||
|
||||
@@ -83,7 +83,7 @@ pub fn admin_router(shared: SharedState, token: Arc<zeroize::Zeroizing<String>>)
|
||||
"/admin/api/config/overrides/{key}",
|
||||
delete(config::delete_config_override),
|
||||
)
|
||||
.route("/admin/api/env", get(config::get_env))
|
||||
.route("/admin/api/env", get(env::get_env))
|
||||
.route("/admin/api/env/import", post(env::import_env))
|
||||
.route("/admin/api/env/export", get(env::export_env))
|
||||
.route("/admin/api/metrics", get(logs::get_metrics))
|
||||
|
||||
@@ -495,103 +495,5 @@ fn ensure_local_discover_host(parsed: &url::Url) -> Result<(), String> {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn custom_request(url: &str) -> DiscoverRequest {
|
||||
DiscoverRequest {
|
||||
source: "custom".to_string(),
|
||||
url: Some(url.to_string()),
|
||||
provider_id: None,
|
||||
api_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_rejects_loopback_url() {
|
||||
let err = resolve_discover_target(&custom_request("http://127.0.0.1:11434"), false)
|
||||
.expect_err("loopback custom discovery URL must be rejected");
|
||||
|
||||
assert!(err.contains("private/loopback"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_rejects_file_scheme() {
|
||||
let err = resolve_discover_target(&custom_request("file:///tmp/anyllm"), false)
|
||||
.expect_err("non-http custom discovery URL must be rejected");
|
||||
|
||||
assert!(err.contains("scheme"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_accepts_https_url_and_appends_models() {
|
||||
let (url, api_key) = resolve_discover_target(&custom_request("https://1.1.1.1"), false)
|
||||
.expect("public HTTPS custom discovery URL should be accepted");
|
||||
|
||||
assert_eq!(url, "https://1.1.1.1/v1/models");
|
||||
assert_eq!(api_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_accepts_existing_models_suffix() {
|
||||
let (url, api_key) =
|
||||
resolve_discover_target(&custom_request("https://1.1.1.1/v1/models"), false)
|
||||
.expect("public HTTPS custom discovery URL should be accepted");
|
||||
|
||||
assert_eq!(url, "https://1.1.1.1/v1/models");
|
||||
assert_eq!(api_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_does_not_double_v1_suffix() {
|
||||
// Local-provider catalog defaults already end in /v1; don't produce /v1/v1/models.
|
||||
let (url, _) =
|
||||
resolve_discover_target(&custom_request("http://192.168.1.72:4444/v1"), true)
|
||||
.expect("local LAN /v1 discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://192.168.1.72:4444/v1/models");
|
||||
|
||||
// Trailing slash on a /v1 base collapses the same way.
|
||||
let (url, _) =
|
||||
resolve_discover_target(&custom_request("http://192.168.1.72:4444/v1/"), true)
|
||||
.expect("trailing-slash /v1 discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://192.168.1.72:4444/v1/models");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_local_allows_loopback_but_keeps_scheme_check() {
|
||||
// allow_local=true: loopback is accepted (LM Studio/Ollama on localhost)...
|
||||
let (url, _) = resolve_discover_target(&custom_request("http://127.0.0.1:1234"), true)
|
||||
.expect("local loopback discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://127.0.0.1:1234/v1/models");
|
||||
|
||||
// ...but a non-http scheme is still rejected.
|
||||
let err = resolve_discover_target(&custom_request("file:///tmp/anyllm"), true)
|
||||
.expect_err("non-http scheme must be rejected even when allow_local");
|
||||
assert!(err.contains("scheme"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_threads_api_key() {
|
||||
let mut req = custom_request("http://192.168.1.72:4444");
|
||||
req.api_key = Some("sk-local".to_string());
|
||||
let (_, api_key) = resolve_discover_target(&req, true)
|
||||
.expect("local LAN discovery URL should be accepted when allow_local");
|
||||
assert_eq!(api_key.as_deref(), Some("sk-local"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_local_rejects_public_and_metadata_hosts() {
|
||||
// allow_local must NOT permit public hosts (would be a general outbound fetch).
|
||||
let err = resolve_discover_target(&custom_request("http://1.1.1.1/v1"), true)
|
||||
.expect_err("public IP must be rejected even when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
// Cloud-metadata / link-local stay blocked here too.
|
||||
let err = resolve_discover_target(&custom_request("http://169.254.169.254/"), true)
|
||||
.expect_err("metadata IP must be rejected even when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
// A bare host name that isn't `localhost` cannot be verified as local.
|
||||
let err = resolve_discover_target(&custom_request("http://evil.example.com/"), true)
|
||||
.expect_err("non-local host name must be rejected when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
}
|
||||
}
|
||||
#[path = "models/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
use super::*;
|
||||
|
||||
fn custom_request(url: &str) -> DiscoverRequest {
|
||||
DiscoverRequest {
|
||||
source: "custom".to_string(),
|
||||
url: Some(url.to_string()),
|
||||
provider_id: None,
|
||||
api_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_rejects_loopback_url() {
|
||||
let err = resolve_discover_target(&custom_request("http://127.0.0.1:11434"), false)
|
||||
.expect_err("loopback custom discovery URL must be rejected");
|
||||
|
||||
assert!(err.contains("private/loopback"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_rejects_file_scheme() {
|
||||
let err = resolve_discover_target(&custom_request("file:///tmp/anyllm"), false)
|
||||
.expect_err("non-http custom discovery URL must be rejected");
|
||||
|
||||
assert!(err.contains("scheme"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_accepts_https_url_and_appends_models() {
|
||||
let (url, api_key) = resolve_discover_target(&custom_request("https://1.1.1.1"), false)
|
||||
.expect("public HTTPS custom discovery URL should be accepted");
|
||||
|
||||
assert_eq!(url, "https://1.1.1.1/v1/models");
|
||||
assert_eq!(api_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_accepts_existing_models_suffix() {
|
||||
let (url, api_key) =
|
||||
resolve_discover_target(&custom_request("https://1.1.1.1/v1/models"), false)
|
||||
.expect("public HTTPS custom discovery URL should be accepted");
|
||||
|
||||
assert_eq!(url, "https://1.1.1.1/v1/models");
|
||||
assert_eq!(api_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_does_not_double_v1_suffix() {
|
||||
// Local-provider catalog defaults already end in /v1; don't produce /v1/v1/models.
|
||||
let (url, _) = resolve_discover_target(&custom_request("http://192.168.1.72:4444/v1"), true)
|
||||
.expect("local LAN /v1 discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://192.168.1.72:4444/v1/models");
|
||||
|
||||
// Trailing slash on a /v1 base collapses the same way.
|
||||
let (url, _) = resolve_discover_target(&custom_request("http://192.168.1.72:4444/v1/"), true)
|
||||
.expect("trailing-slash /v1 discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://192.168.1.72:4444/v1/models");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_local_allows_loopback_but_keeps_scheme_check() {
|
||||
// allow_local=true: loopback is accepted (LM Studio/Ollama on localhost)...
|
||||
let (url, _) = resolve_discover_target(&custom_request("http://127.0.0.1:1234"), true)
|
||||
.expect("local loopback discovery URL should be accepted when allow_local");
|
||||
assert_eq!(url, "http://127.0.0.1:1234/v1/models");
|
||||
|
||||
// ...but a non-http scheme is still rejected.
|
||||
let err = resolve_discover_target(&custom_request("file:///tmp/anyllm"), true)
|
||||
.expect_err("non-http scheme must be rejected even when allow_local");
|
||||
assert!(err.contains("scheme"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_threads_api_key() {
|
||||
let mut req = custom_request("http://192.168.1.72:4444");
|
||||
req.api_key = Some("sk-local".to_string());
|
||||
let (_, api_key) = resolve_discover_target(&req, true)
|
||||
.expect("local LAN discovery URL should be accepted when allow_local");
|
||||
assert_eq!(api_key.as_deref(), Some("sk-local"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_discover_local_rejects_public_and_metadata_hosts() {
|
||||
// allow_local must NOT permit public hosts (would be a general outbound fetch).
|
||||
let err = resolve_discover_target(&custom_request("http://1.1.1.1/v1"), true)
|
||||
.expect_err("public IP must be rejected even when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
// Cloud-metadata / link-local stay blocked here too.
|
||||
let err = resolve_discover_target(&custom_request("http://169.254.169.254/"), true)
|
||||
.expect_err("metadata IP must be rejected even when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
// A bare host name that isn't `localhost` cannot be verified as local.
|
||||
let err = resolve_discover_target(&custom_request("http://evil.example.com/"), true)
|
||||
.expect_err("non-local host name must be rejected when allow_local");
|
||||
assert!(err.contains("loopback/LAN"));
|
||||
}
|
||||
Vendored
+42
@@ -0,0 +1,42 @@
|
||||
/// Configuration for the cache subsystem.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CacheConfig {
|
||||
/// Default TTL in seconds for cached responses.
|
||||
pub ttl_secs: u64,
|
||||
/// Maximum number of entries in the in-memory cache.
|
||||
pub max_entries: u64,
|
||||
/// Optional Redis URL. Used by the Redis L2 cache backend (requires `redis` feature)
|
||||
/// and distributed rate limiting. When set, responses are cached in Redis in addition
|
||||
/// to the in-memory L1 cache, and rate limit state is shared across proxy instances.
|
||||
pub redis_url: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for CacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ttl_secs: 300,
|
||||
max_entries: 10_000,
|
||||
redis_url: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CacheConfig {
|
||||
/// Load cache configuration from environment variables.
|
||||
pub fn from_env() -> Self {
|
||||
let ttl_secs = std::env::var("CACHE_TTL_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(300);
|
||||
let max_entries = std::env::var("CACHE_MAX_ENTRIES")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(10_000);
|
||||
let redis_url = std::env::var("REDIS_URL").ok();
|
||||
Self {
|
||||
ttl_secs,
|
||||
max_entries,
|
||||
redis_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+189
@@ -0,0 +1,189 @@
|
||||
use super::CacheEntry;
|
||||
use super::MAX_TTL_SECS;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Per-request cache controls after combining local and LiteLLM-style fields.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct CacheControl {
|
||||
/// Whether to read an existing cached response before calling upstream.
|
||||
pub lookup: bool,
|
||||
/// Whether to store a successful upstream response after the call.
|
||||
pub store: bool,
|
||||
/// Per-entry store TTL override in seconds.
|
||||
pub ttl_secs: Option<u64>,
|
||||
/// Maximum acceptable age for an existing cached entry.
|
||||
pub max_age_secs: Option<u64>,
|
||||
/// Optional caller namespace for exact-match cache isolation.
|
||||
pub namespace: Option<String>,
|
||||
/// Parsed for LiteLLM compatibility. Currently no behavior change.
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl Default for CacheControl {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
lookup: true,
|
||||
store: true,
|
||||
ttl_secs: None,
|
||||
max_age_secs: None,
|
||||
namespace: None,
|
||||
use_cache: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the optional `cache_ttl_secs` field from a request body.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `Ok(None)` if the field is absent or null (use default TTL).
|
||||
/// - `Ok(Some(0))` if explicitly 0 (bypass cache).
|
||||
/// - `Ok(Some(n))` for valid positive values up to MAX_TTL_SECS.
|
||||
/// - `Err(message)` for negative values, values > MAX_TTL_SECS, or non-numeric.
|
||||
pub fn parse_cache_ttl(body: &serde_json::Value) -> Result<Option<u64>, String> {
|
||||
let Some(val) = body.get("cache_ttl_secs") else {
|
||||
return Ok(None);
|
||||
};
|
||||
if val.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(n) = val.as_u64() {
|
||||
if n > MAX_TTL_SECS {
|
||||
return Err(format!("cache_ttl_secs must be <= {MAX_TTL_SECS}, got {n}"));
|
||||
}
|
||||
return Ok(Some(n));
|
||||
}
|
||||
if let Some(n) = val.as_i64() {
|
||||
// Negative values are invalid
|
||||
return Err(format!("cache_ttl_secs must be non-negative, got {n}"));
|
||||
}
|
||||
if let Some(n) = val.as_f64() {
|
||||
if n < 0.0 {
|
||||
return Err(format!("cache_ttl_secs must be non-negative, got {n}"));
|
||||
}
|
||||
let truncated = n as u64;
|
||||
if truncated > MAX_TTL_SECS {
|
||||
return Err(format!(
|
||||
"cache_ttl_secs must be <= {MAX_TTL_SECS}, got {truncated}"
|
||||
));
|
||||
}
|
||||
return Ok(Some(truncated));
|
||||
}
|
||||
Err(format!("cache_ttl_secs must be a number, got {}", val))
|
||||
}
|
||||
|
||||
/// Parse local `cache_ttl_secs` and LiteLLM-compatible top-level `cache`.
|
||||
pub fn parse_cache_control(body: &serde_json::Value) -> Result<CacheControl, String> {
|
||||
let mut control = CacheControl::default();
|
||||
match parse_cache_ttl(body)? {
|
||||
Some(0) => {
|
||||
control.lookup = false;
|
||||
control.store = false;
|
||||
control.ttl_secs = Some(0);
|
||||
}
|
||||
Some(ttl) => {
|
||||
control.ttl_secs = Some(ttl);
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
let Some(cache_value) = body.get("cache") else {
|
||||
return Ok(control);
|
||||
};
|
||||
if cache_value.is_null() {
|
||||
return Ok(control);
|
||||
}
|
||||
let Some(cache_obj) = cache_value.as_object() else {
|
||||
return Err(format!("cache must be an object, got {}", cache_value));
|
||||
};
|
||||
|
||||
if parse_cache_bool_field(cache_obj, "no-cache")?.unwrap_or(false) {
|
||||
control.lookup = false;
|
||||
}
|
||||
if parse_cache_bool_field(cache_obj, "no-store")?.unwrap_or(false) {
|
||||
control.store = false;
|
||||
}
|
||||
if let Some(use_cache) = parse_cache_bool_field(cache_obj, "use-cache")? {
|
||||
control.use_cache = use_cache;
|
||||
}
|
||||
if let Some(ttl) = parse_cache_secs_field(cache_obj, "ttl")? {
|
||||
control.ttl_secs = Some(ttl);
|
||||
}
|
||||
if let Some(max_age) = parse_cache_secs_field(cache_obj, "s-maxage")? {
|
||||
control.max_age_secs = Some(max_age);
|
||||
}
|
||||
if let Some(max_age) = parse_cache_secs_field(cache_obj, "s-max-age")? {
|
||||
control.max_age_secs = Some(max_age);
|
||||
}
|
||||
if let Some(namespace_value) = cache_obj.get("namespace") {
|
||||
if !namespace_value.is_null() {
|
||||
let Some(namespace) = namespace_value.as_str() else {
|
||||
return Err(format!(
|
||||
"cache.namespace must be a string, got {}",
|
||||
namespace_value
|
||||
));
|
||||
};
|
||||
if !namespace.is_empty() {
|
||||
control.namespace = Some(namespace.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(control)
|
||||
}
|
||||
|
||||
fn parse_cache_bool_field(
|
||||
cache_obj: &serde_json::Map<String, serde_json::Value>,
|
||||
field: &str,
|
||||
) -> Result<Option<bool>, String> {
|
||||
let Some(value) = cache_obj.get(field) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
value
|
||||
.as_bool()
|
||||
.map(Some)
|
||||
.ok_or_else(|| format!("cache.{field} must be a boolean, got {value}"))
|
||||
}
|
||||
|
||||
fn parse_cache_secs_field(
|
||||
cache_obj: &serde_json::Map<String, serde_json::Value>,
|
||||
field: &str,
|
||||
) -> Result<Option<u64>, String> {
|
||||
let Some(value) = cache_obj.get(field) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(n) = value.as_u64() {
|
||||
if n > MAX_TTL_SECS {
|
||||
return Err(format!("cache.{field} must be <= {MAX_TTL_SECS}, got {n}"));
|
||||
}
|
||||
return Ok(Some(n));
|
||||
}
|
||||
if let Some(n) = value.as_i64() {
|
||||
return Err(format!("cache.{field} must be non-negative, got {n}"));
|
||||
}
|
||||
if let Some(n) = value.as_f64() {
|
||||
if n < 0.0 {
|
||||
return Err(format!("cache.{field} must be non-negative, got {n}"));
|
||||
}
|
||||
let truncated = n as u64;
|
||||
if truncated > MAX_TTL_SECS {
|
||||
return Err(format!(
|
||||
"cache.{field} must be <= {MAX_TTL_SECS}, got {truncated}"
|
||||
));
|
||||
}
|
||||
return Ok(Some(truncated));
|
||||
}
|
||||
Err(format!("cache.{field} must be a number, got {}", value))
|
||||
}
|
||||
|
||||
pub fn cache_entry_is_fresh(entry: &CacheEntry, max_age_secs: Option<u64>) -> bool {
|
||||
match max_age_secs {
|
||||
Some(max_age) => entry.created_at.elapsed() <= Duration::from_secs(max_age),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
Vendored
+148
@@ -0,0 +1,148 @@
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::io::{self, Write};
|
||||
|
||||
/// Namespace prefix for cache keys, preventing cross-endpoint collisions.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum CacheNamespace {
|
||||
/// Anthropic /v1/messages endpoint.
|
||||
Anthropic,
|
||||
/// OpenAI /v1/chat/completions endpoint.
|
||||
OpenAI,
|
||||
}
|
||||
|
||||
impl CacheNamespace {
|
||||
fn prefix(self) -> &'static str {
|
||||
match self {
|
||||
Self::Anthropic => "anth",
|
||||
Self::OpenAI => "oai",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CacheScope<'a> {
|
||||
pub backend_name: &'a str,
|
||||
pub auth_identity: &'a str,
|
||||
pub namespace: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// Compute a deterministic cache key for a request body.
|
||||
///
|
||||
/// Extracts the fields that affect response content, sorts them via BTreeMap,
|
||||
/// serializes to canonical JSON, SHA-256 hashes the result, and prepends the
|
||||
/// namespace prefix.
|
||||
pub fn cache_key_for_request(
|
||||
body: &serde_json::Value,
|
||||
ns: CacheNamespace,
|
||||
scope: &CacheScope<'_>,
|
||||
) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
write_canonical_cache_body(&mut hasher, body, scope);
|
||||
let hash = hasher.finalize();
|
||||
let hex = hex::encode(hash);
|
||||
format!("{}:{}", ns.prefix(), hex)
|
||||
}
|
||||
|
||||
enum CacheField<'a> {
|
||||
Json(&'a str, &'a serde_json::Value),
|
||||
Str(&'static str, &'a str),
|
||||
}
|
||||
|
||||
impl<'a> CacheField<'a> {
|
||||
fn key(&self) -> &str {
|
||||
match self {
|
||||
Self::Json(key, _) | Self::Str(key, _) => key,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct HashWriter<'a> {
|
||||
hasher: &'a mut Sha256,
|
||||
}
|
||||
|
||||
impl Write for HashWriter<'_> {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.hasher.update(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn write_canonical_cache_body(
|
||||
hasher: &mut Sha256,
|
||||
body: &serde_json::Value,
|
||||
scope: &CacheScope<'_>,
|
||||
) {
|
||||
let mut fields = Vec::new();
|
||||
if let Some(obj) = body.as_object() {
|
||||
fields.extend(
|
||||
obj.iter()
|
||||
.filter(|(key, value)| should_include_cache_field(key, value))
|
||||
.map(|(key, value)| CacheField::Json(key.as_str(), value)),
|
||||
);
|
||||
}
|
||||
fields.push(CacheField::Str("_scope_auth", scope.auth_identity));
|
||||
fields.push(CacheField::Str("_scope_backend", scope.backend_name));
|
||||
if let Some(namespace) = scope.namespace {
|
||||
fields.push(CacheField::Str("_scope_cache_namespace", namespace));
|
||||
}
|
||||
fields.sort_unstable_by(|a, b| a.key().cmp(b.key()));
|
||||
|
||||
let mut writer = HashWriter { hasher };
|
||||
writer
|
||||
.write_all(b"{")
|
||||
.expect("hash writer should not fail writing object start");
|
||||
for (idx, field) in fields.iter().enumerate() {
|
||||
if idx > 0 {
|
||||
writer
|
||||
.write_all(b",")
|
||||
.expect("hash writer should not fail writing separator");
|
||||
}
|
||||
serde_json::to_writer(&mut writer, field.key())
|
||||
.expect("hash writer should not fail writing key");
|
||||
writer
|
||||
.write_all(b":")
|
||||
.expect("hash writer should not fail writing colon");
|
||||
match field {
|
||||
CacheField::Json(_, value) => serde_json::to_writer(&mut writer, value)
|
||||
.expect("hash writer should not fail writing JSON value"),
|
||||
CacheField::Str(_, value) => serde_json::to_writer(&mut writer, value)
|
||||
.expect("hash writer should not fail writing string value"),
|
||||
}
|
||||
}
|
||||
writer
|
||||
.write_all(b"}")
|
||||
.expect("hash writer should not fail writing object end");
|
||||
}
|
||||
|
||||
fn should_include_cache_field(key: &str, value: &serde_json::Value) -> bool {
|
||||
if value.is_null() {
|
||||
return false;
|
||||
}
|
||||
// Exclude fields that do not affect the backend response:
|
||||
// - stream / stream_options: transport only (a cached non-stream response is
|
||||
// replayed as a stream and vice versa).
|
||||
// - cache: request cache controls, handled outside response-content hashing.
|
||||
// - _scope_*: added separately as scope fields.
|
||||
// - user / metadata: tracking fields documented as "Ignored" (anyllm_translate
|
||||
// ChatCompletionRequest user, anthropic Metadata). Hashing them fragments the
|
||||
// cache per end-user with no correctness benefit (tenant isolation is already
|
||||
// provided by _scope_auth).
|
||||
//
|
||||
// parallel_tool_calls is NOT excluded: backends that honor it (e.g. OpenAI)
|
||||
// produce different output for true vs false, so it must be part of the cache
|
||||
// identity. (Gemini/Vertex have it stripped before dispatch by the tool policy.)
|
||||
!matches!(
|
||||
key,
|
||||
"stream"
|
||||
| "stream_options"
|
||||
| "cache"
|
||||
| "_scope_auth"
|
||||
| "_scope_backend"
|
||||
| "_scope_cache_namespace"
|
||||
| "user"
|
||||
| "metadata"
|
||||
)
|
||||
}
|
||||
Vendored
+12
-881
@@ -15,14 +15,23 @@ pub mod redis;
|
||||
#[cfg(feature = "qdrant")]
|
||||
pub mod semantic;
|
||||
|
||||
mod config;
|
||||
mod control;
|
||||
mod key;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use bytes::Bytes;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::io::{self, Write};
|
||||
use std::time::{Duration, Instant};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Maximum allowed value for per-request `cache_ttl_secs`.
|
||||
pub const MAX_TTL_SECS: u64 = 86_400;
|
||||
|
||||
pub use config::CacheConfig;
|
||||
pub use control::{cache_entry_is_fresh, parse_cache_control, parse_cache_ttl, CacheControl};
|
||||
pub use key::{cache_key_for_request, CacheNamespace, CacheScope};
|
||||
|
||||
/// Cached response entry stored in any cache backend.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct CacheEntry {
|
||||
@@ -37,24 +46,6 @@ pub struct CacheEntry {
|
||||
pub ttl_secs: Option<u64>,
|
||||
}
|
||||
|
||||
/// Namespace prefix for cache keys, preventing cross-endpoint collisions.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum CacheNamespace {
|
||||
/// Anthropic /v1/messages endpoint.
|
||||
Anthropic,
|
||||
/// OpenAI /v1/chat/completions endpoint.
|
||||
OpenAI,
|
||||
}
|
||||
|
||||
impl CacheNamespace {
|
||||
fn prefix(self) -> &'static str {
|
||||
match self {
|
||||
Self::Anthropic => "anth",
|
||||
Self::OpenAI => "oai",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pluggable cache backend trait. Implementations must be Send + Sync
|
||||
/// for use behind Arc in axum handlers.
|
||||
pub trait CacheBackend: Send + Sync {
|
||||
@@ -69,863 +60,3 @@ pub trait CacheBackend: Send + Sync {
|
||||
ttl_secs: u64,
|
||||
) -> impl std::future::Future<Output = ()> + Send;
|
||||
}
|
||||
|
||||
/// Compute a deterministic cache key for a request body.
|
||||
///
|
||||
/// Extracts the fields that affect response content, sorts them via BTreeMap,
|
||||
/// serializes to canonical JSON, SHA-256 hashes the result, and prepends the
|
||||
/// namespace prefix.
|
||||
pub fn cache_key_for_request(
|
||||
body: &serde_json::Value,
|
||||
ns: CacheNamespace,
|
||||
scope: &CacheScope<'_>,
|
||||
) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
write_canonical_cache_body(&mut hasher, body, scope);
|
||||
let hash = hasher.finalize();
|
||||
let hex = hex::encode(hash);
|
||||
format!("{}:{}", ns.prefix(), hex)
|
||||
}
|
||||
|
||||
enum CacheField<'a> {
|
||||
Json(&'a str, &'a serde_json::Value),
|
||||
Str(&'static str, &'a str),
|
||||
}
|
||||
|
||||
impl<'a> CacheField<'a> {
|
||||
fn key(&self) -> &str {
|
||||
match self {
|
||||
Self::Json(key, _) | Self::Str(key, _) => key,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct HashWriter<'a> {
|
||||
hasher: &'a mut Sha256,
|
||||
}
|
||||
|
||||
impl Write for HashWriter<'_> {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.hasher.update(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn write_canonical_cache_body(
|
||||
hasher: &mut Sha256,
|
||||
body: &serde_json::Value,
|
||||
scope: &CacheScope<'_>,
|
||||
) {
|
||||
let mut fields = Vec::new();
|
||||
if let Some(obj) = body.as_object() {
|
||||
fields.extend(
|
||||
obj.iter()
|
||||
.filter(|(key, value)| should_include_cache_field(key, value))
|
||||
.map(|(key, value)| CacheField::Json(key.as_str(), value)),
|
||||
);
|
||||
}
|
||||
fields.push(CacheField::Str("_scope_auth", scope.auth_identity));
|
||||
fields.push(CacheField::Str("_scope_backend", scope.backend_name));
|
||||
if let Some(namespace) = scope.namespace {
|
||||
fields.push(CacheField::Str("_scope_cache_namespace", namespace));
|
||||
}
|
||||
fields.sort_unstable_by(|a, b| a.key().cmp(b.key()));
|
||||
|
||||
let mut writer = HashWriter { hasher };
|
||||
writer
|
||||
.write_all(b"{")
|
||||
.expect("hash writer should not fail writing object start");
|
||||
for (idx, field) in fields.iter().enumerate() {
|
||||
if idx > 0 {
|
||||
writer
|
||||
.write_all(b",")
|
||||
.expect("hash writer should not fail writing separator");
|
||||
}
|
||||
serde_json::to_writer(&mut writer, field.key())
|
||||
.expect("hash writer should not fail writing key");
|
||||
writer
|
||||
.write_all(b":")
|
||||
.expect("hash writer should not fail writing colon");
|
||||
match field {
|
||||
CacheField::Json(_, value) => serde_json::to_writer(&mut writer, value)
|
||||
.expect("hash writer should not fail writing JSON value"),
|
||||
CacheField::Str(_, value) => serde_json::to_writer(&mut writer, value)
|
||||
.expect("hash writer should not fail writing string value"),
|
||||
}
|
||||
}
|
||||
writer
|
||||
.write_all(b"}")
|
||||
.expect("hash writer should not fail writing object end");
|
||||
}
|
||||
|
||||
fn should_include_cache_field(key: &str, value: &serde_json::Value) -> bool {
|
||||
if value.is_null() {
|
||||
return false;
|
||||
}
|
||||
// Exclude fields that do not affect the backend response:
|
||||
// - stream / stream_options: transport only (a cached non-stream response is
|
||||
// replayed as a stream and vice versa).
|
||||
// - cache: request cache controls, handled outside response-content hashing.
|
||||
// - _scope_*: added separately as scope fields.
|
||||
// - user / metadata: tracking fields documented as "Ignored" (anyllm_translate
|
||||
// ChatCompletionRequest user, anthropic Metadata). Hashing them fragments the
|
||||
// cache per end-user with no correctness benefit (tenant isolation is already
|
||||
// provided by _scope_auth).
|
||||
//
|
||||
// parallel_tool_calls is NOT excluded: backends that honor it (e.g. OpenAI)
|
||||
// produce different output for true vs false, so it must be part of the cache
|
||||
// identity. (Gemini/Vertex have it stripped before dispatch by the tool policy.)
|
||||
!matches!(
|
||||
key,
|
||||
"stream"
|
||||
| "stream_options"
|
||||
| "cache"
|
||||
| "_scope_auth"
|
||||
| "_scope_backend"
|
||||
| "_scope_cache_namespace"
|
||||
| "user"
|
||||
| "metadata"
|
||||
)
|
||||
}
|
||||
|
||||
pub struct CacheScope<'a> {
|
||||
pub backend_name: &'a str,
|
||||
pub auth_identity: &'a str,
|
||||
pub namespace: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// Per-request cache controls after combining local and LiteLLM-style fields.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct CacheControl {
|
||||
/// Whether to read an existing cached response before calling upstream.
|
||||
pub lookup: bool,
|
||||
/// Whether to store a successful upstream response after the call.
|
||||
pub store: bool,
|
||||
/// Per-entry store TTL override in seconds.
|
||||
pub ttl_secs: Option<u64>,
|
||||
/// Maximum acceptable age for an existing cached entry.
|
||||
pub max_age_secs: Option<u64>,
|
||||
/// Optional caller namespace for exact-match cache isolation.
|
||||
pub namespace: Option<String>,
|
||||
/// Parsed for LiteLLM compatibility. Currently no behavior change.
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl Default for CacheControl {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
lookup: true,
|
||||
store: true,
|
||||
ttl_secs: None,
|
||||
max_age_secs: None,
|
||||
namespace: None,
|
||||
use_cache: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the optional `cache_ttl_secs` field from a request body.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `Ok(None)` if the field is absent or null (use default TTL).
|
||||
/// - `Ok(Some(0))` if explicitly 0 (bypass cache).
|
||||
/// - `Ok(Some(n))` for valid positive values up to MAX_TTL_SECS.
|
||||
/// - `Err(message)` for negative values, values > MAX_TTL_SECS, or non-numeric.
|
||||
pub fn parse_cache_ttl(body: &serde_json::Value) -> Result<Option<u64>, String> {
|
||||
let Some(val) = body.get("cache_ttl_secs") else {
|
||||
return Ok(None);
|
||||
};
|
||||
if val.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(n) = val.as_u64() {
|
||||
if n > MAX_TTL_SECS {
|
||||
return Err(format!("cache_ttl_secs must be <= {MAX_TTL_SECS}, got {n}"));
|
||||
}
|
||||
return Ok(Some(n));
|
||||
}
|
||||
if let Some(n) = val.as_i64() {
|
||||
// Negative values are invalid
|
||||
return Err(format!("cache_ttl_secs must be non-negative, got {n}"));
|
||||
}
|
||||
if let Some(n) = val.as_f64() {
|
||||
if n < 0.0 {
|
||||
return Err(format!("cache_ttl_secs must be non-negative, got {n}"));
|
||||
}
|
||||
let truncated = n as u64;
|
||||
if truncated > MAX_TTL_SECS {
|
||||
return Err(format!(
|
||||
"cache_ttl_secs must be <= {MAX_TTL_SECS}, got {truncated}"
|
||||
));
|
||||
}
|
||||
return Ok(Some(truncated));
|
||||
}
|
||||
Err(format!("cache_ttl_secs must be a number, got {}", val))
|
||||
}
|
||||
|
||||
/// Parse local `cache_ttl_secs` and LiteLLM-compatible top-level `cache`.
|
||||
pub fn parse_cache_control(body: &serde_json::Value) -> Result<CacheControl, String> {
|
||||
let mut control = CacheControl::default();
|
||||
match parse_cache_ttl(body)? {
|
||||
Some(0) => {
|
||||
control.lookup = false;
|
||||
control.store = false;
|
||||
control.ttl_secs = Some(0);
|
||||
}
|
||||
Some(ttl) => {
|
||||
control.ttl_secs = Some(ttl);
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
let Some(cache_value) = body.get("cache") else {
|
||||
return Ok(control);
|
||||
};
|
||||
if cache_value.is_null() {
|
||||
return Ok(control);
|
||||
}
|
||||
let Some(cache_obj) = cache_value.as_object() else {
|
||||
return Err(format!("cache must be an object, got {}", cache_value));
|
||||
};
|
||||
|
||||
if parse_cache_bool_field(cache_obj, "no-cache")?.unwrap_or(false) {
|
||||
control.lookup = false;
|
||||
}
|
||||
if parse_cache_bool_field(cache_obj, "no-store")?.unwrap_or(false) {
|
||||
control.store = false;
|
||||
}
|
||||
if let Some(use_cache) = parse_cache_bool_field(cache_obj, "use-cache")? {
|
||||
control.use_cache = use_cache;
|
||||
}
|
||||
if let Some(ttl) = parse_cache_secs_field(cache_obj, "ttl")? {
|
||||
control.ttl_secs = Some(ttl);
|
||||
}
|
||||
if let Some(max_age) = parse_cache_secs_field(cache_obj, "s-maxage")? {
|
||||
control.max_age_secs = Some(max_age);
|
||||
}
|
||||
if let Some(max_age) = parse_cache_secs_field(cache_obj, "s-max-age")? {
|
||||
control.max_age_secs = Some(max_age);
|
||||
}
|
||||
if let Some(namespace_value) = cache_obj.get("namespace") {
|
||||
if !namespace_value.is_null() {
|
||||
let Some(namespace) = namespace_value.as_str() else {
|
||||
return Err(format!(
|
||||
"cache.namespace must be a string, got {}",
|
||||
namespace_value
|
||||
));
|
||||
};
|
||||
if !namespace.is_empty() {
|
||||
control.namespace = Some(namespace.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(control)
|
||||
}
|
||||
|
||||
fn parse_cache_bool_field(
|
||||
cache_obj: &serde_json::Map<String, serde_json::Value>,
|
||||
field: &str,
|
||||
) -> Result<Option<bool>, String> {
|
||||
let Some(value) = cache_obj.get(field) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
value
|
||||
.as_bool()
|
||||
.map(Some)
|
||||
.ok_or_else(|| format!("cache.{field} must be a boolean, got {value}"))
|
||||
}
|
||||
|
||||
fn parse_cache_secs_field(
|
||||
cache_obj: &serde_json::Map<String, serde_json::Value>,
|
||||
field: &str,
|
||||
) -> Result<Option<u64>, String> {
|
||||
let Some(value) = cache_obj.get(field) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(n) = value.as_u64() {
|
||||
if n > MAX_TTL_SECS {
|
||||
return Err(format!("cache.{field} must be <= {MAX_TTL_SECS}, got {n}"));
|
||||
}
|
||||
return Ok(Some(n));
|
||||
}
|
||||
if let Some(n) = value.as_i64() {
|
||||
return Err(format!("cache.{field} must be non-negative, got {n}"));
|
||||
}
|
||||
if let Some(n) = value.as_f64() {
|
||||
if n < 0.0 {
|
||||
return Err(format!("cache.{field} must be non-negative, got {n}"));
|
||||
}
|
||||
let truncated = n as u64;
|
||||
if truncated > MAX_TTL_SECS {
|
||||
return Err(format!(
|
||||
"cache.{field} must be <= {MAX_TTL_SECS}, got {truncated}"
|
||||
));
|
||||
}
|
||||
return Ok(Some(truncated));
|
||||
}
|
||||
Err(format!("cache.{field} must be a number, got {}", value))
|
||||
}
|
||||
|
||||
pub fn cache_entry_is_fresh(entry: &CacheEntry, max_age_secs: Option<u64>) -> bool {
|
||||
match max_age_secs {
|
||||
Some(max_age) => entry.created_at.elapsed() <= Duration::from_secs(max_age),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for the cache subsystem.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CacheConfig {
|
||||
/// Default TTL in seconds for cached responses.
|
||||
pub ttl_secs: u64,
|
||||
/// Maximum number of entries in the in-memory cache.
|
||||
pub max_entries: u64,
|
||||
/// Optional Redis URL. Used by the Redis L2 cache backend (requires `redis` feature)
|
||||
/// and distributed rate limiting. When set, responses are cached in Redis in addition
|
||||
/// to the in-memory L1 cache, and rate limit state is shared across proxy instances.
|
||||
pub redis_url: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for CacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ttl_secs: 300,
|
||||
max_entries: 10_000,
|
||||
redis_url: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CacheConfig {
|
||||
/// Load cache configuration from environment variables.
|
||||
pub fn from_env() -> Self {
|
||||
let ttl_secs = std::env::var("CACHE_TTL_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(300);
|
||||
let max_entries = std::env::var("CACHE_MAX_ENTRIES")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(10_000);
|
||||
let redis_url = std::env::var("REDIS_URL").ok();
|
||||
Self {
|
||||
ttl_secs,
|
||||
max_entries,
|
||||
redis_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cache_key_deterministic_same_fields() {
|
||||
let body = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
assert!(key1.starts_with("anth:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_different_for_different_temperature() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.7
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.9
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_field_order() {
|
||||
// JSON object field order should not affect the key because we
|
||||
// extract into a BTreeMap.
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"temperature": 0.5,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"model": "gpt-4o",
|
||||
"temperature": 0.5
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_non_cache_fields() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"stream": true
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_namespace_differs() {
|
||||
let body = serde_json::json!({
|
||||
"model": "test",
|
||||
"messages": []
|
||||
});
|
||||
let anth = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let oai = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(anth, oai);
|
||||
assert!(anth.starts_with("anth:"));
|
||||
assert!(oai.starts_with("oai:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_null_field_same_as_absent() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [],
|
||||
"temperature": null
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": []
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_absent() {
|
||||
let body = serde_json::json!({"model": "test"});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_null() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": null});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_zero() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 0});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_valid() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 600});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(600));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_max() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 86400});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(86400));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_over_max() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 86401});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_negative() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": -1});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_string() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": "not a number"});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_differs_for_different_cache_ttl_secs() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache_ttl_secs": 60
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache_ttl_secs": 3600
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(
|
||||
key1, key2,
|
||||
"different cache_ttl_secs must produce different cache keys"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_litellm_cache_controls() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache": {"ttl": 60, "no-cache": true}
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache": {"ttl": 3600, "no-store": true}
|
||||
});
|
||||
|
||||
assert_eq!(openai_key(&body1), openai_key(&body2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_namespace_control_separates_keys() {
|
||||
let body = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: Some("tenant-a"),
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: Some("tenant-b"),
|
||||
},
|
||||
);
|
||||
|
||||
assert_ne!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_litellm_fields() {
|
||||
let body = serde_json::json!({
|
||||
"cache": {
|
||||
"ttl": 120,
|
||||
"no-cache": true,
|
||||
"no-store": false,
|
||||
"s-maxage": 30,
|
||||
"namespace": "tenant-a",
|
||||
"use-cache": true
|
||||
}
|
||||
});
|
||||
|
||||
let control = parse_cache_control(&body).unwrap();
|
||||
|
||||
assert!(!control.lookup);
|
||||
assert!(control.store);
|
||||
assert_eq!(control.ttl_secs, Some(120));
|
||||
assert_eq!(control.max_age_secs, Some(30));
|
||||
assert_eq!(control.namespace.as_deref(), Some("tenant-a"));
|
||||
assert!(control.use_cache);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_preserves_cache_ttl_secs_bypass() {
|
||||
let body = serde_json::json!({
|
||||
"cache_ttl_secs": 0,
|
||||
"cache": {"ttl": 120}
|
||||
});
|
||||
|
||||
let control = parse_cache_control(&body).unwrap();
|
||||
|
||||
assert!(!control.lookup);
|
||||
assert!(!control.store);
|
||||
assert_eq!(control.ttl_secs, Some(120));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_rejects_invalid_cache_object() {
|
||||
let body = serde_json::json!({"cache": true});
|
||||
|
||||
assert!(parse_cache_control(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_entry_s_maxage_rejects_stale_entries() {
|
||||
let entry = CacheEntry {
|
||||
response_body: Bytes::from_static(b"{}"),
|
||||
model: "gpt-4o".to_string(),
|
||||
created_at: Instant::now() - std::time::Duration::from_secs(10),
|
||||
ttl_secs: None,
|
||||
};
|
||||
|
||||
assert!(!cache_entry_is_fresh(&entry, Some(5)));
|
||||
assert!(cache_entry_is_fresh(&entry, Some(30)));
|
||||
assert!(cache_entry_is_fresh(&entry, None));
|
||||
}
|
||||
|
||||
fn test_scope() -> CacheScope<'static> {
|
||||
CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_key(body: &serde_json::Value) -> String {
|
||||
cache_key_for_request(body, CacheNamespace::Anthropic, &test_scope())
|
||||
}
|
||||
|
||||
fn openai_key(body: &serde_json::Value) -> String {
|
||||
cache_key_for_request(body, CacheNamespace::OpenAI, &test_scope())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_anthropic_response_affecting_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
|
||||
let with_top_k = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"top_k": 10
|
||||
});
|
||||
let with_stop_sequences = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"stop_sequences": ["END"]
|
||||
});
|
||||
let with_thinking = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"thinking": {"type": "enabled", "budget_tokens": 1024}
|
||||
});
|
||||
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_top_k));
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_stop_sequences));
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_thinking));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_unknown_extra_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let with_extra = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"prediction": {"type": "content", "content": "expected"}
|
||||
});
|
||||
|
||||
assert_ne!(openai_key(&base), openai_key(&with_extra));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_tracking_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let with_user = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"user": "end-user-123"
|
||||
});
|
||||
let with_metadata = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"user_id": "session-abc"}
|
||||
});
|
||||
let metadata_base = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
|
||||
assert_eq!(openai_key(&base), openai_key(&with_user));
|
||||
assert_eq!(anthropic_key(&metadata_base), anthropic_key(&with_metadata));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_parallel_tool_calls() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"parameters": {"type": "object", "properties": {}}
|
||||
}
|
||||
}]
|
||||
});
|
||||
let with_parallel_tool_calls = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"parameters": {"type": "object", "properties": {}}
|
||||
}
|
||||
}],
|
||||
"parallel_tool_calls": false
|
||||
});
|
||||
|
||||
assert_ne!(openai_key(&base), openai_key(&with_parallel_tool_calls));
|
||||
}
|
||||
}
|
||||
|
||||
Vendored
+501
@@ -0,0 +1,501 @@
|
||||
use super::*;
|
||||
use bytes::Bytes;
|
||||
use std::time::Instant;
|
||||
|
||||
#[test]
|
||||
fn cache_key_deterministic_same_fields() {
|
||||
let body = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
assert!(key1.starts_with("anth:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_different_for_different_temperature() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.7
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.9
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_field_order() {
|
||||
// JSON object field order should not affect the key because we
|
||||
// extract into a BTreeMap.
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"temperature": 0.5,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"model": "gpt-4o",
|
||||
"temperature": 0.5
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_non_cache_fields() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"stream": true
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_namespace_differs() {
|
||||
let body = serde_json::json!({
|
||||
"model": "test",
|
||||
"messages": []
|
||||
});
|
||||
let anth = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::Anthropic,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let oai = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(anth, oai);
|
||||
assert!(anth.starts_with("anth:"));
|
||||
assert!(oai.starts_with("oai:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_null_field_same_as_absent() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [],
|
||||
"temperature": null
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": []
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_absent() {
|
||||
let body = serde_json::json!({"model": "test"});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_null() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": null});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_zero() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 0});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_valid() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 600});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(600));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_max() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 86400});
|
||||
assert_eq!(parse_cache_ttl(&body).unwrap(), Some(86400));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_over_max() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": 86401});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_negative() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": -1});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_ttl_string() {
|
||||
let body = serde_json::json!({"cache_ttl_secs": "not a number"});
|
||||
assert!(parse_cache_ttl(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_differs_for_different_cache_ttl_secs() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache_ttl_secs": 60
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache_ttl_secs": 3600
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body1,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body2,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
},
|
||||
);
|
||||
assert_ne!(
|
||||
key1, key2,
|
||||
"different cache_ttl_secs must produce different cache keys"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_litellm_cache_controls() {
|
||||
let body1 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache": {"ttl": 60, "no-cache": true}
|
||||
});
|
||||
let body2 = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"cache": {"ttl": 3600, "no-store": true}
|
||||
});
|
||||
|
||||
assert_eq!(openai_key(&body1), openai_key(&body2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_namespace_control_separates_keys() {
|
||||
let body = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let key1 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: Some("tenant-a"),
|
||||
},
|
||||
);
|
||||
let key2 = cache_key_for_request(
|
||||
&body,
|
||||
CacheNamespace::OpenAI,
|
||||
&CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: Some("tenant-b"),
|
||||
},
|
||||
);
|
||||
|
||||
assert_ne!(key1, key2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_litellm_fields() {
|
||||
let body = serde_json::json!({
|
||||
"cache": {
|
||||
"ttl": 120,
|
||||
"no-cache": true,
|
||||
"no-store": false,
|
||||
"s-maxage": 30,
|
||||
"namespace": "tenant-a",
|
||||
"use-cache": true
|
||||
}
|
||||
});
|
||||
|
||||
let control = parse_cache_control(&body).unwrap();
|
||||
|
||||
assert!(!control.lookup);
|
||||
assert!(control.store);
|
||||
assert_eq!(control.ttl_secs, Some(120));
|
||||
assert_eq!(control.max_age_secs, Some(30));
|
||||
assert_eq!(control.namespace.as_deref(), Some("tenant-a"));
|
||||
assert!(control.use_cache);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_preserves_cache_ttl_secs_bypass() {
|
||||
let body = serde_json::json!({
|
||||
"cache_ttl_secs": 0,
|
||||
"cache": {"ttl": 120}
|
||||
});
|
||||
|
||||
let control = parse_cache_control(&body).unwrap();
|
||||
|
||||
assert!(!control.lookup);
|
||||
assert!(!control.store);
|
||||
assert_eq!(control.ttl_secs, Some(120));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cache_control_rejects_invalid_cache_object() {
|
||||
let body = serde_json::json!({"cache": true});
|
||||
|
||||
assert!(parse_cache_control(&body).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_entry_s_maxage_rejects_stale_entries() {
|
||||
let entry = CacheEntry {
|
||||
response_body: Bytes::from_static(b"{}"),
|
||||
model: "gpt-4o".to_string(),
|
||||
created_at: Instant::now() - std::time::Duration::from_secs(10),
|
||||
ttl_secs: None,
|
||||
};
|
||||
|
||||
assert!(!cache_entry_is_fresh(&entry, Some(5)));
|
||||
assert!(cache_entry_is_fresh(&entry, Some(30)));
|
||||
assert!(cache_entry_is_fresh(&entry, None));
|
||||
}
|
||||
|
||||
fn test_scope() -> CacheScope<'static> {
|
||||
CacheScope {
|
||||
backend_name: "openai",
|
||||
auth_identity: "k1",
|
||||
namespace: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_key(body: &serde_json::Value) -> String {
|
||||
cache_key_for_request(body, CacheNamespace::Anthropic, &test_scope())
|
||||
}
|
||||
|
||||
fn openai_key(body: &serde_json::Value) -> String {
|
||||
cache_key_for_request(body, CacheNamespace::OpenAI, &test_scope())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_anthropic_response_affecting_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
|
||||
let with_top_k = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"top_k": 10
|
||||
});
|
||||
let with_stop_sequences = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"stop_sequences": ["END"]
|
||||
});
|
||||
let with_thinking = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"thinking": {"type": "enabled", "budget_tokens": 1024}
|
||||
});
|
||||
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_top_k));
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_stop_sequences));
|
||||
assert_ne!(anthropic_key(&base), anthropic_key(&with_thinking));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_unknown_extra_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let with_extra = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"prediction": {"type": "content", "content": "expected"}
|
||||
});
|
||||
|
||||
assert_ne!(openai_key(&base), openai_key(&with_extra));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_ignores_tracking_fields() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let with_user = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"user": "end-user-123"
|
||||
});
|
||||
let with_metadata = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"user_id": "session-abc"}
|
||||
});
|
||||
let metadata_base = serde_json::json!({
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
|
||||
assert_eq!(openai_key(&base), openai_key(&with_user));
|
||||
assert_eq!(anthropic_key(&metadata_base), anthropic_key(&with_metadata));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_includes_parallel_tool_calls() {
|
||||
let base = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"parameters": {"type": "object", "properties": {}}
|
||||
}
|
||||
}]
|
||||
});
|
||||
let with_parallel_tool_calls = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"parameters": {"type": "object", "properties": {}}
|
||||
}
|
||||
}],
|
||||
"parallel_tool_calls": false
|
||||
});
|
||||
|
||||
assert_ne!(openai_key(&base), openai_key(&with_parallel_tool_calls));
|
||||
}
|
||||
@@ -1,667 +1,18 @@
|
||||
/// LiteLLM config.yaml parser.
|
||||
///
|
||||
/// Accepts LiteLLM's YAML config format (model_list, litellm_settings,
|
||||
/// router_settings, general_settings) and converts it to anyllm-proxy's
|
||||
/// MultiConfig + ModelRouter.
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
//! LiteLLM config.yaml parser.
|
||||
//!
|
||||
//! Accepts LiteLLM's YAML config format (model_list, litellm_settings,
|
||||
//! router_settings, general_settings) and converts it to anyllm-proxy's
|
||||
//! MultiConfig + ModelRouter.
|
||||
//!
|
||||
//! [`types`] holds the Serde structs mapping the LiteLLM schema; [`parser`]
|
||||
//! holds the conversion functions.
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use serde::Deserialize;
|
||||
mod parser;
|
||||
mod types;
|
||||
|
||||
use super::model_router::{Deployment, ModelRouter, RoutingStrategy};
|
||||
use super::single::validate_gcp_identifier;
|
||||
use super::{
|
||||
resolve_env_value, validate_base_url, BackendAuth, BackendConfig, BackendKind, ModelMapping,
|
||||
MultiConfig, OpenAIApiFormat, TlsConfig,
|
||||
};
|
||||
|
||||
// ---- Serde structs for LiteLLM config.yaml ----
|
||||
|
||||
/// Root structure of a LiteLLM `config.yaml` file.
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct LiteLLMConfig {
|
||||
#[serde(default)]
|
||||
model_list: Vec<LiteLLMModelEntry>,
|
||||
#[serde(default)]
|
||||
litellm_settings: Option<LiteLLMSettings>,
|
||||
#[serde(default)]
|
||||
router_settings: Option<RouterSettings>,
|
||||
#[serde(default)]
|
||||
general_settings: Option<GeneralSettings>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LiteLLMModelEntry {
|
||||
model_name: String,
|
||||
litellm_params: LiteLLMParams,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LiteLLMParams {
|
||||
model: String,
|
||||
api_base: Option<String>,
|
||||
api_key: Option<String>,
|
||||
rpm: Option<u32>,
|
||||
tpm: Option<u64>,
|
||||
weight: Option<u32>,
|
||||
// Azure-specific
|
||||
api_version: Option<String>,
|
||||
// Vertex-specific
|
||||
vertex_project: Option<String>,
|
||||
vertex_location: Option<String>,
|
||||
// Bedrock-specific
|
||||
aws_access_key_id: Option<String>,
|
||||
aws_secret_access_key: Option<String>,
|
||||
aws_region_name: Option<String>,
|
||||
// Catch unknown fields silently (LiteLLM has many we don't support).
|
||||
#[serde(flatten)]
|
||||
_extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LiteLLMSettings {
|
||||
#[serde(default)]
|
||||
num_retries: Option<u32>,
|
||||
#[serde(default)]
|
||||
request_timeout: Option<u64>,
|
||||
#[serde(default)]
|
||||
callbacks: Vec<String>,
|
||||
#[serde(flatten)]
|
||||
_extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RouterSettings {
|
||||
#[serde(default)]
|
||||
routing_strategy: Option<String>,
|
||||
#[serde(flatten)]
|
||||
_extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Map LiteLLM routing_strategy string to our enum.
|
||||
pub(crate) fn parse_routing_strategy_str(s: &str) -> RoutingStrategy {
|
||||
match s.to_ascii_lowercase().replace('_', "-").as_str() {
|
||||
"simple-shuffle" | "round-robin" => RoutingStrategy::RoundRobin,
|
||||
"least-busy" => RoutingStrategy::LeastBusy,
|
||||
"latency-based-routing" | "latency-based" => RoutingStrategy::LatencyBased,
|
||||
"usage-based-routing" | "usage-based" => RoutingStrategy::LeastBusy,
|
||||
"weighted" => RoutingStrategy::Weighted,
|
||||
"cost-based" => RoutingStrategy::CostBased,
|
||||
other => {
|
||||
tracing::warn!(
|
||||
strategy = %other,
|
||||
"unknown routing_strategy, falling back to round-robin"
|
||||
);
|
||||
RoutingStrategy::RoundRobin
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeneralSettings {
|
||||
master_key: Option<String>,
|
||||
#[serde(flatten)]
|
||||
_extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
// ---- Provider parsing ----
|
||||
|
||||
/// Parse LiteLLM's "provider/model_name" format.
|
||||
/// No prefix defaults to OpenAI (matches LiteLLM behavior).
|
||||
/// Returns (kind, model_name, stub_provider) where stub_provider is set for
|
||||
/// registry-resolved OpenAI-compatible providers so callers can use their default URL.
|
||||
fn parse_provider_model(
|
||||
model: &str,
|
||||
) -> (
|
||||
BackendKind,
|
||||
String,
|
||||
Option<&'static anyllm_providers::ProviderDef>,
|
||||
) {
|
||||
let (provider, model_name) = model.split_once('/').unwrap_or(("openai", model));
|
||||
let mut stub_provider: Option<&'static anyllm_providers::ProviderDef> = None;
|
||||
let kind = match provider.to_ascii_lowercase().as_str() {
|
||||
"openai" => BackendKind::OpenAI,
|
||||
"azure" => BackendKind::AzureOpenAI,
|
||||
"vertex_ai" | "vertex" => BackendKind::Vertex,
|
||||
"gemini" => BackendKind::Gemini,
|
||||
"anthropic" => {
|
||||
stub_provider = anyllm_providers::get_provider("anthropic");
|
||||
BackendKind::Anthropic
|
||||
}
|
||||
"bedrock" => BackendKind::Bedrock,
|
||||
other => {
|
||||
// Try the provider registry for known OpenAI-compatible providers
|
||||
// (e.g. "groq", "together_ai", "mistral", etc.)
|
||||
let prefix_with_slash = format!("{other}/");
|
||||
if let Some(p) = anyllm_providers::find_by_litellm_prefix(&prefix_with_slash) {
|
||||
let resolved = match anyllm_providers::resolve_backend(p.id) {
|
||||
Some(("openai", _)) => {
|
||||
stub_provider = Some(p);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
Some(("anthropic", _)) => BackendKind::Anthropic,
|
||||
Some(("gemini", _)) => BackendKind::Gemini,
|
||||
Some(("vertex", _)) => BackendKind::Vertex,
|
||||
Some(("azure", _)) => BackendKind::AzureOpenAI,
|
||||
Some(("bedrock", _)) => BackendKind::Bedrock,
|
||||
_ => {
|
||||
tracing::warn!(provider = %other, "provider found in registry but protocol not mappable, treating as openai-compatible");
|
||||
stub_provider = Some(p);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
};
|
||||
resolved
|
||||
} else {
|
||||
tracing::warn!(
|
||||
provider = %other,
|
||||
"unknown LiteLLM provider, treating as openai-compatible"
|
||||
);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
}
|
||||
};
|
||||
(kind, model_name.to_string(), stub_provider)
|
||||
}
|
||||
|
||||
fn provider_id_for_litellm_model(
|
||||
model: &str,
|
||||
kind: &BackendKind,
|
||||
stub_provider: Option<&'static anyllm_providers::ProviderDef>,
|
||||
) -> String {
|
||||
if let Some(provider) = stub_provider {
|
||||
return provider.id.to_string();
|
||||
}
|
||||
|
||||
let raw_provider = model
|
||||
.split_once('/')
|
||||
.map(|(provider, _)| provider)
|
||||
.unwrap_or("openai")
|
||||
.to_ascii_lowercase();
|
||||
|
||||
match kind {
|
||||
BackendKind::OpenAI => raw_provider,
|
||||
BackendKind::AzureOpenAI => "azure".to_string(),
|
||||
BackendKind::Vertex => "vertex_ai".to_string(),
|
||||
BackendKind::Gemini => "gemini".to_string(),
|
||||
BackendKind::Anthropic => "anthropic".to_string(),
|
||||
BackendKind::Bedrock => "bedrock".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Backend deduplication key ----
|
||||
|
||||
/// Unique identity for a backend: same kind + base_url + api_key share one connection pool.
|
||||
#[derive(Hash, PartialEq, Eq, Clone)]
|
||||
struct BackendKey {
|
||||
kind: String,
|
||||
provider_id: String,
|
||||
base_url: String,
|
||||
/// Hash of the API key (not the key itself) to avoid holding secrets in hash keys.
|
||||
api_key_hash: u64,
|
||||
}
|
||||
|
||||
fn hash_string(s: &str) -> u64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
s.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
// ---- Conversion ----
|
||||
|
||||
/// Parse a LiteLLM config.yaml string and produce a MultiConfig + ModelRouter.
|
||||
///
|
||||
/// # Panics
|
||||
/// On invalid YAML, missing required fields, or unresolvable env var references.
|
||||
/// Parsed result from a LiteLLM config file.
|
||||
pub struct LiteLLMParsed {
|
||||
pub multi_config: MultiConfig,
|
||||
pub router: ModelRouter,
|
||||
/// Webhook callback URLs from litellm_settings.callbacks (non-named entries).
|
||||
pub callback_urls: Vec<String>,
|
||||
/// True when "langfuse" appears in litellm_settings.callbacks.
|
||||
pub langfuse_requested: bool,
|
||||
/// Resolved `general_settings.master_key`, if present.
|
||||
/// Caller should apply as PROXY_API_KEYS if that var is not already set.
|
||||
pub master_key: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse a LiteLLM YAML config and return the multi-backend config + model router pair.
|
||||
pub fn from_litellm_yaml(yaml: &str) -> (MultiConfig, ModelRouter) {
|
||||
let parsed = parse_litellm_yaml(yaml);
|
||||
(parsed.multi_config, parsed.router)
|
||||
}
|
||||
|
||||
/// Parse a LiteLLM YAML config into the intermediate `LiteLLMParsed` struct.
|
||||
/// Panics on invalid YAML (startup-time validation; misconfiguration is unrecoverable).
|
||||
pub fn parse_litellm_yaml(yaml: &str) -> LiteLLMParsed {
|
||||
let config: LiteLLMConfig =
|
||||
serde_yaml::from_str(yaml).unwrap_or_else(|e| panic!("invalid LiteLLM config YAML: {e}"));
|
||||
|
||||
if config.model_list.is_empty() {
|
||||
panic!("LiteLLM config must define at least one entry in model_list");
|
||||
}
|
||||
|
||||
// Resolve general_settings.master_key but do not call set_var here.
|
||||
// The caller applies it in the consolidated env override block.
|
||||
let master_key = if let Some(ref gs) = config.general_settings {
|
||||
let mk = gs.master_key.as_ref().map(|mk| {
|
||||
resolve_env_value(mk).unwrap_or_else(|e| panic!("general_settings.master_key: {e}"))
|
||||
});
|
||||
// Log unsupported keys at warn.
|
||||
for key in gs._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported general_settings key (ignored)");
|
||||
}
|
||||
mk
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(ref ls) = config.litellm_settings {
|
||||
for key in ls._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported litellm_settings key (ignored)");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref rs) = config.router_settings {
|
||||
for key in rs._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported router_settings key (ignored)");
|
||||
}
|
||||
}
|
||||
|
||||
let listen_port = std::env::var("LISTEN_PORT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(3000);
|
||||
|
||||
let log_bodies = std::env::var("LOG_BODIES")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let redact_secrets = std::env::var("REDACT_SECRETS")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let anthropic_thinking_repair = std::env::var("ANTHROPIC_THINKING_REPAIR")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let pxpipe_compress = std::env::var("PXPIPE_COMPRESS")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
|
||||
let tls = TlsConfig::from_env();
|
||||
|
||||
// Group model_list entries into deduplicated backends + deployment list.
|
||||
let mut backend_map: HashMap<BackendKey, (String, BackendConfig)> = HashMap::new();
|
||||
let mut backend_counter = 0u32;
|
||||
// model_name -> Vec<(backend_name, actual_model, rpm, tpm)>
|
||||
let mut model_deployments: HashMap<String, Vec<DeploymentSpec>> = HashMap::new();
|
||||
|
||||
for entry in &config.model_list {
|
||||
let (kind, actual_model, stub_provider) = parse_provider_model(&entry.litellm_params.model);
|
||||
let provider_id =
|
||||
provider_id_for_litellm_model(&entry.litellm_params.model, &kind, stub_provider);
|
||||
let params = &entry.litellm_params;
|
||||
|
||||
let api_key = super::sanitize_api_key(
|
||||
¶ms
|
||||
.api_key
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("model_list api_key: {e}")))
|
||||
.unwrap_or_else(|| {
|
||||
// Fall back to the provider's own env vars when no api_key in YAML.
|
||||
stub_provider
|
||||
.and_then(|p| p.env_vars.iter().find_map(|v| std::env::var(v).ok()))
|
||||
.or_else(|| {
|
||||
(kind == BackendKind::Anthropic)
|
||||
.then(|| std::env::var("ANTHROPIC_AUTH_TOKEN").ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
);
|
||||
|
||||
let base_url = resolve_base_url(&kind, params, stub_provider, &actual_model);
|
||||
|
||||
let bk = BackendKey {
|
||||
kind: format!("{kind:?}"),
|
||||
provider_id: provider_id.clone(),
|
||||
base_url: base_url.clone(),
|
||||
api_key_hash: hash_string(&api_key),
|
||||
};
|
||||
|
||||
let backend_name = if let Some((name, _)) = backend_map.get(&bk) {
|
||||
name.clone()
|
||||
} else {
|
||||
let name = format!("litellm_{backend_counter}");
|
||||
backend_counter += 1;
|
||||
|
||||
let bc = build_backend_config(
|
||||
&name,
|
||||
&kind,
|
||||
&provider_id,
|
||||
&api_key,
|
||||
&base_url,
|
||||
params,
|
||||
&tls,
|
||||
log_bodies,
|
||||
&config,
|
||||
);
|
||||
backend_map.insert(bk, (name.clone(), bc));
|
||||
name
|
||||
};
|
||||
|
||||
model_deployments
|
||||
.entry(entry.model_name.clone())
|
||||
.or_default()
|
||||
.push(DeploymentSpec {
|
||||
backend_name,
|
||||
actual_model,
|
||||
rpm: params.rpm,
|
||||
tpm: params.tpm,
|
||||
weight: params.weight,
|
||||
});
|
||||
}
|
||||
|
||||
// Build MultiConfig backends (ordered).
|
||||
let mut backends = IndexMap::new();
|
||||
for (name, bc) in backend_map.values() {
|
||||
backends.insert(name.clone(), bc.clone());
|
||||
}
|
||||
|
||||
let default_backend = backends
|
||||
.keys()
|
||||
.next()
|
||||
.cloned()
|
||||
.expect("at least one backend");
|
||||
|
||||
let multi = MultiConfig {
|
||||
listen_port,
|
||||
log_bodies,
|
||||
redact_secrets,
|
||||
anthropic_thinking_repair,
|
||||
pxpipe_compress,
|
||||
forward_client_auth: crate::config::env_bool_flag("ANTHROPIC_FORWARD_CLIENT_AUTH"),
|
||||
default_backend,
|
||||
backends,
|
||||
expose_degradation_warnings: false, // overridden in MultiConfig::load()
|
||||
};
|
||||
|
||||
// Determine routing strategy from router_settings.
|
||||
let strategy = config
|
||||
.router_settings
|
||||
.as_ref()
|
||||
.and_then(|rs| rs.routing_strategy.as_deref())
|
||||
.map(parse_routing_strategy_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
if strategy != RoutingStrategy::RoundRobin {
|
||||
tracing::info!(strategy = ?strategy, "using routing strategy from config");
|
||||
}
|
||||
|
||||
// Build ModelRouter.
|
||||
let mut routes: HashMap<String, Vec<Arc<Deployment>>> = HashMap::new();
|
||||
for (model_name, specs) in model_deployments {
|
||||
let deployments = specs
|
||||
.into_iter()
|
||||
.map(|s| {
|
||||
Arc::new(Deployment::with_weight(
|
||||
s.backend_name,
|
||||
s.actual_model,
|
||||
s.rpm,
|
||||
s.tpm,
|
||||
s.weight.unwrap_or(1),
|
||||
))
|
||||
})
|
||||
.collect();
|
||||
routes.insert(model_name, deployments);
|
||||
}
|
||||
|
||||
let router = ModelRouter::with_strategy(routes, strategy);
|
||||
|
||||
let callbacks = config
|
||||
.litellm_settings
|
||||
.as_ref()
|
||||
.map(|s| s.callbacks.clone())
|
||||
.unwrap_or_default();
|
||||
let langfuse_requested = callbacks.iter().any(|c| c.eq_ignore_ascii_case("langfuse"));
|
||||
let callback_urls: Vec<String> = callbacks
|
||||
.into_iter()
|
||||
.filter(|c| !c.eq_ignore_ascii_case("langfuse"))
|
||||
.collect();
|
||||
|
||||
LiteLLMParsed {
|
||||
multi_config: multi,
|
||||
router,
|
||||
callback_urls,
|
||||
langfuse_requested,
|
||||
master_key,
|
||||
}
|
||||
}
|
||||
|
||||
struct DeploymentSpec {
|
||||
backend_name: String,
|
||||
actual_model: String,
|
||||
rpm: Option<u32>,
|
||||
tpm: Option<u64>,
|
||||
weight: Option<u32>,
|
||||
}
|
||||
|
||||
/// Extract `general_settings.master_key` from a LiteLLM YAML string without
|
||||
/// performing full config parsing. Used by the synchronous `fn main()` to apply
|
||||
/// the key via `set_var` before the tokio runtime spawns worker threads.
|
||||
pub fn extract_master_key(yaml: &str) -> Option<String> {
|
||||
#[derive(Deserialize)]
|
||||
struct Probe {
|
||||
general_settings: Option<GeneralSettings>,
|
||||
}
|
||||
let probe: Probe = serde_yaml::from_str(yaml).ok()?;
|
||||
let gs = probe.general_settings?;
|
||||
let raw = gs.master_key?;
|
||||
resolve_env_value(&raw).ok()
|
||||
}
|
||||
|
||||
/// Determine the base URL for a deployment, applying provider-specific defaults.
|
||||
fn resolve_base_url(
|
||||
kind: &BackendKind,
|
||||
params: &LiteLLMParams,
|
||||
stub_provider: Option<&'static anyllm_providers::ProviderDef>,
|
||||
actual_model: &str,
|
||||
) -> String {
|
||||
if let Some(ref url) = params.api_base {
|
||||
let resolved =
|
||||
resolve_env_value(url).unwrap_or_else(|e| panic!("model_list api_base: {e}"));
|
||||
if *kind == BackendKind::AzureOpenAI {
|
||||
let api_version = params.api_version.as_deref().unwrap_or("2024-10-21");
|
||||
if !resolved.contains("/openai/deployments/") {
|
||||
let deployment = azure_deployment_from_model(actual_model);
|
||||
return format!(
|
||||
"{}/openai/deployments/{deployment}/chat/completions?api-version={api_version}",
|
||||
resolved.trim_end_matches('/'),
|
||||
);
|
||||
}
|
||||
// api_base is already a deployment URL. Ensure the required api-version
|
||||
// query is present so a partial deployment URL still authenticates.
|
||||
if !resolved.contains("api-version=") {
|
||||
let sep = if resolved.contains('?') { '&' } else { '?' };
|
||||
return format!("{resolved}{sep}api-version={api_version}");
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
match kind {
|
||||
BackendKind::OpenAI => {
|
||||
// Use the stub provider's default URL when available (e.g. groq, xai, mistral).
|
||||
// If a known provider has no safe global default, require explicit api_base.
|
||||
let url = if let Some(provider) = stub_provider {
|
||||
if provider.default_base_url.is_empty() {
|
||||
panic!(
|
||||
"model_list provider '{}' requires api_base because it has no safe global API base URL",
|
||||
provider.id
|
||||
);
|
||||
}
|
||||
provider.default_base_url
|
||||
} else {
|
||||
"https://api.openai.com"
|
||||
};
|
||||
super::strip_v1_suffix(url).to_string()
|
||||
}
|
||||
BackendKind::Gemini => {
|
||||
"https://generativelanguage.googleapis.com/v1beta/openai".to_string()
|
||||
}
|
||||
BackendKind::Anthropic => std::env::var("ANTHROPIC_BASE_URL")
|
||||
.unwrap_or_else(|_| "https://api.anthropic.com".to_string()),
|
||||
BackendKind::Bedrock => {
|
||||
// For Bedrock, base_url stores the region.
|
||||
params
|
||||
.aws_region_name
|
||||
.as_deref()
|
||||
.map(|v| v.to_string())
|
||||
.or_else(|| std::env::var("AWS_REGION").ok())
|
||||
.unwrap_or_else(|| "us-east-1".to_string())
|
||||
}
|
||||
// Azure and Vertex require api_base in the config.
|
||||
BackendKind::AzureOpenAI => {
|
||||
panic!("api_base is required for azure deployments in model_list")
|
||||
}
|
||||
BackendKind::Vertex => {
|
||||
let project = params.vertex_project.as_deref().unwrap_or_else(|| {
|
||||
panic!("vertex_project is required for vertex deployments in model_list")
|
||||
});
|
||||
let location = params.vertex_location.as_deref().unwrap_or_else(|| {
|
||||
panic!("vertex_location is required for vertex deployments in model_list")
|
||||
});
|
||||
validate_gcp_identifier("vertex_project", project);
|
||||
validate_gcp_identifier("vertex_location", location);
|
||||
format!(
|
||||
"https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// LiteLLM Azure model names may carry a route-group marker prefix that selects a
|
||||
// request-shaping path but is not part of the deployment name. This is the complete
|
||||
// set of such markers; they are defined only here (no external producer). Add a
|
||||
// marker here if a new Azure route group is introduced.
|
||||
fn azure_deployment_from_model(model: &str) -> &str {
|
||||
for marker in ["o_series/", "gpt5_series/"] {
|
||||
if let Some(deployment) = model.strip_prefix(marker) {
|
||||
if !deployment.is_empty() {
|
||||
return deployment;
|
||||
}
|
||||
}
|
||||
}
|
||||
model
|
||||
}
|
||||
|
||||
/// Build a BackendConfig from LiteLLM model_list params.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn build_backend_config(
|
||||
name: &str,
|
||||
kind: &BackendKind,
|
||||
provider_id: &str,
|
||||
api_key: &str,
|
||||
base_url: &str,
|
||||
params: &LiteLLMParams,
|
||||
tls: &TlsConfig,
|
||||
log_bodies: bool,
|
||||
config: &LiteLLMConfig,
|
||||
) -> BackendConfig {
|
||||
let backend_auth = match kind {
|
||||
BackendKind::AzureOpenAI => BackendAuth::AzureApiKey(api_key.to_string()),
|
||||
BackendKind::Gemini | BackendKind::Vertex => BackendAuth::GoogleApiKey(api_key.to_string()),
|
||||
BackendKind::Anthropic => BackendAuth::anthropic_from_api_key_like(api_key.to_string()),
|
||||
_ => BackendAuth::BearerToken(api_key.to_string()),
|
||||
};
|
||||
|
||||
// For Azure, resolve_base_url already produced the full deployment URL
|
||||
// (deployment name + api-version), so use it as-is here.
|
||||
let effective_url = if *kind == BackendKind::AzureOpenAI {
|
||||
base_url.to_string()
|
||||
} else {
|
||||
// Validate non-Azure URLs.
|
||||
if *kind != BackendKind::Bedrock {
|
||||
if let Err(e) = validate_base_url(base_url) {
|
||||
panic!("backend '{name}' base_url rejected: {e}");
|
||||
}
|
||||
}
|
||||
base_url.to_string()
|
||||
};
|
||||
|
||||
// Bedrock credentials.
|
||||
let bedrock_credentials = if *kind == BackendKind::Bedrock {
|
||||
let region = params
|
||||
.aws_region_name
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_REGION").ok())
|
||||
.unwrap_or_else(|| "us-east-1".to_string());
|
||||
|
||||
let access_key = params
|
||||
.aws_access_key_id
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_ACCESS_KEY_ID").ok())
|
||||
.unwrap_or_else(|| panic!("backend '{name}': aws_access_key_id required for bedrock"));
|
||||
|
||||
let secret_key = params
|
||||
.aws_secret_access_key
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_SECRET_ACCESS_KEY").ok())
|
||||
.unwrap_or_else(|| {
|
||||
panic!("backend '{name}': aws_secret_access_key required for bedrock")
|
||||
});
|
||||
|
||||
// Store region as base_url for Bedrock (matches existing convention).
|
||||
// The effective_url was already set to the region string.
|
||||
let _ = region; // region is used as base_url via resolve_base_url
|
||||
|
||||
Some(aws_credential_types::Credentials::new(
|
||||
access_key,
|
||||
secret_key,
|
||||
None, // session token not commonly in LiteLLM configs
|
||||
None,
|
||||
"litellm-config",
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Placeholder model mapping: with model router, these are not used for routing.
|
||||
// They serve as fallback for Anthropic model name translation if needed.
|
||||
let model_mapping = ModelMapping {
|
||||
big_model: String::new(),
|
||||
small_model: String::new(),
|
||||
};
|
||||
|
||||
let _num_retries = config.litellm_settings.as_ref().and_then(|s| s.num_retries);
|
||||
let _request_timeout = config
|
||||
.litellm_settings
|
||||
.as_ref()
|
||||
.and_then(|s| s.request_timeout);
|
||||
|
||||
BackendConfig {
|
||||
kind: kind.clone(),
|
||||
provider_id: Some(provider_id.to_string()),
|
||||
api_key: api_key.to_string(),
|
||||
base_url: effective_url,
|
||||
api_format: OpenAIApiFormat::Chat,
|
||||
model_mapping,
|
||||
tls: tls.clone(),
|
||||
backend_auth,
|
||||
log_bodies,
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: std::env::var("REQUEST_TIMEOUT_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(900u64),
|
||||
bedrock_credentials,
|
||||
// Config-file backends are validated against SSRF at load time; no auto-relax.
|
||||
allow_local_ssrf: false,
|
||||
}
|
||||
}
|
||||
pub(crate) use parser::parse_routing_strategy_str;
|
||||
pub use parser::{extract_master_key, from_litellm_yaml, parse_litellm_yaml};
|
||||
pub use types::LiteLLMParsed;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,576 @@
|
||||
//! Conversion logic: turn parsed LiteLLM YAML structs into anyllm-proxy's
|
||||
//! `MultiConfig` + `ModelRouter`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::types::{GeneralSettings, LiteLLMConfig, LiteLLMParams, LiteLLMParsed};
|
||||
use crate::config::model_router::{Deployment, ModelRouter, RoutingStrategy};
|
||||
use crate::config::single::validate_gcp_identifier;
|
||||
use crate::config::{
|
||||
resolve_env_value, sanitize_api_key, strip_v1_suffix, validate_base_url, BackendAuth,
|
||||
BackendConfig, BackendKind, ModelMapping, MultiConfig, OpenAIApiFormat, TlsConfig,
|
||||
};
|
||||
|
||||
// ---- Provider parsing ----
|
||||
|
||||
/// Map LiteLLM routing_strategy string to our enum.
|
||||
pub(crate) fn parse_routing_strategy_str(s: &str) -> RoutingStrategy {
|
||||
match s.to_ascii_lowercase().replace('_', "-").as_str() {
|
||||
"simple-shuffle" | "round-robin" => RoutingStrategy::RoundRobin,
|
||||
"least-busy" => RoutingStrategy::LeastBusy,
|
||||
"latency-based-routing" | "latency-based" => RoutingStrategy::LatencyBased,
|
||||
"usage-based-routing" | "usage-based" => RoutingStrategy::LeastBusy,
|
||||
"weighted" => RoutingStrategy::Weighted,
|
||||
"cost-based" => RoutingStrategy::CostBased,
|
||||
other => {
|
||||
tracing::warn!(
|
||||
strategy = %other,
|
||||
"unknown routing_strategy, falling back to round-robin"
|
||||
);
|
||||
RoutingStrategy::RoundRobin
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse LiteLLM's "provider/model_name" format.
|
||||
/// No prefix defaults to OpenAI (matches LiteLLM behavior).
|
||||
/// Returns (kind, model_name, stub_provider) where stub_provider is set for
|
||||
/// registry-resolved OpenAI-compatible providers so callers can use their default URL.
|
||||
pub(super) fn parse_provider_model(
|
||||
model: &str,
|
||||
) -> (
|
||||
BackendKind,
|
||||
String,
|
||||
Option<&'static anyllm_providers::ProviderDef>,
|
||||
) {
|
||||
let (provider, model_name) = model.split_once('/').unwrap_or(("openai", model));
|
||||
let mut stub_provider: Option<&'static anyllm_providers::ProviderDef> = None;
|
||||
let kind = match provider.to_ascii_lowercase().as_str() {
|
||||
"openai" => BackendKind::OpenAI,
|
||||
"azure" => BackendKind::AzureOpenAI,
|
||||
"vertex_ai" | "vertex" => BackendKind::Vertex,
|
||||
"gemini" => BackendKind::Gemini,
|
||||
"anthropic" => {
|
||||
stub_provider = anyllm_providers::get_provider("anthropic");
|
||||
BackendKind::Anthropic
|
||||
}
|
||||
"bedrock" => BackendKind::Bedrock,
|
||||
other => {
|
||||
// Try the provider registry for known OpenAI-compatible providers
|
||||
// (e.g. "groq", "together_ai", "mistral", etc.)
|
||||
let prefix_with_slash = format!("{other}/");
|
||||
if let Some(p) = anyllm_providers::find_by_litellm_prefix(&prefix_with_slash) {
|
||||
let resolved = match anyllm_providers::resolve_backend(p.id) {
|
||||
Some(("openai", _)) => {
|
||||
stub_provider = Some(p);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
Some(("anthropic", _)) => BackendKind::Anthropic,
|
||||
Some(("gemini", _)) => BackendKind::Gemini,
|
||||
Some(("vertex", _)) => BackendKind::Vertex,
|
||||
Some(("azure", _)) => BackendKind::AzureOpenAI,
|
||||
Some(("bedrock", _)) => BackendKind::Bedrock,
|
||||
_ => {
|
||||
tracing::warn!(provider = %other, "provider found in registry but protocol not mappable, treating as openai-compatible");
|
||||
stub_provider = Some(p);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
};
|
||||
resolved
|
||||
} else {
|
||||
tracing::warn!(
|
||||
provider = %other,
|
||||
"unknown LiteLLM provider, treating as openai-compatible"
|
||||
);
|
||||
BackendKind::OpenAI
|
||||
}
|
||||
}
|
||||
};
|
||||
(kind, model_name.to_string(), stub_provider)
|
||||
}
|
||||
|
||||
fn provider_id_for_litellm_model(
|
||||
model: &str,
|
||||
kind: &BackendKind,
|
||||
stub_provider: Option<&'static anyllm_providers::ProviderDef>,
|
||||
) -> String {
|
||||
if let Some(provider) = stub_provider {
|
||||
return provider.id.to_string();
|
||||
}
|
||||
|
||||
let raw_provider = model
|
||||
.split_once('/')
|
||||
.map(|(provider, _)| provider)
|
||||
.unwrap_or("openai")
|
||||
.to_ascii_lowercase();
|
||||
|
||||
match kind {
|
||||
BackendKind::OpenAI => raw_provider,
|
||||
BackendKind::AzureOpenAI => "azure".to_string(),
|
||||
BackendKind::Vertex => "vertex_ai".to_string(),
|
||||
BackendKind::Gemini => "gemini".to_string(),
|
||||
BackendKind::Anthropic => "anthropic".to_string(),
|
||||
BackendKind::Bedrock => "bedrock".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Backend deduplication key ----
|
||||
|
||||
/// Unique identity for a backend: same kind + base_url + api_key share one connection pool.
|
||||
#[derive(Hash, PartialEq, Eq, Clone)]
|
||||
struct BackendKey {
|
||||
kind: String,
|
||||
provider_id: String,
|
||||
base_url: String,
|
||||
/// Hash of the API key (not the key itself) to avoid holding secrets in hash keys.
|
||||
api_key_hash: u64,
|
||||
}
|
||||
|
||||
fn hash_string(s: &str) -> u64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
s.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
// ---- Conversion ----
|
||||
|
||||
/// Parse a LiteLLM YAML config and return the multi-backend config + model router pair.
|
||||
pub fn from_litellm_yaml(yaml: &str) -> (MultiConfig, ModelRouter) {
|
||||
let parsed = parse_litellm_yaml(yaml);
|
||||
(parsed.multi_config, parsed.router)
|
||||
}
|
||||
|
||||
/// Parse a LiteLLM YAML config into the intermediate `LiteLLMParsed` struct.
|
||||
/// Panics on invalid YAML (startup-time validation; misconfiguration is unrecoverable).
|
||||
pub fn parse_litellm_yaml(yaml: &str) -> LiteLLMParsed {
|
||||
let config: LiteLLMConfig =
|
||||
serde_yaml::from_str(yaml).unwrap_or_else(|e| panic!("invalid LiteLLM config YAML: {e}"));
|
||||
|
||||
if config.model_list.is_empty() {
|
||||
panic!("LiteLLM config must define at least one entry in model_list");
|
||||
}
|
||||
|
||||
// Resolve general_settings.master_key but do not call set_var here.
|
||||
// The caller applies it in the consolidated env override block.
|
||||
let master_key = if let Some(ref gs) = config.general_settings {
|
||||
let mk = gs.master_key.as_ref().map(|mk| {
|
||||
resolve_env_value(mk).unwrap_or_else(|e| panic!("general_settings.master_key: {e}"))
|
||||
});
|
||||
// Log unsupported keys at warn.
|
||||
for key in gs._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported general_settings key (ignored)");
|
||||
}
|
||||
mk
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(ref ls) = config.litellm_settings {
|
||||
for key in ls._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported litellm_settings key (ignored)");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref rs) = config.router_settings {
|
||||
for key in rs._extra.keys() {
|
||||
tracing::warn!(key = %key, "unsupported router_settings key (ignored)");
|
||||
}
|
||||
}
|
||||
|
||||
let listen_port = std::env::var("LISTEN_PORT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(3000);
|
||||
|
||||
let log_bodies = std::env::var("LOG_BODIES")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let redact_secrets = std::env::var("REDACT_SECRETS")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let anthropic_thinking_repair = std::env::var("ANTHROPIC_THINKING_REPAIR")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
let pxpipe_compress = std::env::var("PXPIPE_COMPRESS")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false);
|
||||
|
||||
let tls = TlsConfig::from_env();
|
||||
|
||||
// Group model_list entries into deduplicated backends + deployment list.
|
||||
let mut backend_map: HashMap<BackendKey, (String, BackendConfig)> = HashMap::new();
|
||||
let mut backend_counter = 0u32;
|
||||
// model_name -> Vec<(backend_name, actual_model, rpm, tpm)>
|
||||
let mut model_deployments: HashMap<String, Vec<DeploymentSpec>> = HashMap::new();
|
||||
|
||||
for entry in &config.model_list {
|
||||
let (kind, actual_model, stub_provider) = parse_provider_model(&entry.litellm_params.model);
|
||||
let provider_id =
|
||||
provider_id_for_litellm_model(&entry.litellm_params.model, &kind, stub_provider);
|
||||
let params = &entry.litellm_params;
|
||||
|
||||
let api_key = sanitize_api_key(
|
||||
¶ms
|
||||
.api_key
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("model_list api_key: {e}")))
|
||||
.unwrap_or_else(|| {
|
||||
// Fall back to the provider's own env vars when no api_key in YAML.
|
||||
stub_provider
|
||||
.and_then(|p| p.env_vars.iter().find_map(|v| std::env::var(v).ok()))
|
||||
.or_else(|| {
|
||||
(kind == BackendKind::Anthropic)
|
||||
.then(|| std::env::var("ANTHROPIC_AUTH_TOKEN").ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
);
|
||||
|
||||
let base_url = resolve_base_url(&kind, params, stub_provider, &actual_model);
|
||||
|
||||
let bk = BackendKey {
|
||||
kind: format!("{kind:?}"),
|
||||
provider_id: provider_id.clone(),
|
||||
base_url: base_url.clone(),
|
||||
api_key_hash: hash_string(&api_key),
|
||||
};
|
||||
|
||||
let backend_name = if let Some((name, _)) = backend_map.get(&bk) {
|
||||
name.clone()
|
||||
} else {
|
||||
let name = format!("litellm_{backend_counter}");
|
||||
backend_counter += 1;
|
||||
|
||||
let bc = build_backend_config(
|
||||
&name,
|
||||
&kind,
|
||||
&provider_id,
|
||||
&api_key,
|
||||
&base_url,
|
||||
params,
|
||||
&tls,
|
||||
log_bodies,
|
||||
&config,
|
||||
);
|
||||
backend_map.insert(bk, (name.clone(), bc));
|
||||
name
|
||||
};
|
||||
|
||||
model_deployments
|
||||
.entry(entry.model_name.clone())
|
||||
.or_default()
|
||||
.push(DeploymentSpec {
|
||||
backend_name,
|
||||
actual_model,
|
||||
rpm: params.rpm,
|
||||
tpm: params.tpm,
|
||||
weight: params.weight,
|
||||
});
|
||||
}
|
||||
|
||||
// Build MultiConfig backends (ordered).
|
||||
let mut backends = IndexMap::new();
|
||||
for (name, bc) in backend_map.values() {
|
||||
backends.insert(name.clone(), bc.clone());
|
||||
}
|
||||
|
||||
let default_backend = backends
|
||||
.keys()
|
||||
.next()
|
||||
.cloned()
|
||||
.expect("at least one backend");
|
||||
|
||||
let multi = MultiConfig {
|
||||
listen_port,
|
||||
log_bodies,
|
||||
redact_secrets,
|
||||
anthropic_thinking_repair,
|
||||
pxpipe_compress,
|
||||
forward_client_auth: crate::config::env_bool_flag("ANTHROPIC_FORWARD_CLIENT_AUTH"),
|
||||
default_backend,
|
||||
backends,
|
||||
expose_degradation_warnings: false, // overridden in MultiConfig::load()
|
||||
};
|
||||
|
||||
// Determine routing strategy from router_settings.
|
||||
let strategy = config
|
||||
.router_settings
|
||||
.as_ref()
|
||||
.and_then(|rs| rs.routing_strategy.as_deref())
|
||||
.map(parse_routing_strategy_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
if strategy != RoutingStrategy::RoundRobin {
|
||||
tracing::info!(strategy = ?strategy, "using routing strategy from config");
|
||||
}
|
||||
|
||||
// Build ModelRouter.
|
||||
let mut routes: HashMap<String, Vec<Arc<Deployment>>> = HashMap::new();
|
||||
for (model_name, specs) in model_deployments {
|
||||
let deployments = specs
|
||||
.into_iter()
|
||||
.map(|s| {
|
||||
Arc::new(Deployment::with_weight(
|
||||
s.backend_name,
|
||||
s.actual_model,
|
||||
s.rpm,
|
||||
s.tpm,
|
||||
s.weight.unwrap_or(1),
|
||||
))
|
||||
})
|
||||
.collect();
|
||||
routes.insert(model_name, deployments);
|
||||
}
|
||||
|
||||
let router = ModelRouter::with_strategy(routes, strategy);
|
||||
|
||||
let callbacks = config
|
||||
.litellm_settings
|
||||
.as_ref()
|
||||
.map(|s| s.callbacks.clone())
|
||||
.unwrap_or_default();
|
||||
let langfuse_requested = callbacks.iter().any(|c| c.eq_ignore_ascii_case("langfuse"));
|
||||
let callback_urls: Vec<String> = callbacks
|
||||
.into_iter()
|
||||
.filter(|c| !c.eq_ignore_ascii_case("langfuse"))
|
||||
.collect();
|
||||
|
||||
LiteLLMParsed {
|
||||
multi_config: multi,
|
||||
router,
|
||||
callback_urls,
|
||||
langfuse_requested,
|
||||
master_key,
|
||||
}
|
||||
}
|
||||
|
||||
struct DeploymentSpec {
|
||||
backend_name: String,
|
||||
actual_model: String,
|
||||
rpm: Option<u32>,
|
||||
tpm: Option<u64>,
|
||||
weight: Option<u32>,
|
||||
}
|
||||
|
||||
/// Extract `general_settings.master_key` from a LiteLLM YAML string without
|
||||
/// performing full config parsing. Used by the synchronous `fn main()` to apply
|
||||
/// the key via `set_var` before the tokio runtime spawns worker threads.
|
||||
pub fn extract_master_key(yaml: &str) -> Option<String> {
|
||||
#[derive(Deserialize)]
|
||||
struct Probe {
|
||||
general_settings: Option<GeneralSettings>,
|
||||
}
|
||||
let probe: Probe = serde_yaml::from_str(yaml).ok()?;
|
||||
let gs = probe.general_settings?;
|
||||
let raw = gs.master_key?;
|
||||
resolve_env_value(&raw).ok()
|
||||
}
|
||||
|
||||
/// Determine the base URL for a deployment, applying provider-specific defaults.
|
||||
fn resolve_base_url(
|
||||
kind: &BackendKind,
|
||||
params: &LiteLLMParams,
|
||||
stub_provider: Option<&'static anyllm_providers::ProviderDef>,
|
||||
actual_model: &str,
|
||||
) -> String {
|
||||
if let Some(ref url) = params.api_base {
|
||||
let resolved =
|
||||
resolve_env_value(url).unwrap_or_else(|e| panic!("model_list api_base: {e}"));
|
||||
if *kind == BackendKind::AzureOpenAI {
|
||||
let api_version = params.api_version.as_deref().unwrap_or("2024-10-21");
|
||||
if !resolved.contains("/openai/deployments/") {
|
||||
let deployment = azure_deployment_from_model(actual_model);
|
||||
return format!(
|
||||
"{}/openai/deployments/{deployment}/chat/completions?api-version={api_version}",
|
||||
resolved.trim_end_matches('/'),
|
||||
);
|
||||
}
|
||||
// api_base is already a deployment URL. Ensure the required api-version
|
||||
// query is present so a partial deployment URL still authenticates.
|
||||
if !resolved.contains("api-version=") {
|
||||
let sep = if resolved.contains('?') { '&' } else { '?' };
|
||||
return format!("{resolved}{sep}api-version={api_version}");
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
match kind {
|
||||
BackendKind::OpenAI => {
|
||||
// Use the stub provider's default URL when available (e.g. groq, xai, mistral).
|
||||
// If a known provider has no safe global default, require explicit api_base.
|
||||
let url = if let Some(provider) = stub_provider {
|
||||
if provider.default_base_url.is_empty() {
|
||||
panic!(
|
||||
"model_list provider '{}' requires api_base because it has no safe global API base URL",
|
||||
provider.id
|
||||
);
|
||||
}
|
||||
provider.default_base_url
|
||||
} else {
|
||||
"https://api.openai.com"
|
||||
};
|
||||
strip_v1_suffix(url).to_string()
|
||||
}
|
||||
BackendKind::Gemini => {
|
||||
"https://generativelanguage.googleapis.com/v1beta/openai".to_string()
|
||||
}
|
||||
BackendKind::Anthropic => std::env::var("ANTHROPIC_BASE_URL")
|
||||
.unwrap_or_else(|_| "https://api.anthropic.com".to_string()),
|
||||
BackendKind::Bedrock => {
|
||||
// For Bedrock, base_url stores the region.
|
||||
params
|
||||
.aws_region_name
|
||||
.as_deref()
|
||||
.map(|v| v.to_string())
|
||||
.or_else(|| std::env::var("AWS_REGION").ok())
|
||||
.unwrap_or_else(|| "us-east-1".to_string())
|
||||
}
|
||||
// Azure and Vertex require api_base in the config.
|
||||
BackendKind::AzureOpenAI => {
|
||||
panic!("api_base is required for azure deployments in model_list")
|
||||
}
|
||||
BackendKind::Vertex => {
|
||||
let project = params.vertex_project.as_deref().unwrap_or_else(|| {
|
||||
panic!("vertex_project is required for vertex deployments in model_list")
|
||||
});
|
||||
let location = params.vertex_location.as_deref().unwrap_or_else(|| {
|
||||
panic!("vertex_location is required for vertex deployments in model_list")
|
||||
});
|
||||
validate_gcp_identifier("vertex_project", project);
|
||||
validate_gcp_identifier("vertex_location", location);
|
||||
format!(
|
||||
"https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// LiteLLM Azure model names may carry a route-group marker prefix that selects a
|
||||
// request-shaping path but is not part of the deployment name. This is the complete
|
||||
// set of such markers; they are defined only here (no external producer). Add a
|
||||
// marker here if a new Azure route group is introduced.
|
||||
fn azure_deployment_from_model(model: &str) -> &str {
|
||||
for marker in ["o_series/", "gpt5_series/"] {
|
||||
if let Some(deployment) = model.strip_prefix(marker) {
|
||||
if !deployment.is_empty() {
|
||||
return deployment;
|
||||
}
|
||||
}
|
||||
}
|
||||
model
|
||||
}
|
||||
|
||||
/// Build a BackendConfig from LiteLLM model_list params.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn build_backend_config(
|
||||
name: &str,
|
||||
kind: &BackendKind,
|
||||
provider_id: &str,
|
||||
api_key: &str,
|
||||
base_url: &str,
|
||||
params: &LiteLLMParams,
|
||||
tls: &TlsConfig,
|
||||
log_bodies: bool,
|
||||
config: &LiteLLMConfig,
|
||||
) -> BackendConfig {
|
||||
let backend_auth = match kind {
|
||||
BackendKind::AzureOpenAI => BackendAuth::AzureApiKey(api_key.to_string()),
|
||||
BackendKind::Gemini | BackendKind::Vertex => BackendAuth::GoogleApiKey(api_key.to_string()),
|
||||
BackendKind::Anthropic => BackendAuth::anthropic_from_api_key_like(api_key.to_string()),
|
||||
_ => BackendAuth::BearerToken(api_key.to_string()),
|
||||
};
|
||||
|
||||
// For Azure, resolve_base_url already produced the full deployment URL
|
||||
// (deployment name + api-version), so use it as-is here.
|
||||
let effective_url = if *kind == BackendKind::AzureOpenAI {
|
||||
base_url.to_string()
|
||||
} else {
|
||||
// Validate non-Azure URLs.
|
||||
if *kind != BackendKind::Bedrock {
|
||||
if let Err(e) = validate_base_url(base_url) {
|
||||
panic!("backend '{name}' base_url rejected: {e}");
|
||||
}
|
||||
}
|
||||
base_url.to_string()
|
||||
};
|
||||
|
||||
// Bedrock credentials.
|
||||
let bedrock_credentials = if *kind == BackendKind::Bedrock {
|
||||
let region = params
|
||||
.aws_region_name
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_REGION").ok())
|
||||
.unwrap_or_else(|| "us-east-1".to_string());
|
||||
|
||||
let access_key = params
|
||||
.aws_access_key_id
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_ACCESS_KEY_ID").ok())
|
||||
.unwrap_or_else(|| panic!("backend '{name}': aws_access_key_id required for bedrock"));
|
||||
|
||||
let secret_key = params
|
||||
.aws_secret_access_key
|
||||
.as_deref()
|
||||
.map(|v| resolve_env_value(v).unwrap_or_else(|e| panic!("backend '{name}': {e}")))
|
||||
.or_else(|| std::env::var("AWS_SECRET_ACCESS_KEY").ok())
|
||||
.unwrap_or_else(|| {
|
||||
panic!("backend '{name}': aws_secret_access_key required for bedrock")
|
||||
});
|
||||
|
||||
// Store region as base_url for Bedrock (matches existing convention).
|
||||
// The effective_url was already set to the region string.
|
||||
let _ = region; // region is used as base_url via resolve_base_url
|
||||
|
||||
Some(aws_credential_types::Credentials::new(
|
||||
access_key,
|
||||
secret_key,
|
||||
None, // session token not commonly in LiteLLM configs
|
||||
None,
|
||||
"litellm-config",
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Placeholder model mapping: with model router, these are not used for routing.
|
||||
// They serve as fallback for Anthropic model name translation if needed.
|
||||
let model_mapping = ModelMapping {
|
||||
big_model: String::new(),
|
||||
small_model: String::new(),
|
||||
};
|
||||
|
||||
let _num_retries = config.litellm_settings.as_ref().and_then(|s| s.num_retries);
|
||||
let _request_timeout = config
|
||||
.litellm_settings
|
||||
.as_ref()
|
||||
.and_then(|s| s.request_timeout);
|
||||
|
||||
BackendConfig {
|
||||
kind: kind.clone(),
|
||||
provider_id: Some(provider_id.to_string()),
|
||||
api_key: api_key.to_string(),
|
||||
base_url: effective_url,
|
||||
api_format: OpenAIApiFormat::Chat,
|
||||
model_mapping,
|
||||
tls: tls.clone(),
|
||||
backend_auth,
|
||||
log_bodies,
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: std::env::var("REQUEST_TIMEOUT_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(900u64),
|
||||
bedrock_credentials,
|
||||
// Config-file backends are validated against SSRF at load time; no auto-relax.
|
||||
allow_local_ssrf: false,
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,7 @@
|
||||
use super::parser::parse_provider_model;
|
||||
use super::*;
|
||||
use crate::config::model_router::RoutingStrategy;
|
||||
use crate::config::BackendKind;
|
||||
|
||||
#[test]
|
||||
fn parse_provider_model_openai() {
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
//! Serde structs mapping the LiteLLM `config.yaml` schema, plus the parsed-result
|
||||
//! type returned by the conversion functions in [`super::parser`].
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::config::model_router::ModelRouter;
|
||||
use crate::config::MultiConfig;
|
||||
|
||||
/// Root structure of a LiteLLM `config.yaml` file.
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct LiteLLMConfig {
|
||||
#[serde(default)]
|
||||
pub(super) model_list: Vec<LiteLLMModelEntry>,
|
||||
#[serde(default)]
|
||||
pub(super) litellm_settings: Option<LiteLLMSettings>,
|
||||
#[serde(default)]
|
||||
pub(super) router_settings: Option<RouterSettings>,
|
||||
#[serde(default)]
|
||||
pub(super) general_settings: Option<GeneralSettings>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct LiteLLMModelEntry {
|
||||
pub(super) model_name: String,
|
||||
pub(super) litellm_params: LiteLLMParams,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct LiteLLMParams {
|
||||
pub(super) model: String,
|
||||
pub(super) api_base: Option<String>,
|
||||
pub(super) api_key: Option<String>,
|
||||
pub(super) rpm: Option<u32>,
|
||||
pub(super) tpm: Option<u64>,
|
||||
pub(super) weight: Option<u32>,
|
||||
// Azure-specific
|
||||
pub(super) api_version: Option<String>,
|
||||
// Vertex-specific
|
||||
pub(super) vertex_project: Option<String>,
|
||||
pub(super) vertex_location: Option<String>,
|
||||
// Bedrock-specific
|
||||
pub(super) aws_access_key_id: Option<String>,
|
||||
pub(super) aws_secret_access_key: Option<String>,
|
||||
pub(super) aws_region_name: Option<String>,
|
||||
// Catch unknown fields silently (LiteLLM has many we don't support).
|
||||
#[serde(flatten)]
|
||||
pub(super) _extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct LiteLLMSettings {
|
||||
#[serde(default)]
|
||||
pub(super) num_retries: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub(super) request_timeout: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub(super) callbacks: Vec<String>,
|
||||
#[serde(flatten)]
|
||||
pub(super) _extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct RouterSettings {
|
||||
#[serde(default)]
|
||||
pub(super) routing_strategy: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub(super) _extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct GeneralSettings {
|
||||
pub(super) master_key: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub(super) _extra: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Parsed result from a LiteLLM config file.
|
||||
pub struct LiteLLMParsed {
|
||||
pub multi_config: MultiConfig,
|
||||
pub router: ModelRouter,
|
||||
/// Webhook callback URLs from litellm_settings.callbacks (non-named entries).
|
||||
pub callback_urls: Vec<String>,
|
||||
/// True when "langfuse" appears in litellm_settings.callbacks.
|
||||
pub langfuse_requested: bool,
|
||||
/// Resolved `general_settings.master_key`, if present.
|
||||
/// Caller should apply as PROXY_API_KEYS if that var is not already set.
|
||||
pub master_key: Option<String>,
|
||||
}
|
||||
@@ -346,432 +346,4 @@ fn anthropic_request_requires_tool_choice(req: &anthropic::MessageCreateRequest)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn catalog(json: serde_json::Value) -> ProviderCatalog {
|
||||
ProviderCatalog::from_litellm_json(&json.to_string()).unwrap()
|
||||
}
|
||||
|
||||
fn empty_catalog() -> ProviderCatalog {
|
||||
ProviderCatalog::from_litellm_json("{}").unwrap()
|
||||
}
|
||||
|
||||
fn request(value: serde_json::Value) -> openai::ChatCompletionRequest {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
fn ctx<'a>(
|
||||
backend_kind: BackendKind,
|
||||
provider_id: Option<&'a str>,
|
||||
model: &'a str,
|
||||
provider_catalog: &'a ProviderCatalog,
|
||||
) -> OpenAiToolPolicyContext<'a> {
|
||||
OpenAiToolPolicyContext {
|
||||
backend_kind,
|
||||
provider_id,
|
||||
model,
|
||||
provider_catalog,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_mistral_rewrites_ids_and_matching_tool_results() {
|
||||
let mut req = request(json!({
|
||||
"model": "mistral-large-latest",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{"id": "call_alpha", "type": "function", "function": {"name": "a", "arguments": "{}"}},
|
||||
{"id": "call_alpha", "type": "function", "function": {"name": "b", "arguments": "{}"}}
|
||||
]
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_alpha", "content": "first"},
|
||||
{"role": "tool", "tool_call_id": "call_alpha", "content": "second"}
|
||||
]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let report = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("mistral"),
|
||||
"mistral-large-latest",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let calls = req.messages[0].tool_calls.as_ref().unwrap();
|
||||
assert_eq!(calls[0].id, "000000000");
|
||||
assert_eq!(calls[1].id, "000000001");
|
||||
assert_eq!(req.messages[1].tool_call_id.as_deref(), Some("000000000"));
|
||||
assert_eq!(req.messages[2].tool_call_id.as_deref(), Some("000000001"));
|
||||
assert_eq!(report.duplicate_tool_call_ids, 1);
|
||||
assert_eq!(report.remapped_tool_results, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_gemini_sanitizes_schema_and_removes_strict() {
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"strict": true,
|
||||
"parameters": {
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"q": {"type": "string", "default": ""}
|
||||
}
|
||||
}
|
||||
}
|
||||
}]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::Gemini,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool = req.tools.as_ref().unwrap().first().unwrap();
|
||||
assert_eq!(tool.function.strict, None);
|
||||
let params = tool.function.parameters.as_ref().unwrap();
|
||||
assert!(params.get("$schema").is_none());
|
||||
assert!(params.get("additionalProperties").is_none());
|
||||
assert!(
|
||||
params.pointer("/properties/q/default").is_none(),
|
||||
"Gemini rejects JSON Schema default in nested properties"
|
||||
);
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("tools.function.strict")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_applies_gemini_policy_by_protocol_for_managed_backend() {
|
||||
// A managed/custom backend can run BackendKind::OpenAI while its provider_id
|
||||
// resolves to the Gemini OpenAI shim in the catalog. The policy must key off
|
||||
// the provider protocol, not a hardcoded "gemini" id, so strict-stripping
|
||||
// still applies.
|
||||
let catalog = ProviderCatalog::bundled();
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"strict": true,
|
||||
"parameters": {"type": "object", "additionalProperties": false}
|
||||
}
|
||||
}]
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool = req.tools.as_ref().unwrap().first().unwrap();
|
||||
assert_eq!(tool.function.strict, None);
|
||||
assert!(tool
|
||||
.function
|
||||
.parameters
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.get("additionalProperties")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_gemini_strips_parallel_false_with_multiple_tools() {
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"parallel_tool_calls": false,
|
||||
"tools": [
|
||||
{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}},
|
||||
{"type": "function", "function": {"name": "b", "parameters": {"type": "object"}}}
|
||||
]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::Gemini,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Gemini cannot honor parallel_tool_calls; the field is stripped and the
|
||||
// degradation is reported rather than the request being rejected.
|
||||
assert!(req.parallel_tool_calls.is_none());
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("parallel_tool_calls")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_required_tool_choice_when_model_lacks_it() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "required"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let err = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("tool_choice=required"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_drops_auto_tool_choice_when_model_lacks_it() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "auto"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(req.tool_choice.is_none());
|
||||
assert!(req.tools.is_some());
|
||||
assert_eq!(warnings.as_header_value().as_deref(), Some("tool_choice"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_drops_none_tool_choice_and_tools_without_continuation() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "none"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(req.tool_choice.is_none());
|
||||
assert!(req.tools.is_none());
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("tool_choice, tools")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_tools_when_known_model_lacks_tool_use() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": false,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}]
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let err = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("does not support tools"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_allows_forced_tool_choice_for_self_hosted_provider() {
|
||||
// vllm/lm_studio/llamafile/triton stubs advertise tool_use but a
|
||||
// conservative provider-level tool_choice:false and carry no per-model
|
||||
// metadata. Forced tool_choice must pass through (the backend decides),
|
||||
// not 400 -- the provider-level flag is not authoritative for an unknown
|
||||
// self-hosted model.
|
||||
let catalog = ProviderCatalog::bundled();
|
||||
assert!(
|
||||
catalog.list_models("vllm").is_empty(),
|
||||
"test assumes vllm has no per-model metadata"
|
||||
);
|
||||
let mut req = request(json!({
|
||||
"model": "Qwen/Qwen2.5-7B-Instruct",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "required"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("vllm"),
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
req.tool_choice,
|
||||
Some(ChatToolChoice::Simple(ref v)) if v == "required"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_native_anthropic_tools_when_known_model_lacks_tool_use() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "bedrock",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": false,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let req: anthropic::MessageCreateRequest = serde_json::from_value(json!({
|
||||
"model": "demo-model",
|
||||
"max_tokens": 64,
|
||||
"tools": [{
|
||||
"name": "lookup",
|
||||
"input_schema": {"type": "object"}
|
||||
}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let err = validate_anthropic_tool_request(
|
||||
&req,
|
||||
ctx(
|
||||
BackendKind::Bedrock,
|
||||
Some("bedrock"),
|
||||
"demo-model",
|
||||
&catalog,
|
||||
),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("does not support tools"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_normalizes_streaming_top_level_tool_call_delta() {
|
||||
let chunk = parse_openai_chat_completion_chunk(
|
||||
r#"{
|
||||
"id": "chatcmpl_1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "local",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"name": "lookup",
|
||||
"arguments": {"q":"x"}
|
||||
}]
|
||||
}
|
||||
}]
|
||||
}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool_call = &chunk.choices[0].delta.tool_calls.as_ref().unwrap()[0];
|
||||
assert_eq!(tool_call.call_type.as_deref(), Some("function"));
|
||||
assert_eq!(
|
||||
tool_call.function.as_ref().unwrap().name.as_deref(),
|
||||
Some("lookup")
|
||||
);
|
||||
assert_eq!(
|
||||
tool_call.function.as_ref().unwrap().arguments.as_deref(),
|
||||
Some(r#"{"q":"x"}"#)
|
||||
);
|
||||
}
|
||||
}
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,427 @@
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn catalog(json: serde_json::Value) -> ProviderCatalog {
|
||||
ProviderCatalog::from_litellm_json(&json.to_string()).unwrap()
|
||||
}
|
||||
|
||||
fn empty_catalog() -> ProviderCatalog {
|
||||
ProviderCatalog::from_litellm_json("{}").unwrap()
|
||||
}
|
||||
|
||||
fn request(value: serde_json::Value) -> openai::ChatCompletionRequest {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
fn ctx<'a>(
|
||||
backend_kind: BackendKind,
|
||||
provider_id: Option<&'a str>,
|
||||
model: &'a str,
|
||||
provider_catalog: &'a ProviderCatalog,
|
||||
) -> OpenAiToolPolicyContext<'a> {
|
||||
OpenAiToolPolicyContext {
|
||||
backend_kind,
|
||||
provider_id,
|
||||
model,
|
||||
provider_catalog,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_mistral_rewrites_ids_and_matching_tool_results() {
|
||||
let mut req = request(json!({
|
||||
"model": "mistral-large-latest",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{"id": "call_alpha", "type": "function", "function": {"name": "a", "arguments": "{}"}},
|
||||
{"id": "call_alpha", "type": "function", "function": {"name": "b", "arguments": "{}"}}
|
||||
]
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_alpha", "content": "first"},
|
||||
{"role": "tool", "tool_call_id": "call_alpha", "content": "second"}
|
||||
]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let report = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("mistral"),
|
||||
"mistral-large-latest",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let calls = req.messages[0].tool_calls.as_ref().unwrap();
|
||||
assert_eq!(calls[0].id, "000000000");
|
||||
assert_eq!(calls[1].id, "000000001");
|
||||
assert_eq!(req.messages[1].tool_call_id.as_deref(), Some("000000000"));
|
||||
assert_eq!(req.messages[2].tool_call_id.as_deref(), Some("000000001"));
|
||||
assert_eq!(report.duplicate_tool_call_ids, 1);
|
||||
assert_eq!(report.remapped_tool_results, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_gemini_sanitizes_schema_and_removes_strict() {
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"strict": true,
|
||||
"parameters": {
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"q": {"type": "string", "default": ""}
|
||||
}
|
||||
}
|
||||
}
|
||||
}]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::Gemini,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool = req.tools.as_ref().unwrap().first().unwrap();
|
||||
assert_eq!(tool.function.strict, None);
|
||||
let params = tool.function.parameters.as_ref().unwrap();
|
||||
assert!(params.get("$schema").is_none());
|
||||
assert!(params.get("additionalProperties").is_none());
|
||||
assert!(
|
||||
params.pointer("/properties/q/default").is_none(),
|
||||
"Gemini rejects JSON Schema default in nested properties"
|
||||
);
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("tools.function.strict")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_applies_gemini_policy_by_protocol_for_managed_backend() {
|
||||
// A managed/custom backend can run BackendKind::OpenAI while its provider_id
|
||||
// resolves to the Gemini OpenAI shim in the catalog. The policy must key off
|
||||
// the provider protocol, not a hardcoded "gemini" id, so strict-stripping
|
||||
// still applies.
|
||||
let catalog = ProviderCatalog::bundled();
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"strict": true,
|
||||
"parameters": {"type": "object", "additionalProperties": false}
|
||||
}
|
||||
}]
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool = req.tools.as_ref().unwrap().first().unwrap();
|
||||
assert_eq!(tool.function.strict, None);
|
||||
assert!(tool
|
||||
.function
|
||||
.parameters
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.get("additionalProperties")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_gemini_strips_parallel_false_with_multiple_tools() {
|
||||
let mut req = request(json!({
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"parallel_tool_calls": false,
|
||||
"tools": [
|
||||
{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}},
|
||||
{"type": "function", "function": {"name": "b", "parameters": {"type": "object"}}}
|
||||
]
|
||||
}));
|
||||
let catalog = empty_catalog();
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::Gemini,
|
||||
Some("gemini"),
|
||||
"gemini-2.5-pro",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Gemini cannot honor parallel_tool_calls; the field is stripped and the
|
||||
// degradation is reported rather than the request being rejected.
|
||||
assert!(req.parallel_tool_calls.is_none());
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("parallel_tool_calls")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_required_tool_choice_when_model_lacks_it() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "required"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let err = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("tool_choice=required"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_drops_auto_tool_choice_when_model_lacks_it() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "auto"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(req.tool_choice.is_none());
|
||||
assert!(req.tools.is_some());
|
||||
assert_eq!(warnings.as_header_value().as_deref(), Some("tool_choice"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_drops_none_tool_choice_and_tools_without_continuation() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "none"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(req.tool_choice.is_none());
|
||||
assert!(req.tools.is_none());
|
||||
assert_eq!(
|
||||
warnings.as_header_value().as_deref(),
|
||||
Some("tool_choice, tools")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_tools_when_known_model_lacks_tool_use() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "demo",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": false,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let mut req = request(json!({
|
||||
"model": "demo-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}]
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
let err = prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(BackendKind::OpenAI, Some("demo"), "demo-model", &catalog),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("does not support tools"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_allows_forced_tool_choice_for_self_hosted_provider() {
|
||||
// vllm/lm_studio/llamafile/triton stubs advertise tool_use but a
|
||||
// conservative provider-level tool_choice:false and carry no per-model
|
||||
// metadata. Forced tool_choice must pass through (the backend decides),
|
||||
// not 400 -- the provider-level flag is not authoritative for an unknown
|
||||
// self-hosted model.
|
||||
let catalog = ProviderCatalog::bundled();
|
||||
assert!(
|
||||
catalog.list_models("vllm").is_empty(),
|
||||
"test assumes vllm has no per-model metadata"
|
||||
);
|
||||
let mut req = request(json!({
|
||||
"model": "Qwen/Qwen2.5-7B-Instruct",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "a", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "required"
|
||||
}));
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
prepare_openai_tool_request(
|
||||
&mut req,
|
||||
ctx(
|
||||
BackendKind::OpenAI,
|
||||
Some("vllm"),
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
&catalog,
|
||||
),
|
||||
&mut warnings,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
req.tool_choice,
|
||||
Some(ChatToolChoice::Simple(ref v)) if v == "required"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_rejects_native_anthropic_tools_when_known_model_lacks_tool_use() {
|
||||
let catalog = catalog(json!({
|
||||
"demo-model": {
|
||||
"litellm_provider": "bedrock",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": false,
|
||||
"supports_tool_choice": false
|
||||
}
|
||||
}));
|
||||
let req: anthropic::MessageCreateRequest = serde_json::from_value(json!({
|
||||
"model": "demo-model",
|
||||
"max_tokens": 64,
|
||||
"tools": [{
|
||||
"name": "lookup",
|
||||
"input_schema": {"type": "object"}
|
||||
}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let err = validate_anthropic_tool_request(
|
||||
&req,
|
||||
ctx(
|
||||
BackendKind::Bedrock,
|
||||
Some("bedrock"),
|
||||
"demo-model",
|
||||
&catalog,
|
||||
),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.message().contains("does not support tools"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_tool_policy_normalizes_streaming_top_level_tool_call_delta() {
|
||||
let chunk = parse_openai_chat_completion_chunk(
|
||||
r#"{
|
||||
"id": "chatcmpl_1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "local",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"name": "lookup",
|
||||
"arguments": {"q":"x"}
|
||||
}]
|
||||
}
|
||||
}]
|
||||
}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tool_call = &chunk.choices[0].delta.tool_calls.as_ref().unwrap()[0];
|
||||
assert_eq!(tool_call.call_type.as_deref(), Some("function"));
|
||||
assert_eq!(
|
||||
tool_call.function.as_ref().unwrap().name.as_deref(),
|
||||
Some("lookup")
|
||||
);
|
||||
assert_eq!(
|
||||
tool_call.function.as_ref().unwrap().arguments.as_deref(),
|
||||
Some(r#"{"q":"x"}"#)
|
||||
);
|
||||
}
|
||||
@@ -435,249 +435,5 @@ pub fn resolve_runtime_optimizer_locked(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn long_text() -> String {
|
||||
"The quick brown fox jumps over the lazy dog again and again across the wide \
|
||||
green field toward the distant blue mountains far beyond the winding river."
|
||||
.repeat(2)
|
||||
}
|
||||
|
||||
fn long_openai_body() -> Value {
|
||||
let long = long_text();
|
||||
let mut messages = vec![json!({"role":"system","content":"you are helpful"})];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
json!({"model":"gpt-4o","messages": messages})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_default_mode_unset_is_off() {
|
||||
// Deliberately does not touch process env: relies on OPTIMIZER_MODE being
|
||||
// absent by default in the test environment.
|
||||
if std::env::var("OPTIMIZER_MODE").is_err() {
|
||||
assert_eq!(resolve_default_mode(), Mode::Off);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn off_mode_never_mutates_and_reports_off() {
|
||||
let mut body = long_openai_body();
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before, "off mode must not mutate the body");
|
||||
assert_eq!(report.mode, Mode::Off);
|
||||
assert!(!report.applied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_never_mutates_but_reports_savings() {
|
||||
let mut body = long_openai_body();
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before, "shadow mode must not mutate the body");
|
||||
assert!(
|
||||
report.removed_tokens_est > 0,
|
||||
"shadow should still report would-be savings on a long convo"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_openai_history_and_preserves_latest_and_system() {
|
||||
let mut body = long_openai_body();
|
||||
let orig_msgs = body["messages"].as_array().unwrap().clone();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert!(report.applied || report.rewrite_suffix_tokens == 0);
|
||||
|
||||
let new_msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(orig_msgs.len(), new_msgs.len());
|
||||
assert_eq!(orig_msgs[0], new_msgs[0], "system untouched");
|
||||
assert_eq!(
|
||||
orig_msgs.last().unwrap(),
|
||||
new_msgs.last().unwrap(),
|
||||
"latest message untouched"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_anthropic_history_and_preserves_latest() {
|
||||
let long = long_text();
|
||||
let mut messages = vec![];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
let mut body = json!({
|
||||
"model":"claude-sonnet-5",
|
||||
"system":"you are helpful",
|
||||
"messages": messages,
|
||||
});
|
||||
let orig_msgs = body["messages"].as_array().unwrap().clone();
|
||||
let orig_system = body["system"].clone();
|
||||
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_anthropic(&mut body, "messages");
|
||||
|
||||
assert_eq!(body["system"], orig_system, "system field untouched");
|
||||
let new_msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(orig_msgs.len(), new_msgs.len());
|
||||
assert_eq!(
|
||||
orig_msgs.last().unwrap(),
|
||||
new_msgs.last().unwrap(),
|
||||
"latest message untouched"
|
||||
);
|
||||
// Live mode must place the deepest cache breakpoint at the frontier, not just
|
||||
// compress text (crates/optimizer/CLAUDE.md checklist item 5).
|
||||
assert!(
|
||||
report.frontier > 0,
|
||||
"long history must have a nonzero frontier"
|
||||
);
|
||||
let bp_idx = report.frontier - 1;
|
||||
let bp_msg = &new_msgs[bp_idx];
|
||||
let has_marker = bp_msg["content"]
|
||||
.as_array()
|
||||
.is_some_and(|arr| arr.iter().any(|b| b.get("cache_control").is_some()));
|
||||
assert!(
|
||||
has_marker,
|
||||
"expected a cache_control breakpoint on message {bp_idx}, got {bp_msg}"
|
||||
);
|
||||
}
|
||||
|
||||
fn long_anthropic_body() -> Value {
|
||||
let long = long_text();
|
||||
let mut messages = vec![];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
json!({
|
||||
"model":"claude-sonnet-5",
|
||||
"system":"you are helpful",
|
||||
"messages": messages,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_live_shrinks_and_keeps_cache_control() {
|
||||
let body = Bytes::from(serde_json::to_vec(&long_anthropic_body()).unwrap());
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert!(out.len() < body.len(), "live output must be smaller");
|
||||
// The frontier cache_control breakpoint must survive to the wire bytes --
|
||||
// the whole point of doing this on Bytes instead of the typed round-trip.
|
||||
let root: Value = serde_json::from_slice(&out).unwrap();
|
||||
let has_marker = root["messages"].as_array().unwrap().iter().any(|m| {
|
||||
m["content"]
|
||||
.as_array()
|
||||
.is_some_and(|arr| arr.iter().any(|b| b.get("cache_control").is_some()))
|
||||
});
|
||||
assert!(
|
||||
has_marker,
|
||||
"cache_control breakpoint dropped from wire bytes"
|
||||
);
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_off_is_noop() {
|
||||
let body = Bytes::from(serde_json::to_vec(&long_anthropic_body()).unwrap());
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert_eq!(out, body, "off mode must return the body unchanged");
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_fails_open_on_garbage() {
|
||||
let body = Bytes::from_static(b"not json");
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert_eq!(out, body, "garbage body returned unchanged");
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_mode_override_prefers_override() {
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let overridden = engine.with_mode_override("live");
|
||||
assert_eq!(overridden.policy.mode, Mode::Live);
|
||||
// The original engine is untouched.
|
||||
assert_eq!(engine.policy.mode, Mode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_mode_override_falls_back_on_unparseable_value() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let overridden = engine.with_mode_override("not-a-real-mode");
|
||||
assert_eq!(overridden.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_prefers_runtime_override() {
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "live");
|
||||
assert_eq!(resolved.policy.mode, Mode::Live);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_keeps_static_when_modes_match() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "shadow");
|
||||
assert_eq!(resolved.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_falls_back_on_unparseable_value() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "not-a-real-mode");
|
||||
assert_eq!(resolved.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_precedence_wins_over_runtime_and_static() {
|
||||
// Full three-tier precedence, mirroring `effective_optimizer`'s own logic:
|
||||
// static (env-seeded) engine mode is Off, runtime tier says Shadow, but a
|
||||
// route-level override of Live must win over both.
|
||||
let static_engine = OptimizerEngine::new(Mode::Off);
|
||||
let runtime_tier = resolve_runtime_optimizer(&static_engine, "shadow");
|
||||
assert_eq!(
|
||||
runtime_tier.policy.mode,
|
||||
Mode::Shadow,
|
||||
"runtime beats static"
|
||||
);
|
||||
|
||||
// Route override is applied directly against the static engine (as
|
||||
// `effective_optimizer` does), bypassing the runtime tier entirely.
|
||||
let route_tier = static_engine.with_mode_override("live");
|
||||
assert_eq!(
|
||||
route_tier.policy.mode,
|
||||
Mode::Live,
|
||||
"route beats runtime and static"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fails_open_on_malformed_messages_field() {
|
||||
// "messages" is a string, not an array — adapters degrade to an empty
|
||||
// Conversation rather than panicking, so this must not mutate or panic.
|
||||
let mut body = json!({"model":"gpt-4o","messages":"not an array"});
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before);
|
||||
assert!(!report.applied);
|
||||
}
|
||||
}
|
||||
#[path = "optimizer/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn long_text() -> String {
|
||||
"The quick brown fox jumps over the lazy dog again and again across the wide \
|
||||
green field toward the distant blue mountains far beyond the winding river."
|
||||
.repeat(2)
|
||||
}
|
||||
|
||||
fn long_openai_body() -> Value {
|
||||
let long = long_text();
|
||||
let mut messages = vec![json!({"role":"system","content":"you are helpful"})];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
json!({"model":"gpt-4o","messages": messages})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_default_mode_unset_is_off() {
|
||||
// Deliberately does not touch process env: relies on OPTIMIZER_MODE being
|
||||
// absent by default in the test environment.
|
||||
if std::env::var("OPTIMIZER_MODE").is_err() {
|
||||
assert_eq!(resolve_default_mode(), Mode::Off);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn off_mode_never_mutates_and_reports_off() {
|
||||
let mut body = long_openai_body();
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before, "off mode must not mutate the body");
|
||||
assert_eq!(report.mode, Mode::Off);
|
||||
assert!(!report.applied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_never_mutates_but_reports_savings() {
|
||||
let mut body = long_openai_body();
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before, "shadow mode must not mutate the body");
|
||||
assert!(
|
||||
report.removed_tokens_est > 0,
|
||||
"shadow should still report would-be savings on a long convo"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_openai_history_and_preserves_latest_and_system() {
|
||||
let mut body = long_openai_body();
|
||||
let orig_msgs = body["messages"].as_array().unwrap().clone();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert!(report.applied || report.rewrite_suffix_tokens == 0);
|
||||
|
||||
let new_msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(orig_msgs.len(), new_msgs.len());
|
||||
assert_eq!(orig_msgs[0], new_msgs[0], "system untouched");
|
||||
assert_eq!(
|
||||
orig_msgs.last().unwrap(),
|
||||
new_msgs.last().unwrap(),
|
||||
"latest message untouched"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_anthropic_history_and_preserves_latest() {
|
||||
let long = long_text();
|
||||
let mut messages = vec![];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
let mut body = json!({
|
||||
"model":"claude-sonnet-5",
|
||||
"system":"you are helpful",
|
||||
"messages": messages,
|
||||
});
|
||||
let orig_msgs = body["messages"].as_array().unwrap().clone();
|
||||
let orig_system = body["system"].clone();
|
||||
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_anthropic(&mut body, "messages");
|
||||
|
||||
assert_eq!(body["system"], orig_system, "system field untouched");
|
||||
let new_msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(orig_msgs.len(), new_msgs.len());
|
||||
assert_eq!(
|
||||
orig_msgs.last().unwrap(),
|
||||
new_msgs.last().unwrap(),
|
||||
"latest message untouched"
|
||||
);
|
||||
// Live mode must place the deepest cache breakpoint at the frontier, not just
|
||||
// compress text (crates/optimizer/CLAUDE.md checklist item 5).
|
||||
assert!(
|
||||
report.frontier > 0,
|
||||
"long history must have a nonzero frontier"
|
||||
);
|
||||
let bp_idx = report.frontier - 1;
|
||||
let bp_msg = &new_msgs[bp_idx];
|
||||
let has_marker = bp_msg["content"]
|
||||
.as_array()
|
||||
.is_some_and(|arr| arr.iter().any(|b| b.get("cache_control").is_some()));
|
||||
assert!(
|
||||
has_marker,
|
||||
"expected a cache_control breakpoint on message {bp_idx}, got {bp_msg}"
|
||||
);
|
||||
}
|
||||
|
||||
fn long_anthropic_body() -> Value {
|
||||
let long = long_text();
|
||||
let mut messages = vec![];
|
||||
for _ in 0..16 {
|
||||
messages.push(json!({"role":"user","content": long}));
|
||||
messages.push(json!({"role":"assistant","content": long}));
|
||||
}
|
||||
messages.push(json!({"role":"user","content":"what is the latest?"}));
|
||||
json!({
|
||||
"model":"claude-sonnet-5",
|
||||
"system":"you are helpful",
|
||||
"messages": messages,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_live_shrinks_and_keeps_cache_control() {
|
||||
let body = Bytes::from(serde_json::to_vec(&long_anthropic_body()).unwrap());
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert!(out.len() < body.len(), "live output must be smaller");
|
||||
// The frontier cache_control breakpoint must survive to the wire bytes --
|
||||
// the whole point of doing this on Bytes instead of the typed round-trip.
|
||||
let root: Value = serde_json::from_slice(&out).unwrap();
|
||||
let has_marker = root["messages"].as_array().unwrap().iter().any(|m| {
|
||||
m["content"]
|
||||
.as_array()
|
||||
.is_some_and(|arr| arr.iter().any(|b| b.get("cache_control").is_some()))
|
||||
});
|
||||
assert!(
|
||||
has_marker,
|
||||
"cache_control breakpoint dropped from wire bytes"
|
||||
);
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_off_is_noop() {
|
||||
let body = Bytes::from(serde_json::to_vec(&long_anthropic_body()).unwrap());
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert_eq!(out, body, "off mode must return the body unchanged");
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optimize_anthropic_bytes_fails_open_on_garbage() {
|
||||
let body = Bytes::from_static(b"not json");
|
||||
let metrics = crate::metrics::Metrics::new();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let out = engine.optimize_anthropic_bytes(body.clone(), "messages", &metrics);
|
||||
assert_eq!(out, body, "garbage body returned unchanged");
|
||||
assert_eq!(metrics.snapshot().optimizer_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_mode_override_prefers_override() {
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let overridden = engine.with_mode_override("live");
|
||||
assert_eq!(overridden.policy.mode, Mode::Live);
|
||||
// The original engine is untouched.
|
||||
assert_eq!(engine.policy.mode, Mode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_mode_override_falls_back_on_unparseable_value() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let overridden = engine.with_mode_override("not-a-real-mode");
|
||||
assert_eq!(overridden.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_prefers_runtime_override() {
|
||||
let engine = OptimizerEngine::new(Mode::Off);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "live");
|
||||
assert_eq!(resolved.policy.mode, Mode::Live);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_keeps_static_when_modes_match() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "shadow");
|
||||
assert_eq!(resolved.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_optimizer_falls_back_on_unparseable_value() {
|
||||
let engine = OptimizerEngine::new(Mode::Shadow);
|
||||
let resolved = resolve_runtime_optimizer(&engine, "not-a-real-mode");
|
||||
assert_eq!(resolved.policy.mode, Mode::Shadow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_precedence_wins_over_runtime_and_static() {
|
||||
// Full three-tier precedence, mirroring `effective_optimizer`'s own logic:
|
||||
// static (env-seeded) engine mode is Off, runtime tier says Shadow, but a
|
||||
// route-level override of Live must win over both.
|
||||
let static_engine = OptimizerEngine::new(Mode::Off);
|
||||
let runtime_tier = resolve_runtime_optimizer(&static_engine, "shadow");
|
||||
assert_eq!(
|
||||
runtime_tier.policy.mode,
|
||||
Mode::Shadow,
|
||||
"runtime beats static"
|
||||
);
|
||||
|
||||
// Route override is applied directly against the static engine (as
|
||||
// `effective_optimizer` does), bypassing the runtime tier entirely.
|
||||
let route_tier = static_engine.with_mode_override("live");
|
||||
assert_eq!(
|
||||
route_tier.policy.mode,
|
||||
Mode::Live,
|
||||
"route beats runtime and static"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fails_open_on_malformed_messages_field() {
|
||||
// "messages" is a string, not an array — adapters degrade to an empty
|
||||
// Conversation rather than panicking, so this must not mutate or panic.
|
||||
let mut body = json!({"model":"gpt-4o","messages":"not an array"});
|
||||
let before = body.clone();
|
||||
let engine = OptimizerEngine::new(Mode::Live);
|
||||
let report = engine.optimize_openai(&mut body, "chat_completions");
|
||||
assert_eq!(body, before);
|
||||
assert!(!report.applied);
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
//! Runtime-specific error type (no axum response types).
|
||||
|
||||
use crate::backend::BackendError;
|
||||
use crate::config::BackendKind;
|
||||
use std::fmt;
|
||||
|
||||
/// Runtime errors that do not expose axum response types.
|
||||
#[derive(Debug)]
|
||||
pub enum ChatCompletionError {
|
||||
InvalidRequest(String),
|
||||
Translation(anyllm_translate::TranslateError),
|
||||
Routing(String),
|
||||
UnsupportedBackend {
|
||||
backend_name: String,
|
||||
backend_kind: BackendKind,
|
||||
},
|
||||
Backend(BackendError),
|
||||
StreamRead(String),
|
||||
StreamParse(String),
|
||||
StreamBufferOverflow,
|
||||
StreamTimeout,
|
||||
}
|
||||
|
||||
impl ChatCompletionError {
|
||||
pub fn status_code(&self) -> u16 {
|
||||
match self {
|
||||
Self::InvalidRequest(_) | Self::Translation(_) | Self::UnsupportedBackend { .. } => 400,
|
||||
Self::Routing(_) => 429,
|
||||
Self::Backend(e) => e.status_code(),
|
||||
Self::StreamRead(_)
|
||||
| Self::StreamParse(_)
|
||||
| Self::StreamBufferOverflow
|
||||
| Self::StreamTimeout => 502,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ChatCompletionError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::InvalidRequest(msg) => write!(f, "{msg}"),
|
||||
Self::Translation(e) => write!(f, "{e}"),
|
||||
Self::Routing(msg) => write!(f, "{msg}"),
|
||||
Self::UnsupportedBackend {
|
||||
backend_name,
|
||||
backend_kind,
|
||||
} => write!(
|
||||
f,
|
||||
"backend '{backend_name}' ({backend_kind:?}) does not support Chat Completions runtime"
|
||||
),
|
||||
Self::Backend(e) => write!(f, "{e}"),
|
||||
Self::StreamRead(e) => write!(f, "stream read error: {e}"),
|
||||
Self::StreamParse(e) => write!(f, "stream parse error: {e}"),
|
||||
Self::StreamBufferOverflow => write!(f, "SSE buffer exceeded maximum size"),
|
||||
Self::StreamTimeout => write!(f, "stream exceeded wall-clock timeout"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ChatCompletionError {}
|
||||
|
||||
impl From<anyllm_translate::TranslateError> for ChatCompletionError {
|
||||
fn from(e: anyllm_translate::TranslateError) -> Self {
|
||||
Self::Translation(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BackendError> for ChatCompletionError {
|
||||
fn from(e: BackendError) -> Self {
|
||||
Self::Backend(e)
|
||||
}
|
||||
}
|
||||
+14
-566
@@ -2,572 +2,20 @@
|
||||
//!
|
||||
//! This module exposes model routing and backend dispatch without taking
|
||||
//! ownership of HTTP routing, auth, admin UI, caching, or tool execution.
|
||||
//!
|
||||
//! - [`error`] runtime error type
|
||||
//! - [`types`] result / metadata / service-trait types
|
||||
//! - [`service`] the `ChatCompletionRuntime` implementation
|
||||
//! - [`stream`] SSE chunk stream adapters
|
||||
|
||||
use crate::backend::{BackendClient, BackendError, RateLimitHeaders};
|
||||
use crate::config::{BackendKind, Config, ModelMapping, MultiConfig, OpenAIApiFormat};
|
||||
use crate::openai_tool_policy::{prepare_openai_tool_request, OpenAiToolPolicyContext};
|
||||
use anyllm_providers::ProviderCatalog;
|
||||
use anyllm_translate::{
|
||||
mapping, openai, translate_anthropic_to_openai_response, translate_openai_to_anthropic_request,
|
||||
TranslationWarnings,
|
||||
};
|
||||
use futures::{
|
||||
future::{BoxFuture, FutureExt},
|
||||
Stream,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Stream returned by [`ChatCompletionRuntime::complete_stream`].
|
||||
pub type ChatCompletionChunkStream =
|
||||
Pin<Box<dyn Stream<Item = Result<openai::ChatCompletionChunk, ChatCompletionError>> + Send>>;
|
||||
|
||||
/// Object-safe service API for one Chat Completions call at a time.
|
||||
pub trait ChatCompletionService: Send + Sync {
|
||||
fn complete<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionResult, ChatCompletionError>>;
|
||||
|
||||
fn complete_stream<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionStreamResult, ChatCompletionError>>;
|
||||
}
|
||||
|
||||
/// Non-streaming runtime response.
|
||||
#[derive(Debug)]
|
||||
pub struct ChatCompletionResult {
|
||||
pub response: openai::ChatCompletionResponse,
|
||||
pub usage: Option<openai::ChatUsage>,
|
||||
pub rate_limits: RateLimitHeaders,
|
||||
pub metadata: ChatCompletionMetadata,
|
||||
pub warnings: TranslationWarnings,
|
||||
}
|
||||
|
||||
/// Streaming runtime response.
|
||||
pub struct ChatCompletionStreamResult {
|
||||
pub chunks: ChatCompletionChunkStream,
|
||||
pub rate_limits: RateLimitHeaders,
|
||||
pub metadata: ChatCompletionMetadata,
|
||||
pub warnings: TranslationWarnings,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ChatCompletionStreamResult {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("ChatCompletionStreamResult")
|
||||
.field("chunks", &"<stream>")
|
||||
.field("rate_limits", &self.rate_limits)
|
||||
.field("metadata", &self.metadata)
|
||||
.field("warnings", &self.warnings)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Backend selection metadata for observability and callers that need accounting.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatCompletionMetadata {
|
||||
pub requested_model: String,
|
||||
pub selected_backend: String,
|
||||
pub mapped_model: String,
|
||||
pub backend_kind: BackendKind,
|
||||
pub provider_id: Option<String>,
|
||||
pub api_format: OpenAIApiFormat,
|
||||
pub used_responses_api: bool,
|
||||
}
|
||||
|
||||
/// Runtime errors that do not expose axum response types.
|
||||
#[derive(Debug)]
|
||||
pub enum ChatCompletionError {
|
||||
InvalidRequest(String),
|
||||
Translation(anyllm_translate::TranslateError),
|
||||
Routing(String),
|
||||
UnsupportedBackend {
|
||||
backend_name: String,
|
||||
backend_kind: BackendKind,
|
||||
},
|
||||
Backend(BackendError),
|
||||
StreamRead(String),
|
||||
StreamParse(String),
|
||||
StreamBufferOverflow,
|
||||
StreamTimeout,
|
||||
}
|
||||
|
||||
impl ChatCompletionError {
|
||||
pub fn status_code(&self) -> u16 {
|
||||
match self {
|
||||
Self::InvalidRequest(_) | Self::Translation(_) | Self::UnsupportedBackend { .. } => 400,
|
||||
Self::Routing(_) => 429,
|
||||
Self::Backend(e) => e.status_code(),
|
||||
Self::StreamRead(_)
|
||||
| Self::StreamParse(_)
|
||||
| Self::StreamBufferOverflow
|
||||
| Self::StreamTimeout => 502,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ChatCompletionError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::InvalidRequest(msg) => write!(f, "{msg}"),
|
||||
Self::Translation(e) => write!(f, "{e}"),
|
||||
Self::Routing(msg) => write!(f, "{msg}"),
|
||||
Self::UnsupportedBackend {
|
||||
backend_name,
|
||||
backend_kind,
|
||||
} => write!(
|
||||
f,
|
||||
"backend '{backend_name}' ({backend_kind:?}) does not support Chat Completions runtime"
|
||||
),
|
||||
Self::Backend(e) => write!(f, "{e}"),
|
||||
Self::StreamRead(e) => write!(f, "stream read error: {e}"),
|
||||
Self::StreamParse(e) => write!(f, "stream parse error: {e}"),
|
||||
Self::StreamBufferOverflow => write!(f, "SSE buffer exceeded maximum size"),
|
||||
Self::StreamTimeout => write!(f, "stream exceeded wall-clock timeout"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ChatCompletionError {}
|
||||
|
||||
impl From<anyllm_translate::TranslateError> for ChatCompletionError {
|
||||
fn from(e: anyllm_translate::TranslateError) -> Self {
|
||||
Self::Translation(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BackendError> for ChatCompletionError {
|
||||
fn from(e: BackendError) -> Self {
|
||||
Self::Backend(e)
|
||||
}
|
||||
}
|
||||
|
||||
/// Chat Completions runtime built from proxy backend configuration.
|
||||
#[derive(Clone)]
|
||||
pub struct ChatCompletionRuntime {
|
||||
default_backend: String,
|
||||
backends: Arc<HashMap<String, RuntimeBackend>>,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
provider_catalog: Arc<ProviderCatalog>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RuntimeBackend {
|
||||
backend: BackendClient,
|
||||
backend_name: String,
|
||||
model_mapping: ModelMapping,
|
||||
backend_kind: BackendKind,
|
||||
api_format: OpenAIApiFormat,
|
||||
omit_stream_options: bool,
|
||||
stream_timeout_secs: u64,
|
||||
provider_id: Option<String>,
|
||||
}
|
||||
|
||||
struct ResolvedBackend {
|
||||
state: RuntimeBackend,
|
||||
mapped_model: String,
|
||||
deployment: Option<Arc<crate::config::model_router::Deployment>>,
|
||||
}
|
||||
|
||||
impl ChatCompletionRuntime {
|
||||
/// Build a runtime from a legacy single-backend config.
|
||||
pub fn from_config(config: Config) -> Self {
|
||||
let multi = MultiConfig::from_single_config(&config);
|
||||
let default_backend = multi.default_backend.clone();
|
||||
let (_, bc) = multi
|
||||
.backends
|
||||
.get_key_value(&default_backend)
|
||||
.expect("wrapped config must contain default backend");
|
||||
|
||||
let backend = if config.backend == BackendKind::Bedrock {
|
||||
BackendClient::from_backend_config(bc)
|
||||
} else {
|
||||
BackendClient::new(&config)
|
||||
};
|
||||
|
||||
let mut backends = HashMap::new();
|
||||
backends.insert(
|
||||
default_backend.clone(),
|
||||
RuntimeBackend {
|
||||
backend,
|
||||
backend_name: default_backend.clone(),
|
||||
model_mapping: bc.model_mapping.clone(),
|
||||
backend_kind: bc.kind.clone(),
|
||||
api_format: bc.api_format.clone(),
|
||||
omit_stream_options: bc.omit_stream_options,
|
||||
stream_timeout_secs: bc.stream_timeout_secs,
|
||||
provider_id: config.provider_id.clone(),
|
||||
},
|
||||
);
|
||||
|
||||
Self {
|
||||
default_backend,
|
||||
backends: Arc::new(backends),
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a runtime from multi-backend config without model-list routing.
|
||||
pub fn from_multi_config(config: MultiConfig) -> Self {
|
||||
Self::from_multi_config_with_model_router(config, None)
|
||||
}
|
||||
|
||||
/// Build a runtime from multi-backend config and an optional model router.
|
||||
pub fn from_multi_config_with_model_router(
|
||||
config: MultiConfig,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
) -> Self {
|
||||
let mut backends = HashMap::new();
|
||||
for (name, bc) in &config.backends {
|
||||
backends.insert(
|
||||
name.clone(),
|
||||
RuntimeBackend {
|
||||
backend: BackendClient::from_backend_config(bc),
|
||||
backend_name: name.clone(),
|
||||
model_mapping: bc.model_mapping.clone(),
|
||||
backend_kind: bc.kind.clone(),
|
||||
api_format: bc.api_format.clone(),
|
||||
omit_stream_options: bc.omit_stream_options,
|
||||
stream_timeout_secs: bc.stream_timeout_secs,
|
||||
provider_id: bc.provider_id.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Self {
|
||||
default_backend: config.default_backend,
|
||||
backends: Arc::new(backends),
|
||||
model_router,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve(&self, model: &str) -> Result<ResolvedBackend, ChatCompletionError> {
|
||||
if let Some(ref router_lock) = self.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(routed) = router.route(model) {
|
||||
let state = self
|
||||
.backends
|
||||
.get(routed.backend_name)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ChatCompletionError::Routing(format!(
|
||||
"model '{model}' routed to unknown backend '{}'",
|
||||
routed.backend_name
|
||||
))
|
||||
})?;
|
||||
return Ok(ResolvedBackend {
|
||||
state,
|
||||
mapped_model: routed.actual_model.to_string(),
|
||||
deployment: Some(routed.deployment.clone()),
|
||||
});
|
||||
}
|
||||
if router.has_model(model) {
|
||||
return Err(ChatCompletionError::Routing(
|
||||
"all deployments for this model are at their RPM limit".to_string(),
|
||||
));
|
||||
}
|
||||
return Err(ChatCompletionError::InvalidRequest(format!(
|
||||
"model '{model}' is not configured in model_list"
|
||||
)));
|
||||
}
|
||||
|
||||
let state = self
|
||||
.backends
|
||||
.get(&self.default_backend)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ChatCompletionError::Routing(format!(
|
||||
"default backend '{}' is not configured",
|
||||
self.default_backend
|
||||
))
|
||||
})?;
|
||||
Ok(ResolvedBackend {
|
||||
mapped_model: state.model_mapping.map_model(model),
|
||||
state,
|
||||
deployment: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionService for ChatCompletionRuntime {
|
||||
fn complete<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionResult, ChatCompletionError>> {
|
||||
async move { self.complete_inner(req).await }.boxed()
|
||||
}
|
||||
|
||||
fn complete_stream<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionStreamResult, ChatCompletionError>> {
|
||||
async move { self.complete_stream_inner(req).await }.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionRuntime {
|
||||
async fn complete_inner(
|
||||
&self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> Result<ChatCompletionResult, ChatCompletionError> {
|
||||
if req.stream == Some(true) {
|
||||
return Err(ChatCompletionError::InvalidRequest(
|
||||
"complete does not accept stream=true; use complete_stream".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let requested_model = req.model.clone();
|
||||
let resolved = self.resolve(&requested_model)?;
|
||||
let metadata = metadata(&requested_model, &resolved);
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
match &resolved.state.backend {
|
||||
BackendClient::OpenAI(client)
|
||||
| BackendClient::AzureOpenAI(client)
|
||||
| BackendClient::Vertex(client)
|
||||
| BackendClient::GeminiOpenAI(client) => {
|
||||
let mut openai_req = req;
|
||||
prepare_openai_request(
|
||||
&mut openai_req,
|
||||
&resolved,
|
||||
false,
|
||||
&mut warnings,
|
||||
&self.provider_catalog,
|
||||
)?;
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.chat_completion(&openai_req).await {
|
||||
Ok((response, _status, rate_limits)) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
if let Some(ref deployment) = resolved.deployment {
|
||||
if let Some(ref usage) = response.usage {
|
||||
deployment.record_tokens(usage.total_tokens as u64);
|
||||
}
|
||||
}
|
||||
Ok(ChatCompletionResult {
|
||||
usage: response.usage.clone(),
|
||||
response,
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::OpenAIResponses(client) => {
|
||||
let anthropic_req = translate_openai_to_anthropic_request(&req, &mut warnings)?;
|
||||
let mut responses_req =
|
||||
mapping::responses_message_map::anthropic_to_responses_request(&anthropic_req);
|
||||
responses_req.model = resolved.mapped_model.clone();
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.responses(&responses_req).await {
|
||||
Ok((resp, _status, rate_limits)) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
let anthropic_resp =
|
||||
mapping::responses_message_map::responses_to_anthropic_response(
|
||||
&resp,
|
||||
&requested_model,
|
||||
);
|
||||
if let Some(ref deployment) = resolved.deployment {
|
||||
deployment.record_tokens(
|
||||
anthropic_resp.usage.input_tokens as u64
|
||||
+ anthropic_resp.usage.output_tokens as u64,
|
||||
);
|
||||
}
|
||||
let response = translate_anthropic_to_openai_response(
|
||||
&anthropic_resp,
|
||||
&requested_model,
|
||||
);
|
||||
Ok(ChatCompletionResult {
|
||||
usage: response.usage.clone(),
|
||||
response,
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::Anthropic(_)
|
||||
| BackendClient::Bedrock(_)
|
||||
| BackendClient::GeminiNative(_) => Err(ChatCompletionError::UnsupportedBackend {
|
||||
backend_name: resolved.state.backend_name,
|
||||
backend_kind: resolved.state.backend_kind,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn complete_stream_inner(
|
||||
&self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> Result<ChatCompletionStreamResult, ChatCompletionError> {
|
||||
let requested_model = req.model.clone();
|
||||
let resolved = self.resolve(&requested_model)?;
|
||||
let metadata = metadata(&requested_model, &resolved);
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
match &resolved.state.backend {
|
||||
BackendClient::OpenAI(client)
|
||||
| BackendClient::AzureOpenAI(client)
|
||||
| BackendClient::Vertex(client)
|
||||
| BackendClient::GeminiOpenAI(client) => {
|
||||
let mut openai_req = req;
|
||||
prepare_openai_request(
|
||||
&mut openai_req,
|
||||
&resolved,
|
||||
true,
|
||||
&mut warnings,
|
||||
&self.provider_catalog,
|
||||
)?;
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.chat_completion_stream(&openai_req).await {
|
||||
Ok((response, rate_limits)) => Ok(ChatCompletionStreamResult {
|
||||
chunks: openai_chunk_stream(
|
||||
response,
|
||||
resolved.state.stream_timeout_secs,
|
||||
DeploymentLatencyGuard::from_started(
|
||||
resolved.deployment.clone(),
|
||||
start,
|
||||
),
|
||||
),
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
}),
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::OpenAIResponses(client) => {
|
||||
let anthropic_req = translate_openai_to_anthropic_request(&req, &mut warnings)?;
|
||||
let mut responses_req =
|
||||
mapping::responses_message_map::anthropic_to_responses_request(&anthropic_req);
|
||||
responses_req.model = resolved.mapped_model.clone();
|
||||
responses_req.stream = Some(true);
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.responses_stream(&responses_req).await {
|
||||
Ok((response, rate_limits)) => Ok(ChatCompletionStreamResult {
|
||||
chunks: responses_chunk_stream(
|
||||
response,
|
||||
requested_model,
|
||||
resolved.state.stream_timeout_secs,
|
||||
DeploymentLatencyGuard::from_started(
|
||||
resolved.deployment.clone(),
|
||||
start,
|
||||
),
|
||||
),
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
}),
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::Anthropic(_)
|
||||
| BackendClient::Bedrock(_)
|
||||
| BackendClient::GeminiNative(_) => Err(ChatCompletionError::UnsupportedBackend {
|
||||
backend_name: resolved.state.backend_name,
|
||||
backend_kind: resolved.state.backend_kind,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata(requested_model: &str, resolved: &ResolvedBackend) -> ChatCompletionMetadata {
|
||||
ChatCompletionMetadata {
|
||||
requested_model: requested_model.to_string(),
|
||||
selected_backend: resolved.state.backend_name.clone(),
|
||||
mapped_model: resolved.mapped_model.clone(),
|
||||
backend_kind: resolved.state.backend_kind.clone(),
|
||||
provider_id: resolved.state.provider_id.clone(),
|
||||
api_format: resolved.state.api_format.clone(),
|
||||
used_responses_api: matches!(resolved.state.backend, BackendClient::OpenAIResponses(_)),
|
||||
}
|
||||
}
|
||||
|
||||
fn record_start(
|
||||
deployment: &Option<Arc<crate::config::model_router::Deployment>>,
|
||||
) -> Option<Instant> {
|
||||
if let Some(d) = deployment {
|
||||
d.record_start();
|
||||
Some(Instant::now())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn record_finish(
|
||||
deployment: &Option<Arc<crate::config::model_router::Deployment>>,
|
||||
start: Option<Instant>,
|
||||
) {
|
||||
if let (Some(d), Some(start)) = (deployment, start) {
|
||||
d.record_finish(start.elapsed().as_millis() as u64);
|
||||
}
|
||||
}
|
||||
|
||||
mod error;
|
||||
mod service;
|
||||
pub mod stream;
|
||||
use stream::{openai_chunk_stream, responses_chunk_stream, DeploymentLatencyGuard};
|
||||
mod types;
|
||||
|
||||
fn prepare_openai_request(
|
||||
req: &mut openai::ChatCompletionRequest,
|
||||
resolved: &ResolvedBackend,
|
||||
streaming: bool,
|
||||
warnings: &mut TranslationWarnings,
|
||||
provider_catalog: &ProviderCatalog,
|
||||
) -> Result<(), ChatCompletionError> {
|
||||
req.model = resolved.mapped_model.clone();
|
||||
req.stream = Some(streaming);
|
||||
|
||||
if streaming {
|
||||
if resolved.state.omit_stream_options {
|
||||
if req.stream_options.is_some() {
|
||||
warnings.add("stream_options");
|
||||
}
|
||||
req.stream_options = None;
|
||||
} else {
|
||||
req.stream_options = Some(openai::StreamOptions {
|
||||
include_usage: true,
|
||||
});
|
||||
}
|
||||
} else if resolved.state.omit_stream_options {
|
||||
if req.stream_options.is_some() {
|
||||
warnings.add("stream_options");
|
||||
}
|
||||
req.stream_options = None;
|
||||
}
|
||||
|
||||
prepare_openai_tool_request(
|
||||
req,
|
||||
OpenAiToolPolicyContext {
|
||||
backend_kind: resolved.state.backend_kind.clone(),
|
||||
provider_id: resolved.state.provider_id.as_deref(),
|
||||
model: &resolved.mapped_model,
|
||||
provider_catalog,
|
||||
},
|
||||
warnings,
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(|e| ChatCompletionError::InvalidRequest(e.to_string()))
|
||||
}
|
||||
pub use error::ChatCompletionError;
|
||||
pub use service::ChatCompletionRuntime;
|
||||
pub use types::{
|
||||
ChatCompletionChunkStream, ChatCompletionMetadata, ChatCompletionResult, ChatCompletionService,
|
||||
ChatCompletionStreamResult,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,443 @@
|
||||
//! `ChatCompletionRuntime`: model routing, backend dispatch, and the
|
||||
//! `ChatCompletionService` implementation.
|
||||
|
||||
use super::error::ChatCompletionError;
|
||||
use super::stream::{openai_chunk_stream, responses_chunk_stream, DeploymentLatencyGuard};
|
||||
use super::types::{
|
||||
ChatCompletionMetadata, ChatCompletionResult, ChatCompletionService, ChatCompletionStreamResult,
|
||||
};
|
||||
use crate::backend::{BackendClient, BackendError};
|
||||
use crate::config::{BackendKind, Config, ModelMapping, MultiConfig, OpenAIApiFormat};
|
||||
use crate::openai_tool_policy::{prepare_openai_tool_request, OpenAiToolPolicyContext};
|
||||
use anyllm_providers::ProviderCatalog;
|
||||
use anyllm_translate::{
|
||||
mapping, openai, translate_anthropic_to_openai_response, translate_openai_to_anthropic_request,
|
||||
TranslationWarnings,
|
||||
};
|
||||
use futures::future::{BoxFuture, FutureExt};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Chat Completions runtime built from proxy backend configuration.
|
||||
#[derive(Clone)]
|
||||
pub struct ChatCompletionRuntime {
|
||||
default_backend: String,
|
||||
backends: Arc<HashMap<String, RuntimeBackend>>,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
provider_catalog: Arc<ProviderCatalog>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RuntimeBackend {
|
||||
backend: BackendClient,
|
||||
backend_name: String,
|
||||
model_mapping: ModelMapping,
|
||||
backend_kind: BackendKind,
|
||||
api_format: OpenAIApiFormat,
|
||||
omit_stream_options: bool,
|
||||
stream_timeout_secs: u64,
|
||||
provider_id: Option<String>,
|
||||
}
|
||||
|
||||
struct ResolvedBackend {
|
||||
state: RuntimeBackend,
|
||||
mapped_model: String,
|
||||
deployment: Option<Arc<crate::config::model_router::Deployment>>,
|
||||
}
|
||||
|
||||
impl ChatCompletionRuntime {
|
||||
/// Build a runtime from a legacy single-backend config.
|
||||
pub fn from_config(config: Config) -> Self {
|
||||
let multi = MultiConfig::from_single_config(&config);
|
||||
let default_backend = multi.default_backend.clone();
|
||||
let (_, bc) = multi
|
||||
.backends
|
||||
.get_key_value(&default_backend)
|
||||
.expect("wrapped config must contain default backend");
|
||||
|
||||
let backend = if config.backend == BackendKind::Bedrock {
|
||||
BackendClient::from_backend_config(bc)
|
||||
} else {
|
||||
BackendClient::new(&config)
|
||||
};
|
||||
|
||||
let mut backends = HashMap::new();
|
||||
backends.insert(
|
||||
default_backend.clone(),
|
||||
RuntimeBackend {
|
||||
backend,
|
||||
backend_name: default_backend.clone(),
|
||||
model_mapping: bc.model_mapping.clone(),
|
||||
backend_kind: bc.kind.clone(),
|
||||
api_format: bc.api_format.clone(),
|
||||
omit_stream_options: bc.omit_stream_options,
|
||||
stream_timeout_secs: bc.stream_timeout_secs,
|
||||
provider_id: config.provider_id.clone(),
|
||||
},
|
||||
);
|
||||
|
||||
Self {
|
||||
default_backend,
|
||||
backends: Arc::new(backends),
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a runtime from multi-backend config without model-list routing.
|
||||
pub fn from_multi_config(config: MultiConfig) -> Self {
|
||||
Self::from_multi_config_with_model_router(config, None)
|
||||
}
|
||||
|
||||
/// Build a runtime from multi-backend config and an optional model router.
|
||||
pub fn from_multi_config_with_model_router(
|
||||
config: MultiConfig,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
) -> Self {
|
||||
let mut backends = HashMap::new();
|
||||
for (name, bc) in &config.backends {
|
||||
backends.insert(
|
||||
name.clone(),
|
||||
RuntimeBackend {
|
||||
backend: BackendClient::from_backend_config(bc),
|
||||
backend_name: name.clone(),
|
||||
model_mapping: bc.model_mapping.clone(),
|
||||
backend_kind: bc.kind.clone(),
|
||||
api_format: bc.api_format.clone(),
|
||||
omit_stream_options: bc.omit_stream_options,
|
||||
stream_timeout_secs: bc.stream_timeout_secs,
|
||||
provider_id: bc.provider_id.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Self {
|
||||
default_backend: config.default_backend,
|
||||
backends: Arc::new(backends),
|
||||
model_router,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve(&self, model: &str) -> Result<ResolvedBackend, ChatCompletionError> {
|
||||
if let Some(ref router_lock) = self.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(routed) = router.route(model) {
|
||||
let state = self
|
||||
.backends
|
||||
.get(routed.backend_name)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ChatCompletionError::Routing(format!(
|
||||
"model '{model}' routed to unknown backend '{}'",
|
||||
routed.backend_name
|
||||
))
|
||||
})?;
|
||||
return Ok(ResolvedBackend {
|
||||
state,
|
||||
mapped_model: routed.actual_model.to_string(),
|
||||
deployment: Some(routed.deployment.clone()),
|
||||
});
|
||||
}
|
||||
if router.has_model(model) {
|
||||
return Err(ChatCompletionError::Routing(
|
||||
"all deployments for this model are at their RPM limit".to_string(),
|
||||
));
|
||||
}
|
||||
return Err(ChatCompletionError::InvalidRequest(format!(
|
||||
"model '{model}' is not configured in model_list"
|
||||
)));
|
||||
}
|
||||
|
||||
let state = self
|
||||
.backends
|
||||
.get(&self.default_backend)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ChatCompletionError::Routing(format!(
|
||||
"default backend '{}' is not configured",
|
||||
self.default_backend
|
||||
))
|
||||
})?;
|
||||
Ok(ResolvedBackend {
|
||||
mapped_model: state.model_mapping.map_model(model),
|
||||
state,
|
||||
deployment: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionService for ChatCompletionRuntime {
|
||||
fn complete<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionResult, ChatCompletionError>> {
|
||||
async move { self.complete_inner(req).await }.boxed()
|
||||
}
|
||||
|
||||
fn complete_stream<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionStreamResult, ChatCompletionError>> {
|
||||
async move { self.complete_stream_inner(req).await }.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionRuntime {
|
||||
async fn complete_inner(
|
||||
&self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> Result<ChatCompletionResult, ChatCompletionError> {
|
||||
if req.stream == Some(true) {
|
||||
return Err(ChatCompletionError::InvalidRequest(
|
||||
"complete does not accept stream=true; use complete_stream".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let requested_model = req.model.clone();
|
||||
let resolved = self.resolve(&requested_model)?;
|
||||
let metadata = metadata(&requested_model, &resolved);
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
match &resolved.state.backend {
|
||||
BackendClient::OpenAI(client)
|
||||
| BackendClient::AzureOpenAI(client)
|
||||
| BackendClient::Vertex(client)
|
||||
| BackendClient::GeminiOpenAI(client) => {
|
||||
let mut openai_req = req;
|
||||
prepare_openai_request(
|
||||
&mut openai_req,
|
||||
&resolved,
|
||||
false,
|
||||
&mut warnings,
|
||||
&self.provider_catalog,
|
||||
)?;
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.chat_completion(&openai_req).await {
|
||||
Ok((response, _status, rate_limits)) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
if let Some(ref deployment) = resolved.deployment {
|
||||
if let Some(ref usage) = response.usage {
|
||||
deployment.record_tokens(usage.total_tokens as u64);
|
||||
}
|
||||
}
|
||||
Ok(ChatCompletionResult {
|
||||
usage: response.usage.clone(),
|
||||
response,
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::OpenAIResponses(client) => {
|
||||
let anthropic_req = translate_openai_to_anthropic_request(&req, &mut warnings)?;
|
||||
let mut responses_req =
|
||||
mapping::responses_message_map::anthropic_to_responses_request(&anthropic_req);
|
||||
responses_req.model = resolved.mapped_model.clone();
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.responses(&responses_req).await {
|
||||
Ok((resp, _status, rate_limits)) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
let anthropic_resp =
|
||||
mapping::responses_message_map::responses_to_anthropic_response(
|
||||
&resp,
|
||||
&requested_model,
|
||||
);
|
||||
if let Some(ref deployment) = resolved.deployment {
|
||||
deployment.record_tokens(
|
||||
anthropic_resp.usage.input_tokens as u64
|
||||
+ anthropic_resp.usage.output_tokens as u64,
|
||||
);
|
||||
}
|
||||
let response = translate_anthropic_to_openai_response(
|
||||
&anthropic_resp,
|
||||
&requested_model,
|
||||
);
|
||||
Ok(ChatCompletionResult {
|
||||
usage: response.usage.clone(),
|
||||
response,
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::Anthropic(_)
|
||||
| BackendClient::Bedrock(_)
|
||||
| BackendClient::GeminiNative(_) => Err(ChatCompletionError::UnsupportedBackend {
|
||||
backend_name: resolved.state.backend_name,
|
||||
backend_kind: resolved.state.backend_kind,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn complete_stream_inner(
|
||||
&self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> Result<ChatCompletionStreamResult, ChatCompletionError> {
|
||||
let requested_model = req.model.clone();
|
||||
let resolved = self.resolve(&requested_model)?;
|
||||
let metadata = metadata(&requested_model, &resolved);
|
||||
let mut warnings = TranslationWarnings::default();
|
||||
|
||||
match &resolved.state.backend {
|
||||
BackendClient::OpenAI(client)
|
||||
| BackendClient::AzureOpenAI(client)
|
||||
| BackendClient::Vertex(client)
|
||||
| BackendClient::GeminiOpenAI(client) => {
|
||||
let mut openai_req = req;
|
||||
prepare_openai_request(
|
||||
&mut openai_req,
|
||||
&resolved,
|
||||
true,
|
||||
&mut warnings,
|
||||
&self.provider_catalog,
|
||||
)?;
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.chat_completion_stream(&openai_req).await {
|
||||
Ok((response, rate_limits)) => Ok(ChatCompletionStreamResult {
|
||||
chunks: openai_chunk_stream(
|
||||
response,
|
||||
resolved.state.stream_timeout_secs,
|
||||
DeploymentLatencyGuard::from_started(
|
||||
resolved.deployment.clone(),
|
||||
start,
|
||||
),
|
||||
),
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
}),
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::OpenAIResponses(client) => {
|
||||
let anthropic_req = translate_openai_to_anthropic_request(&req, &mut warnings)?;
|
||||
let mut responses_req =
|
||||
mapping::responses_message_map::anthropic_to_responses_request(&anthropic_req);
|
||||
responses_req.model = resolved.mapped_model.clone();
|
||||
responses_req.stream = Some(true);
|
||||
|
||||
let start = record_start(&resolved.deployment);
|
||||
match client.responses_stream(&responses_req).await {
|
||||
Ok((response, rate_limits)) => Ok(ChatCompletionStreamResult {
|
||||
chunks: responses_chunk_stream(
|
||||
response,
|
||||
requested_model,
|
||||
resolved.state.stream_timeout_secs,
|
||||
DeploymentLatencyGuard::from_started(
|
||||
resolved.deployment.clone(),
|
||||
start,
|
||||
),
|
||||
),
|
||||
rate_limits,
|
||||
metadata,
|
||||
warnings,
|
||||
}),
|
||||
Err(e) => {
|
||||
record_finish(&resolved.deployment, start);
|
||||
Err(ChatCompletionError::Backend(BackendError::from(e)))
|
||||
}
|
||||
}
|
||||
}
|
||||
BackendClient::Anthropic(_)
|
||||
| BackendClient::Bedrock(_)
|
||||
| BackendClient::GeminiNative(_) => Err(ChatCompletionError::UnsupportedBackend {
|
||||
backend_name: resolved.state.backend_name,
|
||||
backend_kind: resolved.state.backend_kind,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata(requested_model: &str, resolved: &ResolvedBackend) -> ChatCompletionMetadata {
|
||||
ChatCompletionMetadata {
|
||||
requested_model: requested_model.to_string(),
|
||||
selected_backend: resolved.state.backend_name.clone(),
|
||||
mapped_model: resolved.mapped_model.clone(),
|
||||
backend_kind: resolved.state.backend_kind.clone(),
|
||||
provider_id: resolved.state.provider_id.clone(),
|
||||
api_format: resolved.state.api_format.clone(),
|
||||
used_responses_api: matches!(resolved.state.backend, BackendClient::OpenAIResponses(_)),
|
||||
}
|
||||
}
|
||||
|
||||
fn record_start(
|
||||
deployment: &Option<Arc<crate::config::model_router::Deployment>>,
|
||||
) -> Option<Instant> {
|
||||
if let Some(d) = deployment {
|
||||
d.record_start();
|
||||
Some(Instant::now())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn record_finish(
|
||||
deployment: &Option<Arc<crate::config::model_router::Deployment>>,
|
||||
start: Option<Instant>,
|
||||
) {
|
||||
if let (Some(d), Some(start)) = (deployment, start) {
|
||||
d.record_finish(start.elapsed().as_millis() as u64);
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_openai_request(
|
||||
req: &mut openai::ChatCompletionRequest,
|
||||
resolved: &ResolvedBackend,
|
||||
streaming: bool,
|
||||
warnings: &mut TranslationWarnings,
|
||||
provider_catalog: &ProviderCatalog,
|
||||
) -> Result<(), ChatCompletionError> {
|
||||
req.model = resolved.mapped_model.clone();
|
||||
req.stream = Some(streaming);
|
||||
|
||||
if streaming {
|
||||
if resolved.state.omit_stream_options {
|
||||
if req.stream_options.is_some() {
|
||||
warnings.add("stream_options");
|
||||
}
|
||||
req.stream_options = None;
|
||||
} else {
|
||||
req.stream_options = Some(openai::StreamOptions {
|
||||
include_usage: true,
|
||||
});
|
||||
}
|
||||
} else if resolved.state.omit_stream_options {
|
||||
if req.stream_options.is_some() {
|
||||
warnings.add("stream_options");
|
||||
}
|
||||
req.stream_options = None;
|
||||
}
|
||||
|
||||
prepare_openai_tool_request(
|
||||
req,
|
||||
OpenAiToolPolicyContext {
|
||||
backend_kind: resolved.state.backend_kind.clone(),
|
||||
provider_id: resolved.state.provider_id.as_deref(),
|
||||
model: &resolved.mapped_model,
|
||||
provider_catalog,
|
||||
},
|
||||
warnings,
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(|e| ChatCompletionError::InvalidRequest(e.to_string()))
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
//! Runtime result, metadata, and service-trait types.
|
||||
|
||||
use super::error::ChatCompletionError;
|
||||
use crate::backend::RateLimitHeaders;
|
||||
use crate::config::{BackendKind, OpenAIApiFormat};
|
||||
use anyllm_translate::{openai, TranslationWarnings};
|
||||
use futures::{future::BoxFuture, Stream};
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
|
||||
/// Stream returned by [`ChatCompletionRuntime::complete_stream`].
|
||||
pub type ChatCompletionChunkStream =
|
||||
Pin<Box<dyn Stream<Item = Result<openai::ChatCompletionChunk, ChatCompletionError>> + Send>>;
|
||||
|
||||
/// Object-safe service API for one Chat Completions call at a time.
|
||||
pub trait ChatCompletionService: Send + Sync {
|
||||
fn complete<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionResult, ChatCompletionError>>;
|
||||
|
||||
fn complete_stream<'a>(
|
||||
&'a self,
|
||||
req: openai::ChatCompletionRequest,
|
||||
) -> BoxFuture<'a, Result<ChatCompletionStreamResult, ChatCompletionError>>;
|
||||
}
|
||||
|
||||
/// Non-streaming runtime response.
|
||||
#[derive(Debug)]
|
||||
pub struct ChatCompletionResult {
|
||||
pub response: openai::ChatCompletionResponse,
|
||||
pub usage: Option<openai::ChatUsage>,
|
||||
pub rate_limits: RateLimitHeaders,
|
||||
pub metadata: ChatCompletionMetadata,
|
||||
pub warnings: TranslationWarnings,
|
||||
}
|
||||
|
||||
/// Streaming runtime response.
|
||||
pub struct ChatCompletionStreamResult {
|
||||
pub chunks: ChatCompletionChunkStream,
|
||||
pub rate_limits: RateLimitHeaders,
|
||||
pub metadata: ChatCompletionMetadata,
|
||||
pub warnings: TranslationWarnings,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ChatCompletionStreamResult {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("ChatCompletionStreamResult")
|
||||
.field("chunks", &"<stream>")
|
||||
.field("rate_limits", &self.rate_limits)
|
||||
.field("metadata", &self.metadata)
|
||||
.field("warnings", &self.warnings)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Backend selection metadata for observability and callers that need accounting.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatCompletionMetadata {
|
||||
pub requested_model: String,
|
||||
pub selected_backend: String,
|
||||
pub mapped_model: String,
|
||||
pub backend_kind: BackendKind,
|
||||
pub provider_id: Option<String>,
|
||||
pub api_format: OpenAIApiFormat,
|
||||
pub used_responses_api: bool,
|
||||
}
|
||||
@@ -509,78 +509,5 @@ pub async fn validate_auth(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod auth_mode_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_new_names() {
|
||||
assert_eq!(AuthMode::from_env_str("oidc"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("oidc-only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("oidc_only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys-only"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys_only"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("both"), AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_legacy_names() {
|
||||
assert_eq!(AuthMode::from_env_str("jwt_only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("jwt_or_keys"), AuthMode::Both);
|
||||
assert_eq!(AuthMode::from_env_str("JWT_ONLY"), AuthMode::OidcOnly);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_unknown_defaults_to_both() {
|
||||
assert_eq!(AuthMode::from_env_str("unknown"), AuthMode::Both);
|
||||
assert_eq!(AuthMode::from_env_str(""), AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_oidc_only() {
|
||||
assert!(AuthMode::OidcOnly.allows_oidc());
|
||||
assert!(!AuthMode::OidcOnly.allows_key_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_keys_only() {
|
||||
assert!(AuthMode::KeysOnly.allows_key_auth());
|
||||
assert!(!AuthMode::KeysOnly.allows_oidc());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_both_allows_all() {
|
||||
assert!(AuthMode::Both.allows_oidc());
|
||||
assert!(AuthMode::Both.allows_key_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_from_env_defaults_to_both() {
|
||||
let mode = AuthMode::from_env_str("unrecognized_value");
|
||||
assert_eq!(mode, AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_rejects_multiple_static_keys_without_open_relay() {
|
||||
assert!(forward_client_auth_misconfigured(2, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_open_relay_even_with_multiple_keys() {
|
||||
// Reflects that `open_relay_active()` can only be true when
|
||||
// `distinct_static_key_count()` is 0 (see the OPEN_RELAY static) --
|
||||
// this combination is unreachable via those two real accessors, but
|
||||
// the pure decision function itself must still handle it sanely.
|
||||
assert!(!forward_client_auth_misconfigured(2, true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_exactly_one_key() {
|
||||
assert!(!forward_client_auth_misconfigured(1, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_zero_keys() {
|
||||
assert!(!forward_client_auth_misconfigured(0, false));
|
||||
}
|
||||
}
|
||||
#[path = "auth/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_new_names() {
|
||||
assert_eq!(AuthMode::from_env_str("oidc"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("oidc-only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("oidc_only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys-only"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("keys_only"), AuthMode::KeysOnly);
|
||||
assert_eq!(AuthMode::from_env_str("both"), AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_legacy_names() {
|
||||
assert_eq!(AuthMode::from_env_str("jwt_only"), AuthMode::OidcOnly);
|
||||
assert_eq!(AuthMode::from_env_str("jwt_or_keys"), AuthMode::Both);
|
||||
assert_eq!(AuthMode::from_env_str("JWT_ONLY"), AuthMode::OidcOnly);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_auth_mode_unknown_defaults_to_both() {
|
||||
assert_eq!(AuthMode::from_env_str("unknown"), AuthMode::Both);
|
||||
assert_eq!(AuthMode::from_env_str(""), AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_oidc_only() {
|
||||
assert!(AuthMode::OidcOnly.allows_oidc());
|
||||
assert!(!AuthMode::OidcOnly.allows_key_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_keys_only() {
|
||||
assert!(AuthMode::KeysOnly.allows_key_auth());
|
||||
assert!(!AuthMode::KeysOnly.allows_oidc());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_both_allows_all() {
|
||||
assert!(AuthMode::Both.allows_oidc());
|
||||
assert!(AuthMode::Both.allows_key_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_from_env_defaults_to_both() {
|
||||
let mode = AuthMode::from_env_str("unrecognized_value");
|
||||
assert_eq!(mode, AuthMode::Both);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_rejects_multiple_static_keys_without_open_relay() {
|
||||
assert!(forward_client_auth_misconfigured(2, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_open_relay_even_with_multiple_keys() {
|
||||
// Reflects that `open_relay_active()` can only be true when
|
||||
// `distinct_static_key_count()` is 0 (see the OPEN_RELAY static) --
|
||||
// this combination is unreachable via those two real accessors, but
|
||||
// the pure decision function itself must still handle it sanely.
|
||||
assert!(!forward_client_auth_misconfigured(2, true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_exactly_one_key() {
|
||||
assert!(!forward_client_auth_misconfigured(1, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_client_auth_allows_zero_keys() {
|
||||
assert!(!forward_client_auth_misconfigured(0, false));
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
use crate::server::middleware::ClientAuthPath;
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
/// Selects the exact incoming credential to forward upstream when
|
||||
/// `ANTHROPIC_FORWARD_CLIENT_AUTH` is enabled. Same precedence as
|
||||
/// `validate_auth` (`x-api-key` / `x-goog-api-key` win over `authorization`,
|
||||
/// see `server/middleware/auth.rs`'s `api_key = headers.get("x-api-key")
|
||||
/// .or_else(|| headers.get("x-goog-api-key"))`) so this always forwards the
|
||||
/// credential that actually gated the request into the proxy, never a
|
||||
/// second, unrelated header the client also happened to send.
|
||||
///
|
||||
/// `x-goog-api-key` (Gemini-CLI compatibility) is folded into the
|
||||
/// `x-api-key` slot rather than forwarded under its own name: Anthropic's API
|
||||
/// only recognizes `x-api-key`/`authorization`, so a client authenticated via
|
||||
/// `x-goog-api-key` must still have its value sent upstream as `x-api-key`,
|
||||
/// not as a header name Anthropic would silently ignore. Beyond that one
|
||||
/// rename, no shape detection/conversion happens -- forwarded byte-for-byte,
|
||||
/// unlike LiteLLM's `optionally_handle_anthropic_oauth()`, which mis-converts
|
||||
/// a Bearer token into `x-api-key`.
|
||||
fn select_client_auth_override(headers: &HeaderMap) -> Option<(&'static str, &str)> {
|
||||
if let Some(v) = headers
|
||||
.get("x-api-key")
|
||||
.or_else(|| headers.get("x-goog-api-key"))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if !v.is_empty() {
|
||||
return Some(("x-api-key", v));
|
||||
}
|
||||
}
|
||||
if let Some(v) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
|
||||
if !v.is_empty() {
|
||||
return Some(("authorization", v));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Only `StaticKey`/`OpenRelay` mean "the credential that gated this request
|
||||
/// IS the operator's own secret" (a single-key/BYOK deployment). A virtual
|
||||
/// key is deliberately not a real Anthropic credential and a JWT is a
|
||||
/// proxy-auth artifact, so those must never be forwarded upstream regardless
|
||||
/// of the `ANTHROPIC_FORWARD_CLIENT_AUTH` toggle.
|
||||
fn client_auth_forwardable(auth_path: Option<ClientAuthPath>) -> bool {
|
||||
matches!(
|
||||
auth_path,
|
||||
Some(ClientAuthPath::StaticKey) | Some(ClientAuthPath::OpenRelay)
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolves the client-credential override shared by both passthrough
|
||||
/// handlers. `vk_ctx`/`claims` are checked directly, not just via
|
||||
/// `client_auth_forwardable(auth_path)`: `ClientAuthPath` and
|
||||
/// `VirtualKeyContext`/`JwtClaims` are inserted as two independent
|
||||
/// `request.extensions_mut().insert()` calls in `validate_auth`, with
|
||||
/// nothing structurally coupling them, so a future edit to one of those
|
||||
/// branches could desync them without a compile error. Re-checking presence
|
||||
/// of the extension that actually gates virtual-key/OIDC requests fails
|
||||
/// closed instead of silently forwarding a non-operator credential if that
|
||||
/// ever happens.
|
||||
pub(crate) fn resolve_client_auth_override<'h>(
|
||||
forward_client_auth: bool,
|
||||
auth_path: Option<ClientAuthPath>,
|
||||
vk_ctx: &Option<crate::server::middleware::VirtualKeyContext>,
|
||||
claims: &Option<crate::server::oidc::JwtClaims>,
|
||||
headers: &'h HeaderMap,
|
||||
) -> Option<(&'static str, &'h str)> {
|
||||
if forward_client_auth
|
||||
&& client_auth_forwardable(auth_path)
|
||||
&& vk_ctx.is_none()
|
||||
&& claims.is_none()
|
||||
{
|
||||
select_client_auth_override(headers)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
client_auth_forwardable, resolve_client_auth_override, select_client_auth_override,
|
||||
ClientAuthPath,
|
||||
};
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
fn headers(pairs: &[(&str, &str)]) -> HeaderMap {
|
||||
let mut h = HeaderMap::new();
|
||||
for (k, v) in pairs {
|
||||
h.insert(
|
||||
axum::http::HeaderName::from_bytes(k.as_bytes()).unwrap(),
|
||||
v.parse().unwrap(),
|
||||
);
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_x_api_key_when_only_that_is_sent() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_authorization_when_only_that_is_sent() {
|
||||
let h = headers(&[("authorization", "Bearer sk-ant-oat-abc")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("authorization", "Bearer sk-ant-oat-abc"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_x_api_key_when_both_sent_matching_validate_auth_precedence() {
|
||||
let h = headers(&[
|
||||
("x-api-key", "client-key"),
|
||||
("authorization", "Bearer sk-ant-oat-abc"),
|
||||
]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_neither_header_sent() {
|
||||
let h = headers(&[]);
|
||||
assert_eq!(select_client_auth_override(&h), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_x_goog_api_key_forwarded_as_x_api_key() {
|
||||
// validate_auth (server/middleware/auth.rs) treats x-goog-api-key as
|
||||
// fully equivalent to x-api-key for authentication, but Anthropic's
|
||||
// API only understands x-api-key -- the value must be forwarded
|
||||
// under the x-api-key name, not the literal x-goog-api-key name.
|
||||
let h = headers(&[("x-goog-api-key", "gemini-cli-key")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "gemini-cli-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_x_api_key_over_x_goog_api_key_matching_validate_auth_precedence() {
|
||||
let h = headers(&[
|
||||
("x-api-key", "primary-key"),
|
||||
("x-goog-api-key", "secondary-key"),
|
||||
]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "primary-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_auth_forwardable_only_for_static_key_and_open_relay() {
|
||||
assert!(client_auth_forwardable(Some(ClientAuthPath::StaticKey)));
|
||||
assert!(client_auth_forwardable(Some(ClientAuthPath::OpenRelay)));
|
||||
assert!(!client_auth_forwardable(Some(ClientAuthPath::VirtualKey)));
|
||||
assert!(!client_auth_forwardable(Some(ClientAuthPath::OidcJwt)));
|
||||
assert!(!client_auth_forwardable(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_forwards_on_static_key() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(true, Some(ClientAuthPath::StaticKey), &None, &None, &h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_refuses_when_vk_ctx_present_even_if_auth_path_says_static_key()
|
||||
{
|
||||
// Regression guard for the ClientAuthPath/VirtualKeyContext desync
|
||||
// risk: even if a future bug leaves auth_path reporting StaticKey
|
||||
// while a VirtualKeyContext extension is also present, forwarding
|
||||
// must still be refused.
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
let vk_ctx = Some(crate::server::middleware::VirtualKeyContext {
|
||||
key_id: 1,
|
||||
#[cfg(feature = "redis")]
|
||||
key_hash_hex: String::new(),
|
||||
rate_state: std::sync::Arc::new(crate::admin::keys::RateLimitState::new()),
|
||||
allowed_models: None,
|
||||
allowed_routes: None,
|
||||
period_reset: None,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(true, Some(ClientAuthPath::StaticKey), &vk_ctx, &None, &h),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_refuses_when_feature_disabled() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(false, Some(ClientAuthPath::StaticKey), &None, &None, &h),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
+12
-218
@@ -1,12 +1,10 @@
|
||||
// Anthropic passthrough handler: forwards raw request bytes to the real Anthropic API.
|
||||
// No translation: the proxy receives Anthropic format and returns Anthropic format.
|
||||
|
||||
use crate::backend::{BackendClient, MAX_SSE_BUFFER_SIZE};
|
||||
use crate::openai_tool_policy::{
|
||||
backend_kind_for_policy, validate_anthropic_tool_request, OpenAiToolPolicyContext,
|
||||
};
|
||||
use crate::server::middleware::ClientAuthPath;
|
||||
use crate::server::routes::{log_request, record_virtual_key_usage, RequestCtx};
|
||||
use crate::server::state::ConcurrencyPermit;
|
||||
use crate::server::state::{AppState, ConcurrencyPermit};
|
||||
use crate::server::streaming::{observe_anthropic_sse_frames, AnthropicStreamUsage, StreamOutcome};
|
||||
use anyllm_translate::{anthropic, mapping};
|
||||
use axum::{
|
||||
@@ -20,89 +18,14 @@ use futures::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use super::middleware::ClientAuthPath;
|
||||
use super::state::AppState;
|
||||
|
||||
/// Selects the exact incoming credential to forward upstream when
|
||||
/// `ANTHROPIC_FORWARD_CLIENT_AUTH` is enabled. Same precedence as
|
||||
/// `validate_auth` (`x-api-key` / `x-goog-api-key` win over `authorization`,
|
||||
/// see `server/middleware/auth.rs`'s `api_key = headers.get("x-api-key")
|
||||
/// .or_else(|| headers.get("x-goog-api-key"))`) so this always forwards the
|
||||
/// credential that actually gated the request into the proxy, never a
|
||||
/// second, unrelated header the client also happened to send.
|
||||
///
|
||||
/// `x-goog-api-key` (Gemini-CLI compatibility) is folded into the
|
||||
/// `x-api-key` slot rather than forwarded under its own name: Anthropic's API
|
||||
/// only recognizes `x-api-key`/`authorization`, so a client authenticated via
|
||||
/// `x-goog-api-key` must still have its value sent upstream as `x-api-key`,
|
||||
/// not as a header name Anthropic would silently ignore. Beyond that one
|
||||
/// rename, no shape detection/conversion happens -- forwarded byte-for-byte,
|
||||
/// unlike LiteLLM's `optionally_handle_anthropic_oauth()`, which mis-converts
|
||||
/// a Bearer token into `x-api-key`.
|
||||
fn select_client_auth_override(headers: &axum::http::HeaderMap) -> Option<(&'static str, &str)> {
|
||||
if let Some(v) = headers
|
||||
.get("x-api-key")
|
||||
.or_else(|| headers.get("x-goog-api-key"))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if !v.is_empty() {
|
||||
return Some(("x-api-key", v));
|
||||
}
|
||||
}
|
||||
if let Some(v) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
|
||||
if !v.is_empty() {
|
||||
return Some(("authorization", v));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Only `StaticKey`/`OpenRelay` mean "the credential that gated this request
|
||||
/// IS the operator's own secret" (a single-key/BYOK deployment). A virtual
|
||||
/// key is deliberately not a real Anthropic credential and a JWT is a
|
||||
/// proxy-auth artifact, so those must never be forwarded upstream regardless
|
||||
/// of the `ANTHROPIC_FORWARD_CLIENT_AUTH` toggle.
|
||||
fn client_auth_forwardable(auth_path: Option<ClientAuthPath>) -> bool {
|
||||
matches!(
|
||||
auth_path,
|
||||
Some(ClientAuthPath::StaticKey) | Some(ClientAuthPath::OpenRelay)
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolves the client-credential override shared by both passthrough
|
||||
/// handlers. `vk_ctx`/`claims` are checked directly, not just via
|
||||
/// `client_auth_forwardable(auth_path)`: `ClientAuthPath` and
|
||||
/// `VirtualKeyContext`/`JwtClaims` are inserted as two independent
|
||||
/// `request.extensions_mut().insert()` calls in `validate_auth`, with
|
||||
/// nothing structurally coupling them, so a future edit to one of those
|
||||
/// branches could desync them without a compile error. Re-checking presence
|
||||
/// of the extension that actually gates virtual-key/OIDC requests fails
|
||||
/// closed instead of silently forwarding a non-operator credential if that
|
||||
/// ever happens.
|
||||
fn resolve_client_auth_override<'h>(
|
||||
forward_client_auth: bool,
|
||||
auth_path: Option<ClientAuthPath>,
|
||||
vk_ctx: &Option<super::middleware::VirtualKeyContext>,
|
||||
claims: &Option<crate::server::oidc::JwtClaims>,
|
||||
headers: &'h axum::http::HeaderMap,
|
||||
) -> Option<(&'static str, &'h str)> {
|
||||
if forward_client_auth
|
||||
&& client_auth_forwardable(auth_path)
|
||||
&& vk_ctx.is_none()
|
||||
&& claims.is_none()
|
||||
{
|
||||
select_client_auth_override(headers)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
use super::auth::resolve_client_auth_override;
|
||||
|
||||
/// Forward an Anthropic-format request byte-for-byte to the upstream Anthropic API.
|
||||
/// No translation is performed. Only active when `BACKEND=anthropic`.
|
||||
pub(crate) async fn anthropic_passthrough(
|
||||
State(state): State<AppState>,
|
||||
permit: Option<axum::Extension<ConcurrencyPermit>>,
|
||||
vk_ctx: Option<axum::Extension<super::middleware::VirtualKeyContext>>,
|
||||
vk_ctx: Option<axum::Extension<crate::server::middleware::VirtualKeyContext>>,
|
||||
auth_path: Option<axum::Extension<ClientAuthPath>>,
|
||||
claims: Option<axum::Extension<crate::server::oidc::JwtClaims>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
@@ -196,7 +119,7 @@ pub(crate) async fn anthropic_passthrough(
|
||||
if let Some(ref ctx) = vk_ctx {
|
||||
match &peek.model {
|
||||
Some(m) => {
|
||||
if !super::policy::is_model_allowed(m, &ctx.allowed_models) {
|
||||
if !crate::server::policy::is_model_allowed(m, &ctx.allowed_models) {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::PermissionError,
|
||||
format!("Model '{}' is not allowed for this API key.", m),
|
||||
@@ -303,7 +226,7 @@ pub(crate) async fn anthropic_passthrough(
|
||||
}
|
||||
}
|
||||
|
||||
let mut body = match super::secret_redaction::redact_body_with_content_type(
|
||||
let mut body = match crate::server::secret_redaction::redact_body_with_content_type(
|
||||
state.redact_secrets(),
|
||||
Some("application/json"),
|
||||
body,
|
||||
@@ -311,7 +234,7 @@ pub(crate) async fn anthropic_passthrough(
|
||||
.await
|
||||
{
|
||||
Ok(body) => body,
|
||||
Err(err) => return super::secret_redaction::error_response(err),
|
||||
Err(err) => return crate::server::secret_redaction::error_response(err),
|
||||
};
|
||||
|
||||
// FFEC prompt compression runs FIRST: it compresses conversation history text
|
||||
@@ -576,7 +499,7 @@ pub(crate) async fn anthropic_passthrough(
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn anthropic_generic_passthrough(
|
||||
State(state): State<AppState>,
|
||||
vk_ctx: Option<axum::Extension<super::middleware::VirtualKeyContext>>,
|
||||
vk_ctx: Option<axum::Extension<crate::server::middleware::VirtualKeyContext>>,
|
||||
auth_path: Option<axum::Extension<ClientAuthPath>>,
|
||||
claims: Option<axum::Extension<crate::server::oidc::JwtClaims>>,
|
||||
OriginalUri(uri): OriginalUri,
|
||||
@@ -650,9 +573,11 @@ pub(crate) async fn anthropic_generic_passthrough(
|
||||
}
|
||||
|
||||
let body =
|
||||
match super::secret_redaction::redact_body(state.redact_secrets(), &headers, body).await {
|
||||
match crate::server::secret_redaction::redact_body(state.redact_secrets(), &headers, body)
|
||||
.await
|
||||
{
|
||||
Ok(body) => body,
|
||||
Err(err) => return super::secret_redaction::error_response(err),
|
||||
Err(err) => return crate::server::secret_redaction::error_response(err),
|
||||
};
|
||||
|
||||
match client
|
||||
@@ -730,134 +655,3 @@ fn virtual_key_accounting_parse_error() -> Response {
|
||||
);
|
||||
(StatusCode::BAD_GATEWAY, Json(err)).into_response()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
client_auth_forwardable, resolve_client_auth_override, select_client_auth_override,
|
||||
ClientAuthPath,
|
||||
};
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
fn headers(pairs: &[(&str, &str)]) -> HeaderMap {
|
||||
let mut h = HeaderMap::new();
|
||||
for (k, v) in pairs {
|
||||
h.insert(
|
||||
axum::http::HeaderName::from_bytes(k.as_bytes()).unwrap(),
|
||||
v.parse().unwrap(),
|
||||
);
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_x_api_key_when_only_that_is_sent() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_authorization_when_only_that_is_sent() {
|
||||
let h = headers(&[("authorization", "Bearer sk-ant-oat-abc")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("authorization", "Bearer sk-ant-oat-abc"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_x_api_key_when_both_sent_matching_validate_auth_precedence() {
|
||||
let h = headers(&[
|
||||
("x-api-key", "client-key"),
|
||||
("authorization", "Bearer sk-ant-oat-abc"),
|
||||
]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_neither_header_sent() {
|
||||
let h = headers(&[]);
|
||||
assert_eq!(select_client_auth_override(&h), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_x_goog_api_key_forwarded_as_x_api_key() {
|
||||
// validate_auth (server/middleware/auth.rs) treats x-goog-api-key as
|
||||
// fully equivalent to x-api-key for authentication, but Anthropic's
|
||||
// API only understands x-api-key -- the value must be forwarded
|
||||
// under the x-api-key name, not the literal x-goog-api-key name.
|
||||
let h = headers(&[("x-goog-api-key", "gemini-cli-key")]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "gemini-cli-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_x_api_key_over_x_goog_api_key_matching_validate_auth_precedence() {
|
||||
let h = headers(&[
|
||||
("x-api-key", "primary-key"),
|
||||
("x-goog-api-key", "secondary-key"),
|
||||
]);
|
||||
assert_eq!(
|
||||
select_client_auth_override(&h),
|
||||
Some(("x-api-key", "primary-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_auth_forwardable_only_for_static_key_and_open_relay() {
|
||||
assert!(client_auth_forwardable(Some(ClientAuthPath::StaticKey)));
|
||||
assert!(client_auth_forwardable(Some(ClientAuthPath::OpenRelay)));
|
||||
assert!(!client_auth_forwardable(Some(ClientAuthPath::VirtualKey)));
|
||||
assert!(!client_auth_forwardable(Some(ClientAuthPath::OidcJwt)));
|
||||
assert!(!client_auth_forwardable(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_forwards_on_static_key() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(true, Some(ClientAuthPath::StaticKey), &None, &None, &h),
|
||||
Some(("x-api-key", "client-key"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_refuses_when_vk_ctx_present_even_if_auth_path_says_static_key()
|
||||
{
|
||||
// Regression guard for the ClientAuthPath/VirtualKeyContext desync
|
||||
// risk: even if a future bug leaves auth_path reporting StaticKey
|
||||
// while a VirtualKeyContext extension is also present, forwarding
|
||||
// must still be refused.
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
let vk_ctx = Some(crate::server::middleware::VirtualKeyContext {
|
||||
key_id: 1,
|
||||
#[cfg(feature = "redis")]
|
||||
key_hash_hex: String::new(),
|
||||
rate_state: std::sync::Arc::new(crate::admin::keys::RateLimitState::new()),
|
||||
allowed_models: None,
|
||||
allowed_routes: None,
|
||||
period_reset: None,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(true, Some(ClientAuthPath::StaticKey), &vk_ctx, &None, &h),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_client_auth_override_refuses_when_feature_disabled() {
|
||||
let h = headers(&[("x-api-key", "client-key")]);
|
||||
assert_eq!(
|
||||
resolve_client_auth_override(false, Some(ClientAuthPath::StaticKey), &None, &None, &h),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
// Anthropic passthrough handlers module: forwards raw request bytes to the real Anthropic API.
|
||||
// No translation: the proxy receives Anthropic format and returns Anthropic format.
|
||||
|
||||
pub(crate) mod auth;
|
||||
pub(crate) mod handlers;
|
||||
|
||||
pub(crate) use handlers::{anthropic_generic_passthrough, anthropic_passthrough};
|
||||
@@ -1,977 +0,0 @@
|
||||
// Shared state types for request handlers: AppState, AnthropicJson, ResolvedModel, etc.
|
||||
// Extracted from routes.rs so consumers can import state independently of the router setup.
|
||||
|
||||
use crate::admin::state::{RuntimeConfig, SharedState};
|
||||
use crate::backend::BackendClient;
|
||||
use crate::metrics::Metrics;
|
||||
use anyllm_providers::ProviderCatalog;
|
||||
use anyllm_translate::{anthropic, mapping, openai};
|
||||
use axum::{
|
||||
extract::{rejection::JsonRejection, FromRequest},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Json, Response},
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
/// Custom JSON extractor that returns Anthropic-shaped error responses on
|
||||
/// parse failure. Axum's built-in Json returns its own error format, which
|
||||
/// would break clients expecting Anthropic error shapes.
|
||||
pub(crate) struct AnthropicJson<T>(pub T);
|
||||
|
||||
impl<S, T> FromRequest<S> for AnthropicJson<T>
|
||||
where
|
||||
Json<T>: FromRequest<S, Rejection = JsonRejection>,
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Rejection = Response;
|
||||
|
||||
async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
|
||||
match Json::<T>::from_request(req, state).await {
|
||||
Ok(Json(value)) => Ok(AnthropicJson(value)),
|
||||
Err(rejection) => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::InvalidRequestError,
|
||||
rejection.body_text(),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::BAD_REQUEST, Json(err)).into_response())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of resolving a model name through the model router.
|
||||
pub(crate) enum ResolvedModel {
|
||||
/// Routed via model_list to a specific backend and actual model name.
|
||||
Routed {
|
||||
backend_name: String,
|
||||
model: String,
|
||||
/// The deployment Arc for recording in-flight/latency stats.
|
||||
deployment: Arc<crate::config::model_router::Deployment>,
|
||||
/// Per-route option overrides when routed via a DB route; `None` when
|
||||
/// routed via the LiteLLM model_router (inherit global config).
|
||||
options: Option<Arc<crate::config::route_router::RouteOptions>>,
|
||||
},
|
||||
/// Model is known but all deployments are at their RPM limit.
|
||||
AllAtLimit,
|
||||
/// Model router is active but the model alias is not configured.
|
||||
UnknownModel,
|
||||
/// No model router, or model not in router. Used legacy ModelMapping.
|
||||
Legacy(String),
|
||||
}
|
||||
|
||||
/// Shared state for tool execution, stored in AppState.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolEngineState {
|
||||
pub registry: Arc<crate::tools::ToolRegistry>,
|
||||
pub policy: Arc<crate::tools::ToolExecutionPolicy>,
|
||||
pub loop_config: crate::tools::LoopConfig,
|
||||
pub guardrails: crate::tools::ToolGuardrailConfig,
|
||||
pub mcp_manager: Option<Arc<crate::tools::McpServerManager>>,
|
||||
}
|
||||
|
||||
/// Per-backend state shared across request handlers.
|
||||
///
|
||||
/// In single-backend mode, one `AppState` serves all routes. In multi-backend mode,
|
||||
/// each backend gets its own `AppState` mounted under a prefix path (e.g., `/openai/v1/messages`).
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub backend: BackendClient,
|
||||
pub metrics: Metrics,
|
||||
/// Runtime config (model mappings, body logging, redaction) read on every request.
|
||||
/// Shared with admin server so config changes take effect immediately.
|
||||
pub runtime_config: Arc<RwLock<RuntimeConfig>>,
|
||||
/// Shared admin state for request logging and live updates. None in tests.
|
||||
pub shared: Option<SharedState>,
|
||||
/// Per-route option overrides for the request that produced this (cloned)
|
||||
/// state, set by `resolve_model_and_state` when a DB route was selected.
|
||||
/// `None` means "no route override; use the global RuntimeConfig value".
|
||||
/// Read by the option accessors (`redact_secrets`, `effective_tool_guardrails`,
|
||||
/// `active_pxpipe`, `pxpipe_models`).
|
||||
pub route_options: Option<Arc<crate::config::route_router::RouteOptions>>,
|
||||
/// Backend name for logging purposes.
|
||||
pub backend_name: String,
|
||||
/// Canonical provider id used for provider/model policy decisions.
|
||||
pub provider_id: Option<String>,
|
||||
/// Concurrency limiter. Uses try_acquire (fail-fast) instead of queueing
|
||||
/// to prevent cascading latency under load. Requests exceeding the limit
|
||||
/// get 429 immediately, matching Anthropic's rate limiting behavior.
|
||||
pub concurrency: Arc<Semaphore>,
|
||||
/// Strip `stream_options` from streaming requests for local LLM compat.
|
||||
pub omit_stream_options: bool,
|
||||
/// Wall-clock cap for streaming responses in seconds. 0 = disabled.
|
||||
/// Prevents resource exhaustion from stalled backends.
|
||||
pub stream_timeout_secs: u64,
|
||||
/// When true, set `x-anyllm-degradation` header on responses that silently drop features.
|
||||
/// Mirrors Config::expose_degradation_warnings / MultiConfig::expose_degradation_warnings.
|
||||
pub expose_degradation_warnings: bool,
|
||||
/// Optional response cache for non-streaming requests.
|
||||
pub cache: Option<Arc<crate::cache::memory::MemoryCache>>,
|
||||
/// Anthropic thinking-block record-and-restore repair store. `None`
|
||||
/// unless `backend` is `BackendClient::Anthropic`; only consulted by
|
||||
/// `anthropic_passthrough`. Always `Some` for Anthropic backends
|
||||
/// regardless of whether the feature is enabled -- use
|
||||
/// `thinking_repair_enabled()` to check the live toggle before using it.
|
||||
pub thinking_repair: Option<Arc<crate::thinking_repair::ThinkingRepairStore>>,
|
||||
/// Text-to-image context compression engine (pxpipe). `None` unless
|
||||
/// `backend` is `BackendClient::Anthropic`; only consulted by
|
||||
/// `anthropic_passthrough`. Always `Some` for Anthropic backends regardless
|
||||
/// of the live toggle -- use `active_pxpipe()`, which checks
|
||||
/// `RuntimeConfig.pxpipe_compress`, before using it.
|
||||
pub pxpipe: Option<Arc<crate::pxpipe::PxpipeEngine>>,
|
||||
/// Command-aware tool-output compression engine (RTK). `Some` for Anthropic
|
||||
/// and Translate modes; only consulted when `RuntimeConfig.rtk_compress` is
|
||||
/// on -- use `rtk_engine_for(model)`.
|
||||
pub rtk: Option<Arc<crate::rtk::RtkEngine>>,
|
||||
/// FFEC prompt-compression engine (`OptimizerEngine`). `Some` for Anthropic
|
||||
/// and Translate modes, mirroring `rtk`; baked with the static
|
||||
/// `OPTIMIZER_MODE`-env default at startup -- use `effective_optimizer()`,
|
||||
/// which applies the live `RouteOptions.optimizer_mode` override on top.
|
||||
pub optimizer: Option<Arc<crate::optimizer::OptimizerEngine>>,
|
||||
/// Model-level router for LiteLLM model_list configs. None for TOML/env configs.
|
||||
/// Wrapped in RwLock for dynamic model management via admin API.
|
||||
pub model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
/// Immutable provider/model catalog used for runtime model metadata.
|
||||
pub provider_catalog: Arc<ProviderCatalog>,
|
||||
/// All backend states, for cross-backend model routing. None unless model_router is set.
|
||||
pub all_backends: Option<Arc<HashMap<String, AppState>>>,
|
||||
/// Tool execution engine state. None when tool execution is not configured.
|
||||
pub tool_engine: Option<Arc<ToolEngineState>>,
|
||||
/// Batch orchestration engine. None in test configs that don't need batch.
|
||||
pub batch_engine: Option<
|
||||
Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
/// Map a model name through the current runtime config for this backend.
|
||||
pub(crate) fn map_model(&self, model: &str) -> String {
|
||||
let config = self
|
||||
.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(mapping) = config.model_mappings.get(&self.backend_name) {
|
||||
mapping.map_model(model)
|
||||
} else {
|
||||
model.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a model name to a backend.
|
||||
///
|
||||
/// Precedence: (1) admin-DB routes (`RouteRouter`), (2) LiteLLM model_router,
|
||||
/// (3) legacy ModelMapping. An empty route router falls straight through so
|
||||
/// installs without routes behave exactly as before.
|
||||
pub(crate) fn resolve_model(&self, model: &str) -> ResolvedModel {
|
||||
if let Some(shared) = self.shared.as_ref() {
|
||||
if let Some(ref rr_lock) = shared.route_router {
|
||||
use crate::config::route_router::RouteResolution;
|
||||
let rr = rr_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if !rr.is_empty() {
|
||||
match rr.resolve(model) {
|
||||
RouteResolution::Routed(res) => {
|
||||
return ResolvedModel::Routed {
|
||||
backend_name: res.backend_name,
|
||||
model: res.model,
|
||||
deployment: res.deployment,
|
||||
options: Some(res.options),
|
||||
};
|
||||
}
|
||||
RouteResolution::AllAtLimit => return ResolvedModel::AllAtLimit,
|
||||
// No route serves this model: fall through to the layers below.
|
||||
RouteResolution::NoRoute => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(ref router_lock) = self.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(routed) = router.route(model) {
|
||||
return ResolvedModel::Routed {
|
||||
backend_name: routed.backend_name.to_string(),
|
||||
model: routed.actual_model.to_string(),
|
||||
deployment: routed.deployment.clone(),
|
||||
options: None,
|
||||
};
|
||||
}
|
||||
if router.has_model(model) {
|
||||
return ResolvedModel::AllAtLimit;
|
||||
}
|
||||
return ResolvedModel::UnknownModel;
|
||||
}
|
||||
ResolvedModel::Legacy(self.map_model(model))
|
||||
}
|
||||
|
||||
/// Resolve model and return (mapped_model, effective AppState, optional deployment).
|
||||
/// If the model routes to a different backend, the returned state is cloned from
|
||||
/// all_backends. Returns Err with a 429 response if all deployments are at limit.
|
||||
/// The deployment Arc is returned so handlers can call record_start/record_finish.
|
||||
#[allow(clippy::result_large_err)]
|
||||
pub(crate) fn resolve_model_and_state(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Result<
|
||||
(
|
||||
String,
|
||||
AppState,
|
||||
Option<Arc<crate::config::model_router::Deployment>>,
|
||||
),
|
||||
Response,
|
||||
> {
|
||||
match self.resolve_model(model) {
|
||||
ResolvedModel::Routed {
|
||||
backend_name,
|
||||
model: mapped,
|
||||
deployment,
|
||||
options,
|
||||
} => {
|
||||
let mut effective = self
|
||||
.all_backends
|
||||
.as_ref()
|
||||
.and_then(|m| m.get(&backend_name))
|
||||
.cloned()
|
||||
.or_else(|| {
|
||||
// Check managed backends (SQLite-backed, zero-restart)
|
||||
self.shared.as_ref().and_then(|s| {
|
||||
let guard = s.managed_backends
|
||||
.read()
|
||||
.ok()
|
||||
.or_else(|| {
|
||||
tracing::warn!("managed_backends RwLock is poisoned; skipping managed backend lookup");
|
||||
None
|
||||
})?;
|
||||
guard.get(&backend_name).map(|(row, client)| {
|
||||
let mut state = self.clone();
|
||||
state.backend = client.clone();
|
||||
state.backend_name = backend_name.clone();
|
||||
state.provider_id = Some(row.provider_id.clone());
|
||||
state
|
||||
})
|
||||
})
|
||||
})
|
||||
.unwrap_or_else(|| self.clone());
|
||||
// Carry the per-route option overrides onto the effective state so
|
||||
// the option accessors resolve route-first, global-fallback.
|
||||
effective.route_options = options;
|
||||
Ok((mapped, effective, Some(deployment)))
|
||||
}
|
||||
ResolvedModel::AllAtLimit => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::RateLimitError,
|
||||
"all deployments for this model are at their RPM limit".to_string(),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response())
|
||||
}
|
||||
ResolvedModel::UnknownModel => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::InvalidRequestError,
|
||||
format!("model '{model}' is not configured in model_list"),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::BAD_REQUEST, Json(err)).into_response())
|
||||
}
|
||||
ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone(), None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether request/response body logging is enabled.
|
||||
pub(crate) fn log_bodies(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.log_bodies
|
||||
}
|
||||
|
||||
/// Whether upstream JSON/text request payloads should be redacted.
|
||||
/// Route override (if set) wins over the global RuntimeConfig value.
|
||||
pub(crate) fn redact_secrets(&self) -> bool {
|
||||
if let Some(v) = self.route_options.as_ref().and_then(|o| o.redact_secrets) {
|
||||
return v;
|
||||
}
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.redact_secrets
|
||||
}
|
||||
|
||||
/// Effective tool-call guardrail config for this request: the runtime,
|
||||
/// admin-tunable override (`RuntimeConfig.tool_guardrail_mode`, no
|
||||
/// restart required) applied on top of `engine.guardrails` (the static
|
||||
/// preset built from YAML/env at startup). See
|
||||
/// `crate::tools::resolve_runtime_guardrails`.
|
||||
pub(crate) fn effective_tool_guardrails(
|
||||
&self,
|
||||
engine: &ToolEngineState,
|
||||
) -> crate::tools::ToolGuardrailConfig {
|
||||
// Route override (if set) wins over the live global RuntimeConfig mode.
|
||||
if let Some(mode) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.guardrail_mode.as_deref())
|
||||
{
|
||||
return crate::tools::resolve_runtime_guardrails(&engine.guardrails, mode);
|
||||
}
|
||||
crate::tools::resolve_runtime_guardrails_locked(&self.runtime_config, &engine.guardrails)
|
||||
}
|
||||
|
||||
/// Whether Anthropic thinking-block repair (record + restore) is active.
|
||||
/// `self.thinking_repair` may be `Some` even when this is `false` -- the
|
||||
/// store is always constructed for Anthropic backends; only this flag
|
||||
/// gates whether it's actually used.
|
||||
pub(crate) fn thinking_repair_enabled(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.anthropic_thinking_repair
|
||||
}
|
||||
|
||||
/// Whether Anthropic passthrough forwards the client's own incoming
|
||||
/// credential upstream instead of the operator's (`ANTHROPIC_FORWARD_CLIENT_AUTH`,
|
||||
/// live-toggleable via `RuntimeConfig.forward_client_auth`). Read fresh on
|
||||
/// every request -- unlike the old frozen `AppState` field this replaced,
|
||||
/// this reflects an admin-UI change immediately without a restart, and
|
||||
/// applies uniformly to every `BackendKind::Anthropic` backend since they
|
||||
/// all share one `RuntimeConfig`.
|
||||
pub(crate) fn forward_client_auth_enabled(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.forward_client_auth
|
||||
}
|
||||
|
||||
/// The thinking-repair store, but only when the live admin-toggleable
|
||||
/// flag is actually on. `None` both when repair is entirely absent (non-
|
||||
/// Anthropic backend) and when it's present-but-disabled -- single
|
||||
/// accessor so call sites collapse to `if let Some(store) = ...` instead
|
||||
/// of separately checking `thinking_repair_enabled()` and
|
||||
/// `thinking_repair.is_some()`.
|
||||
pub(crate) fn active_thinking_repair(
|
||||
&self,
|
||||
) -> Option<Arc<crate::thinking_repair::ThinkingRepairStore>> {
|
||||
if self.thinking_repair_enabled() {
|
||||
self.thinking_repair.clone()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// The pxpipe compression engine, but only when the live admin-toggleable
|
||||
/// flag (`RuntimeConfig.pxpipe_compress`) is on. `None` both when the engine
|
||||
/// is absent (non-Anthropic backend) and when present-but-disabled.
|
||||
pub(crate) fn active_pxpipe(&self) -> Option<Arc<crate::pxpipe::PxpipeEngine>> {
|
||||
let enabled = match self.route_options.as_ref().and_then(|o| o.pxpipe_compress) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pxpipe_compress
|
||||
}
|
||||
};
|
||||
if enabled {
|
||||
self.pxpipe.clone()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Live model-scope CSV for pxpipe. Route override wins over the global
|
||||
/// `RuntimeConfig.pxpipe_models` value.
|
||||
pub(crate) fn pxpipe_models(&self) -> String {
|
||||
if let Some(csv) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.pxpipe_models.clone())
|
||||
{
|
||||
return csv;
|
||||
}
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pxpipe_models
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Vision gate: if the catalog knows this model and says it is NOT
|
||||
/// vision-capable, refuse (fail-closed). Unknown models fall back to the
|
||||
/// scope list only — a Claude passthrough model is vision-capable in
|
||||
/// practice, and the scope list is the operator's explicit control.
|
||||
fn pxpipe_vision_ok(&self, model: &str) -> bool {
|
||||
match self
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.and_then(|pid| self.provider_catalog.get_model(pid, model))
|
||||
{
|
||||
Some(def) => def.capabilities.vision,
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// The pxpipe engine for `model`, or `None` if compression shouldn't run:
|
||||
/// the master toggle is off, the engine is absent (non-Anthropic backend),
|
||||
/// the model is out of the live scope CSV, or it isn't vision-capable.
|
||||
/// Single accessor so `passthrough` collapses to
|
||||
/// `if let Some(engine) = state.pxpipe_engine_for(model)`.
|
||||
pub(crate) fn pxpipe_engine_for(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Option<Arc<crate::pxpipe::PxpipeEngine>> {
|
||||
let engine = self.active_pxpipe()?;
|
||||
if crate::pxpipe::model_in_scope(model, &self.pxpipe_models())
|
||||
&& self.pxpipe_vision_ok(model)
|
||||
{
|
||||
Some(engine)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// The RTK engine for `model`, or `None` if compression shouldn't run: the
|
||||
/// toggle is off, the engine is absent, or the model is out of scope. RTK is
|
||||
/// not vision-gated, so there is no capability check.
|
||||
///
|
||||
/// Reads the toggle and scope from a single RwLock critical section for
|
||||
/// consistency, and checks `route_options` first (matching the pxpipe pattern)
|
||||
/// so per-route overrides take precedence over the global RuntimeConfig.
|
||||
pub(crate) fn rtk_engine_for(&self, model: &str) -> Option<Arc<crate::rtk::RtkEngine>> {
|
||||
let cfg = self
|
||||
.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
let enabled = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.rtk_compress)
|
||||
.unwrap_or(cfg.rtk_compress);
|
||||
if !enabled {
|
||||
return None;
|
||||
}
|
||||
let engine = self.rtk.clone()?;
|
||||
let models_csv = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.rtk_models.as_deref())
|
||||
.unwrap_or(&cfg.rtk_models);
|
||||
if crate::rtk::model_in_scope(model, models_csv) {
|
||||
Some(engine)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Effective FFEC prompt-compression engine for this request, or `None`
|
||||
/// when optimization is unconfigured for this backend/mode (`self.optimizer`
|
||||
/// is `None`). Precedence, mirroring `effective_tool_guardrails` /
|
||||
/// `resolve_runtime_guardrails_locked`: (1) route override
|
||||
/// (`RouteOptions.optimizer_mode`, if set) wins outright; (2) otherwise the
|
||||
/// live `RuntimeConfig.optimizer_mode` admin toggle (no restart required);
|
||||
/// (3) otherwise the static per-process engine baked with the
|
||||
/// `OPTIMIZER_MODE`-env default at startup.
|
||||
pub(crate) fn effective_optimizer(&self) -> Option<Arc<crate::optimizer::OptimizerEngine>> {
|
||||
let engine = self.optimizer.as_ref()?;
|
||||
if let Some(mode_str) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.optimizer_mode.as_deref())
|
||||
{
|
||||
return Some(Arc::new(engine.with_mode_override(mode_str)));
|
||||
}
|
||||
Some(Arc::new(
|
||||
crate::optimizer::resolve_runtime_optimizer_locked(&self.runtime_config, engine),
|
||||
))
|
||||
}
|
||||
|
||||
/// Apply RTK tool-output compression to an OpenAI-format request and record
|
||||
/// metrics. Shared helper used by both the /v1/chat/completions and /v1/messages
|
||||
/// translate paths (streaming and non-streaming). No-op when the engine is
|
||||
/// unavailable, disabled, or no tool messages are present.
|
||||
pub(crate) fn apply_rtk_to_openai(&self, req: &mut openai::ChatCompletionRequest, model: &str) {
|
||||
let engine = match self.rtk_engine_for(model) {
|
||||
Some(e) => e,
|
||||
None => return,
|
||||
};
|
||||
// Pre-check: only serialize when there are tool messages to compress.
|
||||
if !req
|
||||
.messages
|
||||
.iter()
|
||||
.any(|m| m.role == openai::ChatRole::Tool)
|
||||
{
|
||||
return;
|
||||
}
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let Some((blocks, saved)) = engine.compress_openai_chat(&mut v) else {
|
||||
return;
|
||||
};
|
||||
match serde_json::from_value::<openai::ChatCompletionRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_rtk_compression(blocks, saved);
|
||||
tracing::info!(
|
||||
model,
|
||||
blocks,
|
||||
chars_saved = saved,
|
||||
"rtk: compressed OpenAI request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"rtk: failed to re-deserialize compressed OpenAI request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply FFEC prompt compression (`effective_optimizer()`) to an OpenAI-format
|
||||
/// request at the parsed-body seam. Client-sent history only -- callers must
|
||||
/// never invoke this on proxy-appended tool-loop turns (see
|
||||
/// `crates/optimizer/CLAUDE.md` "Streaming & tool-loop decision"). `Shadow`
|
||||
/// mode logs the `OptimizationReport` and leaves `req` unchanged; `Live` mode
|
||||
/// applies the rendered body in place. No-op when optimization is
|
||||
/// unconfigured or resolves to `Mode::Off` for this request.
|
||||
pub(crate) fn apply_optimizer_to_openai(
|
||||
&self,
|
||||
req: &mut openai::ChatCompletionRequest,
|
||||
route: &str,
|
||||
) {
|
||||
let Some(engine) = self.effective_optimizer() else {
|
||||
return;
|
||||
};
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let report = engine.optimize_openai(&mut v, route);
|
||||
if report.mode == anyllm_optimize_core::Mode::Shadow {
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
failure = report.failure.as_deref().unwrap_or(""),
|
||||
"optimizer: shadow report (not applied)"
|
||||
);
|
||||
}
|
||||
if !report.applied {
|
||||
return;
|
||||
}
|
||||
match serde_json::from_value::<openai::ChatCompletionRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_optimization(
|
||||
report.messages_compressed as u64,
|
||||
report.removed_tokens_est,
|
||||
);
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
"optimizer: compressed OpenAI request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"optimizer: failed to re-deserialize compressed OpenAI request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply FFEC prompt compression (`effective_optimizer()`) to an Anthropic
|
||||
/// Messages request at the parsed-body seam. Same contract as
|
||||
/// [`Self::apply_optimizer_to_openai`]: client-sent history only, fails open,
|
||||
/// `Shadow` never mutates `req`.
|
||||
pub(crate) fn apply_optimizer_to_anthropic(
|
||||
&self,
|
||||
req: &mut anthropic::MessageCreateRequest,
|
||||
route: &str,
|
||||
) {
|
||||
let Some(engine) = self.effective_optimizer() else {
|
||||
return;
|
||||
};
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let report = engine.optimize_anthropic(&mut v, route);
|
||||
if report.mode == anyllm_optimize_core::Mode::Shadow {
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
failure = report.failure.as_deref().unwrap_or(""),
|
||||
"optimizer: shadow report (not applied)"
|
||||
);
|
||||
}
|
||||
if !report.applied {
|
||||
return;
|
||||
}
|
||||
match serde_json::from_value::<anthropic::MessageCreateRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_optimization(
|
||||
report.messages_compressed as u64,
|
||||
report.removed_tokens_est,
|
||||
);
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
"optimizer: compressed Anthropic request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"optimizer: failed to re-deserialize compressed Anthropic request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod optimizer_seam_tests {
|
||||
use super::*;
|
||||
use crate::config::{
|
||||
BackendAuth, BackendKind, Config, ModelMapping, OpenAIApiFormat, TlsConfig,
|
||||
};
|
||||
use anyllm_optimize_core::Mode;
|
||||
|
||||
/// Long enough that FFEC's min-length gate actually has something to compress.
|
||||
fn long_text() -> String {
|
||||
"The quick brown fox jumps over the lazy dog again and again across the wide \
|
||||
green field toward the distant blue mountains far beyond the winding river."
|
||||
.repeat(4)
|
||||
}
|
||||
|
||||
fn minimal_state(optimizer_mode: Mode) -> AppState {
|
||||
let config = Config {
|
||||
backend: BackendKind::OpenAI,
|
||||
openai_api_key: "test".into(),
|
||||
openai_base_url: "https://api.openai.com".into(),
|
||||
listen_port: 3000,
|
||||
model_mapping: ModelMapping {
|
||||
big_model: "gpt-4o".into(),
|
||||
small_model: "gpt-4o-mini".into(),
|
||||
},
|
||||
tls: TlsConfig::default(),
|
||||
backend_auth: BackendAuth::BearerToken("test".into()),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
expose_degradation_warnings: false,
|
||||
openai_api_format: OpenAIApiFormat::Chat,
|
||||
provider_id: None,
|
||||
};
|
||||
let backend = crate::backend::BackendClient::OpenAI(
|
||||
crate::backend::openai_client::OpenAIClient::new(&config),
|
||||
);
|
||||
let runtime_config = Arc::new(RwLock::new(RuntimeConfig {
|
||||
model_mappings: indexmap::IndexMap::new(),
|
||||
log_level: "info".to_string(),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
pxpipe_models: String::new(),
|
||||
rtk_compress: false,
|
||||
rtk_models: String::new(),
|
||||
forward_client_auth: false,
|
||||
tool_guardrail_mode: "disabled".to_string(),
|
||||
optimizer_mode: optimizer_mode.as_str().to_string(),
|
||||
}));
|
||||
AppState {
|
||||
backend,
|
||||
metrics: Metrics::new(),
|
||||
runtime_config,
|
||||
shared: None,
|
||||
route_options: None,
|
||||
backend_name: "openai".to_string(),
|
||||
provider_id: None,
|
||||
concurrency: Arc::new(Semaphore::new(64)),
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: 0,
|
||||
expose_degradation_warnings: false,
|
||||
cache: None,
|
||||
thinking_repair: None,
|
||||
pxpipe: None,
|
||||
rtk: None,
|
||||
optimizer: Some(Arc::new(crate::optimizer::OptimizerEngine::new(
|
||||
optimizer_mode,
|
||||
))),
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
all_backends: None,
|
||||
tool_engine: None,
|
||||
batch_engine: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn long_openai_request() -> openai::ChatCompletionRequest {
|
||||
let long = long_text();
|
||||
let mut body = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "system", "content": "you are helpful"}],
|
||||
});
|
||||
let msgs = body["messages"].as_array_mut().unwrap();
|
||||
for _ in 0..16 {
|
||||
msgs.push(serde_json::json!({"role": "user", "content": long}));
|
||||
msgs.push(serde_json::json!({"role": "assistant", "content": long}));
|
||||
}
|
||||
msgs.push(serde_json::json!({"role": "user", "content": "what is the latest?"}));
|
||||
serde_json::from_value(body).expect("valid ChatCompletionRequest")
|
||||
}
|
||||
|
||||
fn long_anthropic_request() -> anthropic::MessageCreateRequest {
|
||||
let long = long_text();
|
||||
let mut body = serde_json::json!({
|
||||
"model": "claude-sonnet-5",
|
||||
"max_tokens": 1024,
|
||||
"messages": [],
|
||||
});
|
||||
let msgs = body["messages"].as_array_mut().unwrap();
|
||||
for _ in 0..16 {
|
||||
msgs.push(serde_json::json!({"role": "user", "content": long}));
|
||||
msgs.push(serde_json::json!({"role": "assistant", "content": long}));
|
||||
}
|
||||
msgs.push(serde_json::json!({"role": "user", "content": "what is the latest?"}));
|
||||
serde_json::from_value(body).expect("valid MessageCreateRequest")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_forwards_openai_body_unchanged() {
|
||||
let state = minimal_state(Mode::Shadow);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(before, after, "shadow mode must forward the original body");
|
||||
assert_eq!(
|
||||
state.metrics.snapshot().optimizer_compressed_total,
|
||||
0,
|
||||
"shadow mode must never record a metrics-visible compression"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_forwards_anthropic_body_unchanged() {
|
||||
let state = minimal_state(Mode::Shadow);
|
||||
let mut req = long_anthropic_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_anthropic(&mut req, "messages");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(before, after, "shadow mode must forward the original body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_openai_history_and_preserves_latest() {
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(
|
||||
before["messages"].as_array().unwrap().last(),
|
||||
after["messages"].as_array().unwrap().last(),
|
||||
"the latest turn must never be rewritten"
|
||||
);
|
||||
assert_eq!(
|
||||
state.metrics.snapshot().optimizer_compressed_total,
|
||||
1,
|
||||
"an applied Live compression must be recorded in metrics"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_anthropic_history_and_preserves_latest() {
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req = long_anthropic_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_anthropic(&mut req, "messages");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(
|
||||
before["messages"].as_array().unwrap().last(),
|
||||
after["messages"].as_array().unwrap().last(),
|
||||
"the latest turn must never be rewritten"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn off_mode_is_noop_and_engine_absent_is_noop() {
|
||||
// Off mode: engine present, mode Off -> never applied.
|
||||
let state = minimal_state(Mode::Off);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
|
||||
// No engine at all (e.g. non-Anthropic/Translate mode backend): no panic, no-op.
|
||||
let mut state_no_engine = minimal_state(Mode::Live);
|
||||
state_no_engine.optimizer = None;
|
||||
let mut req2 = long_openai_request();
|
||||
let before2 = serde_json::to_value(&req2).unwrap();
|
||||
state_no_engine.apply_optimizer_to_openai(&mut req2, "chat_completions");
|
||||
assert_eq!(before2, serde_json::to_value(&req2).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_history_below_min_len_gate_is_a_noop_not_a_panic() {
|
||||
// A short request has nothing worth compressing (below FFEC's min-length
|
||||
// gate) -- the seam must still round-trip cleanly without panicking or
|
||||
// corrupting the body, i.e. it fails open when there's nothing to do.
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req: openai::ChatCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}))
|
||||
.unwrap();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod rtk_seam_tests {
|
||||
use super::*;
|
||||
|
||||
/// Build a minimal `AppState` with the RTK engine present and the runtime
|
||||
/// `rtk_compress` toggle set to `enabled` (scope left empty = all models).
|
||||
fn state_with_rtk(enabled: bool) -> AppState {
|
||||
use crate::config::{
|
||||
BackendAuth, BackendKind, Config, ModelMapping, OpenAIApiFormat, TlsConfig,
|
||||
};
|
||||
let config = Config {
|
||||
backend: BackendKind::OpenAI,
|
||||
openai_api_key: "test".into(),
|
||||
openai_base_url: "https://api.openai.com".into(),
|
||||
listen_port: 3000,
|
||||
model_mapping: ModelMapping {
|
||||
big_model: "gpt-4o".into(),
|
||||
small_model: "gpt-4o-mini".into(),
|
||||
},
|
||||
tls: TlsConfig::default(),
|
||||
backend_auth: BackendAuth::BearerToken("test".into()),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
expose_degradation_warnings: false,
|
||||
openai_api_format: OpenAIApiFormat::Chat,
|
||||
provider_id: None,
|
||||
};
|
||||
let backend = crate::backend::BackendClient::OpenAI(
|
||||
crate::backend::openai_client::OpenAIClient::new(&config),
|
||||
);
|
||||
let runtime_config = Arc::new(RwLock::new(RuntimeConfig {
|
||||
model_mappings: indexmap::IndexMap::new(),
|
||||
log_level: "info".to_string(),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
pxpipe_models: String::new(),
|
||||
rtk_compress: enabled,
|
||||
rtk_models: String::new(),
|
||||
forward_client_auth: false,
|
||||
tool_guardrail_mode: "disabled".to_string(),
|
||||
optimizer_mode: "off".to_string(),
|
||||
}));
|
||||
AppState {
|
||||
backend,
|
||||
metrics: Metrics::new(),
|
||||
runtime_config,
|
||||
shared: None,
|
||||
route_options: None,
|
||||
backend_name: "openai".to_string(),
|
||||
provider_id: None,
|
||||
concurrency: Arc::new(Semaphore::new(64)),
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: 0,
|
||||
expose_degradation_warnings: false,
|
||||
cache: None,
|
||||
thinking_repair: None,
|
||||
pxpipe: None,
|
||||
rtk: Some(Arc::new(crate::rtk::RtkEngine::new())),
|
||||
optimizer: None,
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
all_backends: None,
|
||||
tool_engine: None,
|
||||
batch_engine: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// An OpenAI request whose `role: tool` message carries compressible git noise.
|
||||
fn request_with_noisy_tool_output() -> openai::ChatCompletionRequest {
|
||||
let mut noise = String::from("On branch main\nChanges not staged for commit:\n");
|
||||
for i in 0..200 {
|
||||
noise.push_str(&format!(" (use \"git add ...\" file {i})\n"));
|
||||
}
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [
|
||||
{"role": "user", "content": "run git status"},
|
||||
{"role": "assistant", "content": null, "tool_calls": [
|
||||
{"id": "t1", "type": "function",
|
||||
"function": {"name": "bash", "arguments": "{\"cmd\":\"git status\"}"}}
|
||||
]},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": noise},
|
||||
],
|
||||
}))
|
||||
.expect("valid ChatCompletionRequest")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enabled_compresses_tool_output_and_records_metrics() {
|
||||
let state = state_with_rtk(true);
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_ne!(before, after, "tool output should have been compressed");
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_is_a_noop() {
|
||||
let state = state_with_rtk(false);
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_absent_is_a_noop() {
|
||||
let mut state = state_with_rtk(true);
|
||||
state.rtk = None;
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Global state for the multi-backend metrics endpoint.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct GlobalState {
|
||||
pub(crate) backend_metrics: Arc<HashMap<String, Metrics>>,
|
||||
}
|
||||
|
||||
/// Wrapper so OwnedSemaphorePermit can be stored in request extensions.
|
||||
/// The field is never read directly; it exists as an RAII guard to hold
|
||||
/// the permit until the struct is dropped.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ConcurrencyPermit(
|
||||
#[allow(dead_code)] pub(crate) Arc<tokio::sync::OwnedSemaphorePermit>,
|
||||
);
|
||||
@@ -0,0 +1,33 @@
|
||||
use anyllm_translate::{anthropic, mapping};
|
||||
use axum::{
|
||||
extract::{rejection::JsonRejection, FromRequest},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Json, Response},
|
||||
};
|
||||
|
||||
/// Custom JSON extractor that returns Anthropic-shaped error responses on
|
||||
/// parse failure. Axum's built-in Json returns its own error format, which
|
||||
/// would break clients expecting Anthropic error shapes.
|
||||
pub(crate) struct AnthropicJson<T>(pub T);
|
||||
|
||||
impl<S, T> FromRequest<S> for AnthropicJson<T>
|
||||
where
|
||||
Json<T>: FromRequest<S, Rejection = JsonRejection>,
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Rejection = Response;
|
||||
|
||||
async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
|
||||
match Json::<T>::from_request(req, state).await {
|
||||
Ok(Json(value)) => Ok(AnthropicJson(value)),
|
||||
Err(rejection) => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::InvalidRequestError,
|
||||
rejection.body_text(),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::BAD_REQUEST, Json(err)).into_response())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
use crate::admin::state::{RuntimeConfig, SharedState};
|
||||
use crate::backend::BackendClient;
|
||||
use crate::metrics::Metrics;
|
||||
use anyllm_providers::ProviderCatalog;
|
||||
use anyllm_translate::{anthropic, mapping};
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Json, Response},
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use super::resolved_model::ResolvedModel;
|
||||
use super::tool_engine::ToolEngineState;
|
||||
|
||||
/// Per-backend state shared across request handlers.
|
||||
///
|
||||
/// In single-backend mode, one `AppState` serves all routes. In multi-backend mode,
|
||||
/// each backend gets its own `AppState` mounted under a prefix path (e.g., `/openai/v1/messages`).
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub backend: BackendClient,
|
||||
pub metrics: Metrics,
|
||||
/// Runtime config (model mappings, body logging, redaction) read on every request.
|
||||
/// Shared with admin server so config changes take effect immediately.
|
||||
pub runtime_config: Arc<RwLock<RuntimeConfig>>,
|
||||
/// Shared admin state for request logging and live updates. None in tests.
|
||||
pub shared: Option<SharedState>,
|
||||
/// Per-route option overrides for the request that produced this (cloned)
|
||||
/// state, set by `resolve_model_and_state` when a DB route was selected.
|
||||
/// `None` means "no route override; use the global RuntimeConfig value".
|
||||
/// Read by the option accessors (`redact_secrets`, `effective_tool_guardrails`,
|
||||
/// `active_pxpipe`, `pxpipe_models`).
|
||||
pub route_options: Option<Arc<crate::config::route_router::RouteOptions>>,
|
||||
/// Backend name for logging purposes.
|
||||
pub backend_name: String,
|
||||
/// Canonical provider id used for provider/model policy decisions.
|
||||
pub provider_id: Option<String>,
|
||||
/// Concurrency limiter. Uses try_acquire (fail-fast) instead of queueing
|
||||
/// to prevent cascading latency under load. Requests exceeding the limit
|
||||
/// get 429 immediately, matching Anthropic's rate limiting behavior.
|
||||
pub concurrency: Arc<Semaphore>,
|
||||
/// Strip `stream_options` from streaming requests for local LLM compat.
|
||||
pub omit_stream_options: bool,
|
||||
/// Wall-clock cap for streaming responses in seconds. 0 = disabled.
|
||||
/// Prevents resource exhaustion from stalled backends.
|
||||
pub stream_timeout_secs: u64,
|
||||
/// When true, set `x-anyllm-degradation` header on responses that silently drop features.
|
||||
/// Mirrors Config::expose_degradation_warnings / MultiConfig::expose_degradation_warnings.
|
||||
pub expose_degradation_warnings: bool,
|
||||
/// Optional response cache for non-streaming requests.
|
||||
pub cache: Option<Arc<crate::cache::memory::MemoryCache>>,
|
||||
/// Anthropic thinking-block record-and-restore repair store. `None`
|
||||
/// unless `backend` is `BackendClient::Anthropic`; only consulted by
|
||||
/// `anthropic_passthrough`. Always `Some` for Anthropic backends
|
||||
/// regardless of whether the feature is enabled -- use
|
||||
/// `thinking_repair_enabled()` to check the live toggle before using it.
|
||||
pub thinking_repair: Option<Arc<crate::thinking_repair::ThinkingRepairStore>>,
|
||||
/// Text-to-image context compression engine (pxpipe). `None` unless
|
||||
/// `backend` is `BackendClient::Anthropic`; only consulted by
|
||||
/// `anthropic_passthrough`. Always `Some` for Anthropic backends regardless
|
||||
/// of the live toggle -- use `active_pxpipe()`, which checks
|
||||
/// `RuntimeConfig.pxpipe_compress`, before using it.
|
||||
pub pxpipe: Option<Arc<crate::pxpipe::PxpipeEngine>>,
|
||||
/// Command-aware tool-output compression engine (RTK). `Some` for Anthropic
|
||||
/// and Translate modes; only consulted when `RuntimeConfig.rtk_compress` is
|
||||
/// on -- use `rtk_engine_for(model)`.
|
||||
pub rtk: Option<Arc<crate::rtk::RtkEngine>>,
|
||||
/// FFEC prompt-compression engine (`OptimizerEngine`). `Some` for Anthropic
|
||||
/// and Translate modes, mirroring `rtk`; baked with the static
|
||||
/// `OPTIMIZER_MODE`-env default at startup -- use `effective_optimizer()`,
|
||||
/// which applies the live `RouteOptions.optimizer_mode` override on top.
|
||||
pub optimizer: Option<Arc<crate::optimizer::OptimizerEngine>>,
|
||||
/// Model-level router for LiteLLM model_list configs. None for TOML/env configs.
|
||||
/// Wrapped in RwLock for dynamic model management via admin API.
|
||||
pub model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
/// Immutable provider/model catalog used for runtime model metadata.
|
||||
pub provider_catalog: Arc<ProviderCatalog>,
|
||||
/// All backend states, for cross-backend model routing. None unless model_router is set.
|
||||
pub all_backends: Option<Arc<HashMap<String, AppState>>>,
|
||||
/// Tool execution engine state. None when tool execution is not configured.
|
||||
pub tool_engine: Option<Arc<ToolEngineState>>,
|
||||
/// Batch orchestration engine. None in test configs that don't need batch.
|
||||
pub batch_engine: Option<
|
||||
Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
/// Map a model name through the current runtime config for this backend.
|
||||
pub(crate) fn map_model(&self, model: &str) -> String {
|
||||
let config = self
|
||||
.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(mapping) = config.model_mappings.get(&self.backend_name) {
|
||||
mapping.map_model(model)
|
||||
} else {
|
||||
model.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a model name to a backend.
|
||||
///
|
||||
/// Precedence: (1) admin-DB routes (`RouteRouter`), (2) LiteLLM model_router,
|
||||
/// (3) legacy ModelMapping. An empty route router falls straight through so
|
||||
/// installs without routes behave exactly as before.
|
||||
pub(crate) fn resolve_model(&self, model: &str) -> ResolvedModel {
|
||||
if let Some(shared) = self.shared.as_ref() {
|
||||
if let Some(ref rr_lock) = shared.route_router {
|
||||
use crate::config::route_router::RouteResolution;
|
||||
let rr = rr_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if !rr.is_empty() {
|
||||
match rr.resolve(model) {
|
||||
RouteResolution::Routed(res) => {
|
||||
return ResolvedModel::Routed {
|
||||
backend_name: res.backend_name,
|
||||
model: res.model,
|
||||
deployment: res.deployment,
|
||||
options: Some(res.options),
|
||||
};
|
||||
}
|
||||
RouteResolution::AllAtLimit => return ResolvedModel::AllAtLimit,
|
||||
// No route serves this model: fall through to the layers below.
|
||||
RouteResolution::NoRoute => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(ref router_lock) = self.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(routed) = router.route(model) {
|
||||
return ResolvedModel::Routed {
|
||||
backend_name: routed.backend_name.to_string(),
|
||||
model: routed.actual_model.to_string(),
|
||||
deployment: routed.deployment.clone(),
|
||||
options: None,
|
||||
};
|
||||
}
|
||||
if router.has_model(model) {
|
||||
return ResolvedModel::AllAtLimit;
|
||||
}
|
||||
return ResolvedModel::UnknownModel;
|
||||
}
|
||||
ResolvedModel::Legacy(self.map_model(model))
|
||||
}
|
||||
|
||||
/// Resolve model and return (mapped_model, effective AppState, optional deployment).
|
||||
/// If the model routes to a different backend, the returned state is cloned from
|
||||
/// all_backends. Returns Err with a 429 response if all deployments are at limit.
|
||||
/// The deployment Arc is returned so handlers can call record_start/record_finish.
|
||||
#[allow(clippy::result_large_err)]
|
||||
pub(crate) fn resolve_model_and_state(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Result<
|
||||
(
|
||||
String,
|
||||
AppState,
|
||||
Option<Arc<crate::config::model_router::Deployment>>,
|
||||
),
|
||||
Response,
|
||||
> {
|
||||
match self.resolve_model(model) {
|
||||
ResolvedModel::Routed {
|
||||
backend_name,
|
||||
model: mapped,
|
||||
deployment,
|
||||
options,
|
||||
} => {
|
||||
let mut effective = self
|
||||
.all_backends
|
||||
.as_ref()
|
||||
.and_then(|m| m.get(&backend_name))
|
||||
.cloned()
|
||||
.or_else(|| {
|
||||
// Check managed backends (SQLite-backed, zero-restart)
|
||||
self.shared.as_ref().and_then(|s| {
|
||||
let guard = s.managed_backends
|
||||
.read()
|
||||
.ok()
|
||||
.or_else(|| {
|
||||
tracing::warn!("managed_backends RwLock is poisoned; skipping managed backend lookup");
|
||||
None
|
||||
})?;
|
||||
guard.get(&backend_name).map(|(row, client)| {
|
||||
let mut state = self.clone();
|
||||
state.backend = client.clone();
|
||||
state.backend_name = backend_name.clone();
|
||||
state.provider_id = Some(row.provider_id.clone());
|
||||
state
|
||||
})
|
||||
})
|
||||
})
|
||||
.unwrap_or_else(|| self.clone());
|
||||
// Carry the per-route option overrides onto the effective state so
|
||||
// the option accessors resolve route-first, global-fallback.
|
||||
effective.route_options = options;
|
||||
Ok((mapped, effective, Some(deployment)))
|
||||
}
|
||||
ResolvedModel::AllAtLimit => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::RateLimitError,
|
||||
"all deployments for this model are at their RPM limit".to_string(),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response())
|
||||
}
|
||||
ResolvedModel::UnknownModel => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
anthropic::ErrorType::InvalidRequestError,
|
||||
format!("model '{model}' is not configured in model_list"),
|
||||
None,
|
||||
);
|
||||
Err((StatusCode::BAD_REQUEST, Json(err)).into_response())
|
||||
}
|
||||
ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone(), None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether request/response body logging is enabled.
|
||||
pub(crate) fn log_bodies(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.log_bodies
|
||||
}
|
||||
|
||||
/// Whether upstream JSON/text request payloads should be redacted.
|
||||
/// Route override (if set) wins over the global RuntimeConfig value.
|
||||
pub(crate) fn redact_secrets(&self) -> bool {
|
||||
if let Some(v) = self.route_options.as_ref().and_then(|o| o.redact_secrets) {
|
||||
return v;
|
||||
}
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.redact_secrets
|
||||
}
|
||||
|
||||
/// Effective tool-call guardrail config for this request: the runtime,
|
||||
/// admin-tunable override (`RuntimeConfig.tool_guardrail_mode`, no
|
||||
/// restart required) applied on top of `engine.guardrails` (the static
|
||||
/// preset built from YAML/env at startup). See
|
||||
/// `crate::tools::resolve_runtime_guardrails`.
|
||||
pub(crate) fn effective_tool_guardrails(
|
||||
&self,
|
||||
engine: &ToolEngineState,
|
||||
) -> crate::tools::ToolGuardrailConfig {
|
||||
// Route override (if set) wins over the live global RuntimeConfig mode.
|
||||
if let Some(mode) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.guardrail_mode.as_deref())
|
||||
{
|
||||
return crate::tools::resolve_runtime_guardrails(&engine.guardrails, mode);
|
||||
}
|
||||
crate::tools::resolve_runtime_guardrails_locked(&self.runtime_config, &engine.guardrails)
|
||||
}
|
||||
|
||||
/// Whether Anthropic thinking-block repair (record + restore) is active.
|
||||
/// `self.thinking_repair` may be `Some` even when this is `false` -- the
|
||||
/// store is always constructed for Anthropic backends; only this flag
|
||||
/// gates whether it's actually used.
|
||||
pub(crate) fn thinking_repair_enabled(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.anthropic_thinking_repair
|
||||
}
|
||||
|
||||
/// Whether Anthropic passthrough forwards the client's own incoming
|
||||
/// credential upstream instead of the operator's (`ANTHROPIC_FORWARD_CLIENT_AUTH`,
|
||||
/// live-toggleable via `RuntimeConfig.forward_client_auth`). Read fresh on
|
||||
/// every request -- unlike the old frozen `AppState` field this replaced,
|
||||
/// this reflects an admin-UI change immediately without a restart, and
|
||||
/// applies uniformly to every `BackendKind::Anthropic` backend since they
|
||||
/// all share one `RuntimeConfig`.
|
||||
pub(crate) fn forward_client_auth_enabled(&self) -> bool {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.forward_client_auth
|
||||
}
|
||||
|
||||
/// The thinking-repair store, but only when the live admin-toggleable
|
||||
/// flag is actually on. `None` both when repair is entirely absent (non-
|
||||
/// Anthropic backend) and when it's present-but-disabled -- single
|
||||
/// accessor so call sites collapse to `if let Some(store) = ...` instead
|
||||
/// of separately checking `thinking_repair_enabled()` and
|
||||
/// `thinking_repair.is_some()`.
|
||||
pub(crate) fn active_thinking_repair(
|
||||
&self,
|
||||
) -> Option<Arc<crate::thinking_repair::ThinkingRepairStore>> {
|
||||
if self.thinking_repair_enabled() {
|
||||
self.thinking_repair.clone()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
use anyllm_translate::{anthropic, openai};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::app_state::AppState;
|
||||
|
||||
impl AppState {
|
||||
/// The pxpipe compression engine, but only when the live admin-toggleable
|
||||
/// flag (`RuntimeConfig.pxpipe_compress`) is on. `None` both when the engine
|
||||
/// is absent (non-Anthropic backend) and when present-but-disabled.
|
||||
pub(crate) fn active_pxpipe(&self) -> Option<Arc<crate::pxpipe::PxpipeEngine>> {
|
||||
let enabled = match self.route_options.as_ref().and_then(|o| o.pxpipe_compress) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pxpipe_compress
|
||||
}
|
||||
};
|
||||
if enabled {
|
||||
self.pxpipe.clone()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Live model-scope CSV for pxpipe. Route override wins over the global
|
||||
/// `RuntimeConfig.pxpipe_models` value.
|
||||
pub(crate) fn pxpipe_models(&self) -> String {
|
||||
if let Some(csv) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.pxpipe_models.clone())
|
||||
{
|
||||
return csv;
|
||||
}
|
||||
self.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pxpipe_models
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Vision gate: if the catalog knows this model and says it is NOT
|
||||
/// vision-capable, refuse (fail-closed). Unknown models fall back to the
|
||||
/// scope list only — a Claude passthrough model is vision-capable in
|
||||
/// practice, and the scope list is the operator's explicit control.
|
||||
fn pxpipe_vision_ok(&self, model: &str) -> bool {
|
||||
match self
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.and_then(|pid| self.provider_catalog.get_model(pid, model))
|
||||
{
|
||||
Some(def) => def.capabilities.vision,
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// The pxpipe engine for `model`, or `None` if compression shouldn't run:
|
||||
/// the master toggle is off, the engine is absent (non-Anthropic backend),
|
||||
/// the model is out of the live scope CSV, or it isn't vision-capable.
|
||||
/// Single accessor so `passthrough` collapses to
|
||||
/// `if let Some(engine) = state.pxpipe_engine_for(model)`.
|
||||
pub(crate) fn pxpipe_engine_for(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Option<Arc<crate::pxpipe::PxpipeEngine>> {
|
||||
let engine = self.active_pxpipe()?;
|
||||
if crate::pxpipe::model_in_scope(model, &self.pxpipe_models())
|
||||
&& self.pxpipe_vision_ok(model)
|
||||
{
|
||||
Some(engine)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// The RTK engine for `model`, or `None` if compression shouldn't run: the
|
||||
/// toggle is off, the engine is absent, or the model is out of scope. RTK is
|
||||
/// not vision-gated, so there is no capability check.
|
||||
///
|
||||
/// Reads the toggle and scope from a single RwLock critical section for
|
||||
/// consistency, and checks `route_options` first (matching the pxpipe pattern)
|
||||
/// so per-route overrides take precedence over the global RuntimeConfig.
|
||||
pub(crate) fn rtk_engine_for(&self, model: &str) -> Option<Arc<crate::rtk::RtkEngine>> {
|
||||
let cfg = self
|
||||
.runtime_config
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
let enabled = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.rtk_compress)
|
||||
.unwrap_or(cfg.rtk_compress);
|
||||
if !enabled {
|
||||
return None;
|
||||
}
|
||||
let engine = self.rtk.clone()?;
|
||||
let models_csv = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.rtk_models.as_deref())
|
||||
.unwrap_or(&cfg.rtk_models);
|
||||
if crate::rtk::model_in_scope(model, models_csv) {
|
||||
Some(engine)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Effective FFEC prompt-compression engine for this request, or `None`
|
||||
/// when optimization is unconfigured for this backend/mode (`self.optimizer`
|
||||
/// is `None`). Precedence, mirroring `effective_tool_guardrails` /
|
||||
/// `resolve_runtime_guardrails_locked`: (1) route override
|
||||
/// (`RouteOptions.optimizer_mode`, if set) wins outright; (2) otherwise the
|
||||
/// live `RuntimeConfig.optimizer_mode` admin toggle (no restart required);
|
||||
/// (3) otherwise the static per-process engine baked with the
|
||||
/// `OPTIMIZER_MODE`-env default at startup.
|
||||
pub(crate) fn effective_optimizer(&self) -> Option<Arc<crate::optimizer::OptimizerEngine>> {
|
||||
let engine = self.optimizer.as_ref()?;
|
||||
if let Some(mode_str) = self
|
||||
.route_options
|
||||
.as_ref()
|
||||
.and_then(|o| o.optimizer_mode.as_deref())
|
||||
{
|
||||
return Some(Arc::new(engine.with_mode_override(mode_str)));
|
||||
}
|
||||
Some(Arc::new(
|
||||
crate::optimizer::resolve_runtime_optimizer_locked(&self.runtime_config, engine),
|
||||
))
|
||||
}
|
||||
|
||||
/// Apply RTK tool-output compression to an OpenAI-format request and record
|
||||
/// metrics. Shared helper used by both the /v1/chat/completions and /v1/messages
|
||||
/// translate paths (streaming and non-streaming). No-op when the engine is
|
||||
/// unavailable, disabled, or no tool messages are present.
|
||||
pub(crate) fn apply_rtk_to_openai(&self, req: &mut openai::ChatCompletionRequest, model: &str) {
|
||||
let engine = match self.rtk_engine_for(model) {
|
||||
Some(e) => e,
|
||||
None => return,
|
||||
};
|
||||
// Pre-check: only serialize when there are tool messages to compress.
|
||||
if !req
|
||||
.messages
|
||||
.iter()
|
||||
.any(|m| m.role == openai::ChatRole::Tool)
|
||||
{
|
||||
return;
|
||||
}
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let Some((blocks, saved)) = engine.compress_openai_chat(&mut v) else {
|
||||
return;
|
||||
};
|
||||
match serde_json::from_value::<openai::ChatCompletionRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_rtk_compression(blocks, saved);
|
||||
tracing::info!(
|
||||
model,
|
||||
blocks,
|
||||
chars_saved = saved,
|
||||
"rtk: compressed OpenAI request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"rtk: failed to re-deserialize compressed OpenAI request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply FFEC prompt compression (`effective_optimizer()`) to an OpenAI-format
|
||||
/// request at the parsed-body seam. Client-sent history only -- callers must
|
||||
/// never invoke this on proxy-appended tool-loop turns (see
|
||||
/// `crates/optimizer/CLAUDE.md` "Streaming & tool-loop decision"). `Shadow`
|
||||
/// mode logs the `OptimizationReport` and leaves `req` unchanged; `Live` mode
|
||||
/// applies the rendered body in place. No-op when optimization is
|
||||
/// unconfigured or resolves to `Mode::Off` for this request.
|
||||
pub(crate) fn apply_optimizer_to_openai(
|
||||
&self,
|
||||
req: &mut openai::ChatCompletionRequest,
|
||||
route: &str,
|
||||
) {
|
||||
let Some(engine) = self.effective_optimizer() else {
|
||||
return;
|
||||
};
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let report = engine.optimize_openai(&mut v, route);
|
||||
if report.mode == anyllm_optimize_core::Mode::Shadow {
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
failure = report.failure.as_deref().unwrap_or(""),
|
||||
"optimizer: shadow report (not applied)"
|
||||
);
|
||||
}
|
||||
if !report.applied {
|
||||
return;
|
||||
}
|
||||
match serde_json::from_value::<openai::ChatCompletionRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_optimization(
|
||||
report.messages_compressed as u64,
|
||||
report.removed_tokens_est,
|
||||
);
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
"optimizer: compressed OpenAI request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"optimizer: failed to re-deserialize compressed OpenAI request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply FFEC prompt compression (`effective_optimizer()`) to an Anthropic
|
||||
/// Messages request at the parsed-body seam. Same contract as
|
||||
/// [`Self::apply_optimizer_to_openai`]: client-sent history only, fails open,
|
||||
/// `Shadow` never mutates `req`.
|
||||
pub(crate) fn apply_optimizer_to_anthropic(
|
||||
&self,
|
||||
req: &mut anthropic::MessageCreateRequest,
|
||||
route: &str,
|
||||
) {
|
||||
let Some(engine) = self.effective_optimizer() else {
|
||||
return;
|
||||
};
|
||||
let mut v = match serde_json::to_value(&*req) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
let report = engine.optimize_anthropic(&mut v, route);
|
||||
if report.mode == anyllm_optimize_core::Mode::Shadow {
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
failure = report.failure.as_deref().unwrap_or(""),
|
||||
"optimizer: shadow report (not applied)"
|
||||
);
|
||||
}
|
||||
if !report.applied {
|
||||
return;
|
||||
}
|
||||
match serde_json::from_value::<anthropic::MessageCreateRequest>(v) {
|
||||
Ok(patched) => {
|
||||
*req = patched;
|
||||
self.metrics.record_optimization(
|
||||
report.messages_compressed as u64,
|
||||
report.removed_tokens_est,
|
||||
);
|
||||
tracing::info!(
|
||||
route,
|
||||
removed_tokens_est = report.removed_tokens_est,
|
||||
messages_compressed = report.messages_compressed,
|
||||
"optimizer: compressed Anthropic request"
|
||||
);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"optimizer: failed to re-deserialize compressed Anthropic request; forwarding original"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "compression/tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,315 @@
|
||||
use super::*;
|
||||
use crate::admin::state::RuntimeConfig;
|
||||
use crate::config::{BackendAuth, BackendKind, Config, ModelMapping, OpenAIApiFormat, TlsConfig};
|
||||
use crate::metrics::Metrics;
|
||||
use anyllm_optimize_core::Mode;
|
||||
use anyllm_providers::ProviderCatalog;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
/// Long enough that FFEC's min-length gate actually has something to compress.
|
||||
fn long_text() -> String {
|
||||
"The quick brown fox jumps over the lazy dog again and again across the wide \
|
||||
green field toward the distant blue mountains far beyond the winding river."
|
||||
.repeat(4)
|
||||
}
|
||||
|
||||
fn minimal_state(optimizer_mode: Mode) -> AppState {
|
||||
let config = Config {
|
||||
backend: BackendKind::OpenAI,
|
||||
openai_api_key: "test".into(),
|
||||
openai_base_url: "https://api.openai.com".into(),
|
||||
listen_port: 3000,
|
||||
model_mapping: ModelMapping {
|
||||
big_model: "gpt-4o".into(),
|
||||
small_model: "gpt-4o-mini".into(),
|
||||
},
|
||||
tls: TlsConfig::default(),
|
||||
backend_auth: BackendAuth::BearerToken("test".into()),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
expose_degradation_warnings: false,
|
||||
openai_api_format: OpenAIApiFormat::Chat,
|
||||
provider_id: None,
|
||||
};
|
||||
let backend = crate::backend::BackendClient::OpenAI(
|
||||
crate::backend::openai_client::OpenAIClient::new(&config),
|
||||
);
|
||||
let runtime_config = Arc::new(RwLock::new(RuntimeConfig {
|
||||
model_mappings: indexmap::IndexMap::new(),
|
||||
log_level: "info".to_string(),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
pxpipe_models: String::new(),
|
||||
rtk_compress: false,
|
||||
rtk_models: String::new(),
|
||||
forward_client_auth: false,
|
||||
tool_guardrail_mode: "disabled".to_string(),
|
||||
optimizer_mode: optimizer_mode.as_str().to_string(),
|
||||
}));
|
||||
AppState {
|
||||
backend,
|
||||
metrics: Metrics::new(),
|
||||
runtime_config,
|
||||
shared: None,
|
||||
route_options: None,
|
||||
backend_name: "openai".to_string(),
|
||||
provider_id: None,
|
||||
concurrency: Arc::new(Semaphore::new(64)),
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: 0,
|
||||
expose_degradation_warnings: false,
|
||||
cache: None,
|
||||
thinking_repair: None,
|
||||
pxpipe: None,
|
||||
rtk: None,
|
||||
optimizer: Some(Arc::new(crate::optimizer::OptimizerEngine::new(
|
||||
optimizer_mode,
|
||||
))),
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
all_backends: None,
|
||||
tool_engine: None,
|
||||
batch_engine: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn long_openai_request() -> openai::ChatCompletionRequest {
|
||||
let long = long_text();
|
||||
let mut body = serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "system", "content": "you are helpful"}],
|
||||
});
|
||||
let msgs = body["messages"].as_array_mut().unwrap();
|
||||
for _ in 0..16 {
|
||||
msgs.push(serde_json::json!({"role": "user", "content": long}));
|
||||
msgs.push(serde_json::json!({"role": "assistant", "content": long}));
|
||||
}
|
||||
msgs.push(serde_json::json!({"role": "user", "content": "what is the latest?"}));
|
||||
serde_json::from_value(body).expect("valid ChatCompletionRequest")
|
||||
}
|
||||
|
||||
fn long_anthropic_request() -> anthropic::MessageCreateRequest {
|
||||
let long = long_text();
|
||||
let mut body = serde_json::json!({
|
||||
"model": "claude-sonnet-5",
|
||||
"max_tokens": 1024,
|
||||
"messages": [],
|
||||
});
|
||||
let msgs = body["messages"].as_array_mut().unwrap();
|
||||
for _ in 0..16 {
|
||||
msgs.push(serde_json::json!({"role": "user", "content": long}));
|
||||
msgs.push(serde_json::json!({"role": "assistant", "content": long}));
|
||||
}
|
||||
msgs.push(serde_json::json!({"role": "user", "content": "what is the latest?"}));
|
||||
serde_json::from_value(body).expect("valid MessageCreateRequest")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_forwards_openai_body_unchanged() {
|
||||
let state = minimal_state(Mode::Shadow);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(before, after, "shadow mode must forward the original body");
|
||||
assert_eq!(
|
||||
state.metrics.snapshot().optimizer_compressed_total,
|
||||
0,
|
||||
"shadow mode must never record a metrics-visible compression"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shadow_mode_forwards_anthropic_body_unchanged() {
|
||||
let state = minimal_state(Mode::Shadow);
|
||||
let mut req = long_anthropic_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_anthropic(&mut req, "messages");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(before, after, "shadow mode must forward the original body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_openai_history_and_preserves_latest() {
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(
|
||||
before["messages"].as_array().unwrap().last(),
|
||||
after["messages"].as_array().unwrap().last(),
|
||||
"the latest turn must never be rewritten"
|
||||
);
|
||||
assert_eq!(
|
||||
state.metrics.snapshot().optimizer_compressed_total,
|
||||
1,
|
||||
"an applied Live compression must be recorded in metrics"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_mode_compresses_anthropic_history_and_preserves_latest() {
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req = long_anthropic_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_anthropic(&mut req, "messages");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_eq!(
|
||||
before["messages"].as_array().unwrap().last(),
|
||||
after["messages"].as_array().unwrap().last(),
|
||||
"the latest turn must never be rewritten"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn off_mode_is_noop_and_engine_absent_is_noop() {
|
||||
// Off mode: engine present, mode Off -> never applied.
|
||||
let state = minimal_state(Mode::Off);
|
||||
let mut req = long_openai_request();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
|
||||
// No engine at all (e.g. non-Anthropic/Translate mode backend): no panic, no-op.
|
||||
let mut state_no_engine = minimal_state(Mode::Live);
|
||||
state_no_engine.optimizer = None;
|
||||
let mut req2 = long_openai_request();
|
||||
let before2 = serde_json::to_value(&req2).unwrap();
|
||||
state_no_engine.apply_optimizer_to_openai(&mut req2, "chat_completions");
|
||||
assert_eq!(before2, serde_json::to_value(&req2).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_history_below_min_len_gate_is_a_noop_not_a_panic() {
|
||||
// A short request has nothing worth compressing (below FFEC's min-length
|
||||
// gate) -- the seam must still round-trip cleanly without panicking or
|
||||
// corrupting the body, i.e. it fails open when there's nothing to do.
|
||||
let state = minimal_state(Mode::Live);
|
||||
let mut req: openai::ChatCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}))
|
||||
.unwrap();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_optimizer_to_openai(&mut req, "chat_completions");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
}
|
||||
|
||||
fn state_with_rtk(enabled: bool) -> AppState {
|
||||
let config = Config {
|
||||
backend: BackendKind::OpenAI,
|
||||
openai_api_key: "test".into(),
|
||||
openai_base_url: "https://api.openai.com".into(),
|
||||
listen_port: 3000,
|
||||
model_mapping: ModelMapping {
|
||||
big_model: "gpt-4o".into(),
|
||||
small_model: "gpt-4o-mini".into(),
|
||||
},
|
||||
tls: TlsConfig::default(),
|
||||
backend_auth: BackendAuth::BearerToken("test".into()),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
expose_degradation_warnings: false,
|
||||
openai_api_format: OpenAIApiFormat::Chat,
|
||||
provider_id: None,
|
||||
};
|
||||
let backend = crate::backend::BackendClient::OpenAI(
|
||||
crate::backend::openai_client::OpenAIClient::new(&config),
|
||||
);
|
||||
let runtime_config = Arc::new(RwLock::new(RuntimeConfig {
|
||||
model_mappings: indexmap::IndexMap::new(),
|
||||
log_level: "info".to_string(),
|
||||
log_bodies: false,
|
||||
redact_secrets: false,
|
||||
anthropic_thinking_repair: false,
|
||||
pxpipe_compress: false,
|
||||
pxpipe_models: String::new(),
|
||||
rtk_compress: enabled,
|
||||
rtk_models: String::new(),
|
||||
forward_client_auth: false,
|
||||
tool_guardrail_mode: "disabled".to_string(),
|
||||
optimizer_mode: "off".to_string(),
|
||||
}));
|
||||
AppState {
|
||||
backend,
|
||||
metrics: Metrics::new(),
|
||||
runtime_config,
|
||||
shared: None,
|
||||
route_options: None,
|
||||
backend_name: "openai".to_string(),
|
||||
provider_id: None,
|
||||
concurrency: Arc::new(Semaphore::new(64)),
|
||||
omit_stream_options: false,
|
||||
stream_timeout_secs: 0,
|
||||
expose_degradation_warnings: false,
|
||||
cache: None,
|
||||
thinking_repair: None,
|
||||
pxpipe: None,
|
||||
rtk: Some(Arc::new(crate::rtk::RtkEngine::new())),
|
||||
optimizer: None,
|
||||
model_router: None,
|
||||
provider_catalog: Arc::new(ProviderCatalog::bundled()),
|
||||
all_backends: None,
|
||||
tool_engine: None,
|
||||
batch_engine: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn request_with_noisy_tool_output() -> openai::ChatCompletionRequest {
|
||||
let mut noise = String::from("On branch main\nChanges not staged for commit:\n");
|
||||
for i in 0..200 {
|
||||
noise.push_str(&format!(" (use \"git add ...\" file {i})\n"));
|
||||
}
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"model": "gpt-4o",
|
||||
"messages": [
|
||||
{"role": "user", "content": "run git status"},
|
||||
{"role": "assistant", "content": null, "tool_calls": [
|
||||
{"id": "t1", "type": "function",
|
||||
"function": {"name": "bash", "arguments": "{\"cmd\":\"git status\"}"}}
|
||||
]},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": noise},
|
||||
],
|
||||
}))
|
||||
.expect("valid ChatCompletionRequest")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enabled_compresses_tool_output_and_records_metrics() {
|
||||
let state = state_with_rtk(true);
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
let after = serde_json::to_value(&req).unwrap();
|
||||
assert_ne!(before, after, "tool output should have been compressed");
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_is_a_noop() {
|
||||
let state = state_with_rtk(false);
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_absent_is_a_noop() {
|
||||
let mut state = state_with_rtk(true);
|
||||
state.rtk = None;
|
||||
let mut req = request_with_noisy_tool_output();
|
||||
let before = serde_json::to_value(&req).unwrap();
|
||||
state.apply_rtk_to_openai(&mut req, "gpt-4o");
|
||||
assert_eq!(before, serde_json::to_value(&req).unwrap());
|
||||
assert_eq!(state.metrics.snapshot().rtk_compressed_total, 0);
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
use crate::metrics::Metrics;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Global state for the multi-backend metrics endpoint.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct GlobalState {
|
||||
pub(crate) backend_metrics: Arc<HashMap<String, Metrics>>,
|
||||
}
|
||||
|
||||
/// Wrapper so OwnedSemaphorePermit can be stored in request extensions.
|
||||
/// The field is never read directly; it exists as an RAII guard to hold
|
||||
/// the permit until the struct is dropped.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ConcurrencyPermit(
|
||||
#[allow(dead_code)] pub(crate) Arc<tokio::sync::OwnedSemaphorePermit>,
|
||||
);
|
||||
@@ -0,0 +1,15 @@
|
||||
// Shared state types for request handlers: AppState, AnthropicJson, ResolvedModel, etc.
|
||||
// Extracted from routes.rs so consumers can import state independently of the router setup.
|
||||
|
||||
pub(crate) mod anthropic_json;
|
||||
pub(crate) mod app_state;
|
||||
pub(crate) mod compression;
|
||||
pub(crate) mod concurrency;
|
||||
pub(crate) mod resolved_model;
|
||||
pub(crate) mod tool_engine;
|
||||
|
||||
pub(crate) use anthropic_json::AnthropicJson;
|
||||
pub use app_state::AppState;
|
||||
pub(crate) use concurrency::{ConcurrencyPermit, GlobalState};
|
||||
pub(crate) use resolved_model::ResolvedModel;
|
||||
pub use tool_engine::ToolEngineState;
|
||||
@@ -0,0 +1,21 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Result of resolving a model name through the model router.
|
||||
pub(crate) enum ResolvedModel {
|
||||
/// Routed via model_list to a specific backend and actual model name.
|
||||
Routed {
|
||||
backend_name: String,
|
||||
model: String,
|
||||
/// The deployment Arc for recording in-flight/latency stats.
|
||||
deployment: Arc<crate::config::model_router::Deployment>,
|
||||
/// Per-route option overrides when routed via a DB route; `None` when
|
||||
/// routed via the LiteLLM model_router (inherit global config).
|
||||
options: Option<Arc<crate::config::route_router::RouteOptions>>,
|
||||
},
|
||||
/// Model is known but all deployments are at their RPM limit.
|
||||
AllAtLimit,
|
||||
/// Model router is active but the model alias is not configured.
|
||||
UnknownModel,
|
||||
/// No model router, or model not in router. Used legacy ModelMapping.
|
||||
Legacy(String),
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Shared state for tool execution, stored in AppState.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolEngineState {
|
||||
pub registry: Arc<crate::tools::ToolRegistry>,
|
||||
pub policy: Arc<crate::tools::ToolExecutionPolicy>,
|
||||
pub loop_config: crate::tools::LoopConfig,
|
||||
pub guardrails: crate::tools::ToolGuardrailConfig,
|
||||
pub mcp_manager: Option<Arc<crate::tools::McpServerManager>>,
|
||||
}
|
||||
@@ -248,335 +248,5 @@ fn thinking_eq(a: &ContentBlock, b: &ContentBlock) -> bool {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::thinking_repair::store::ThinkingRepairStore;
|
||||
use anyllm_translate::anthropic::{Content, InputMessage, MessageCreateRequest};
|
||||
use serde_json::json;
|
||||
|
||||
fn thinking(text: &str, sig: &str) -> ContentBlock {
|
||||
ContentBlock::Thinking {
|
||||
thinking: text.to_string(),
|
||||
signature: Some(sig.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_use(id: &str) -> ContentBlock {
|
||||
ContentBlock::ToolUse {
|
||||
id: id.to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
input: json!({"city": "nyc"}),
|
||||
}
|
||||
}
|
||||
|
||||
fn text(s: &str) -> ContentBlock {
|
||||
ContentBlock::Text {
|
||||
text: s.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn req_with_last_assistant(blocks: Vec<ContentBlock>) -> MessageCreateRequest {
|
||||
MessageCreateRequest {
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
max_tokens: 1024,
|
||||
messages: vec![
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("hi".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
content: Content::Blocks(blocks),
|
||||
},
|
||||
],
|
||||
system: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
stop_sequences: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
metadata: None,
|
||||
thinking: None,
|
||||
stream: None,
|
||||
extra: serde_json::Map::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn last_blocks(req: &MessageCreateRequest) -> &Vec<ContentBlock> {
|
||||
match &req.messages.last().unwrap().content {
|
||||
Content::Blocks(b) => b,
|
||||
Content::Text(_) => panic!("expected blocks"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier0_byte_identical_passthrough() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("hmm", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![thinking("hmm", "sig_1"), tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"byte-identical replay should not be touched"
|
||||
);
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier1_restores_mutated_text_under_known_signature() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("original thought", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("merged garbled thought", "sig_1"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("restored 1"));
|
||||
match &last_blocks(&req)[0] {
|
||||
ContentBlock::Thinking { thinking, .. } => assert_eq!(thinking, "original thought"),
|
||||
other => panic!("expected restored Thinking block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier2_drops_intruder_block_from_different_message() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("owner thought", "sig_owner"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
store
|
||||
.commit("ns1", "msg_2", vec![thinking("other thought", "sig_other")])
|
||||
.await;
|
||||
|
||||
// Replay claims to be msg_1's turn (tool_use toolu_1) but carries an
|
||||
// intruder thinking block whose signature belongs to msg_2.
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("other thought", "sig_other"),
|
||||
thinking("owner thought", "sig_owner"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("dropped 1"));
|
||||
let blocks = last_blocks(&req);
|
||||
assert_eq!(blocks.len(), 2);
|
||||
assert!(
|
||||
matches!(&blocks[0], ContentBlock::Thinking { signature: Some(s), .. } if s == "sig_owner")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_signature_with_known_owner_is_dropped() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("owner thought", "sig_owner"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("garbage, no matching record", "sig_unknown"),
|
||||
thinking("owner thought", "sig_owner"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("dropped 1"));
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_signature_with_no_owner_fails_open() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
// No tool_use in this turn, so no owner is resolvable.
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("standalone thought", "sig_unknown"),
|
||||
text("done"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"no owner evidence -> fail open, pass through"
|
||||
);
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reinserts_lost_redacted_thinking_block() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
let redacted = ContentBlock::RedactedThinking {
|
||||
data: "encrypted-blob".to_string(),
|
||||
};
|
||||
store
|
||||
.commit("ns1", "msg_1", vec![redacted.clone(), tool_use("toolu_1")])
|
||||
.await;
|
||||
|
||||
// Client replay is missing the redacted_thinking block entirely
|
||||
// (Claude Code doesn't persist it to JSONL).
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
let msg = result.unwrap();
|
||||
assert!(msg.contains("reinserted"), "got: {msg}");
|
||||
let blocks = last_blocks(&req);
|
||||
assert_eq!(blocks.len(), 2);
|
||||
assert!(
|
||||
matches!(&blocks[0], ContentBlock::RedactedThinking { data } if data == "encrypted-blob")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn surplus_current_nonthinking_block_disables_reinsert() {
|
||||
// rec has [thinking, tool_use_1]; current turn replays tool_use_1
|
||||
// AND adds an extra tool_use_2 that was never recorded, while also
|
||||
// dropping the thinking block. Non-thinking counts don't line up
|
||||
// (2 vs. 1), so positional reinsertion would either drop toolu_2 or
|
||||
// misalign it against rec -- must fail open instead.
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("hmm", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1"), tool_use("toolu_2")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
let blocks = last_blocks(&req);
|
||||
|
||||
assert!(result.is_none(), "count mismatch should fail open");
|
||||
assert_eq!(
|
||||
blocks.len(),
|
||||
2,
|
||||
"current turn's blocks must be preserved as-is"
|
||||
);
|
||||
assert!(
|
||||
blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { id, .. } if id == "toolu_2")),
|
||||
"toolu_2 must not be dropped"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_current_nonthinking_block_disables_reinsert() {
|
||||
// rec has [thinking, tool_use_1, tool_use_2]; current turn only
|
||||
// replays tool_use_1 (dropped tool_use_2 AND the thinking block).
|
||||
// Non-thinking counts don't line up (1 vs. 2), so reinserting
|
||||
// positionally would resurrect stale toolu_2 -- must fail open.
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![
|
||||
thinking("hmm", "sig_1"),
|
||||
tool_use("toolu_1"),
|
||||
tool_use("toolu_2"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
let blocks = last_blocks(&req);
|
||||
|
||||
assert!(result.is_none(), "count mismatch should fail open");
|
||||
assert_eq!(
|
||||
blocks.len(),
|
||||
1,
|
||||
"current turn's blocks must be preserved as-is"
|
||||
);
|
||||
assert!(
|
||||
!blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { id, .. } if id == "toolu_2")),
|
||||
"stale recorded toolu_2 must not be reinserted"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_before_last_assistant_are_never_touched() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit("ns1", "msg_2", vec![thinking("turn 2", "sig_2")])
|
||||
.await;
|
||||
|
||||
let mut req = MessageCreateRequest {
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
max_tokens: 1024,
|
||||
messages: vec![
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("first".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
// Deliberately mutated/corrupt — must be left alone
|
||||
// because it is NOT the last assistant message.
|
||||
content: Content::Blocks(vec![thinking("mutated turn 1", "sig_1_unknown")]),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("second".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
content: Content::Blocks(vec![thinking("turn 2", "sig_2")]),
|
||||
},
|
||||
],
|
||||
system: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
stop_sequences: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
metadata: None,
|
||||
thinking: None,
|
||||
stream: None,
|
||||
extra: serde_json::Map::new(),
|
||||
};
|
||||
|
||||
repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
match &req.messages[1].content {
|
||||
Content::Blocks(b) => match &b[0] {
|
||||
ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
} => {
|
||||
assert_eq!(thinking, "mutated turn 1");
|
||||
assert_eq!(signature.as_deref(), Some("sig_1_unknown"));
|
||||
}
|
||||
other => panic!("unexpected block: {other:?}"),
|
||||
},
|
||||
Content::Text(_) => panic!("expected blocks"),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[path = "repair/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
use super::*;
|
||||
use crate::thinking_repair::store::ThinkingRepairStore;
|
||||
use anyllm_translate::anthropic::{Content, InputMessage, MessageCreateRequest};
|
||||
use serde_json::json;
|
||||
|
||||
fn thinking(text: &str, sig: &str) -> ContentBlock {
|
||||
ContentBlock::Thinking {
|
||||
thinking: text.to_string(),
|
||||
signature: Some(sig.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_use(id: &str) -> ContentBlock {
|
||||
ContentBlock::ToolUse {
|
||||
id: id.to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
input: json!({"city": "nyc"}),
|
||||
}
|
||||
}
|
||||
|
||||
fn text(s: &str) -> ContentBlock {
|
||||
ContentBlock::Text {
|
||||
text: s.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn req_with_last_assistant(blocks: Vec<ContentBlock>) -> MessageCreateRequest {
|
||||
MessageCreateRequest {
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
max_tokens: 1024,
|
||||
messages: vec![
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("hi".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
content: Content::Blocks(blocks),
|
||||
},
|
||||
],
|
||||
system: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
stop_sequences: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
metadata: None,
|
||||
thinking: None,
|
||||
stream: None,
|
||||
extra: serde_json::Map::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn last_blocks(req: &MessageCreateRequest) -> &Vec<ContentBlock> {
|
||||
match &req.messages.last().unwrap().content {
|
||||
Content::Blocks(b) => b,
|
||||
Content::Text(_) => panic!("expected blocks"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier0_byte_identical_passthrough() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("hmm", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![thinking("hmm", "sig_1"), tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"byte-identical replay should not be touched"
|
||||
);
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier1_restores_mutated_text_under_known_signature() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("original thought", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("merged garbled thought", "sig_1"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("restored 1"));
|
||||
match &last_blocks(&req)[0] {
|
||||
ContentBlock::Thinking { thinking, .. } => assert_eq!(thinking, "original thought"),
|
||||
other => panic!("expected restored Thinking block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier2_drops_intruder_block_from_different_message() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("owner thought", "sig_owner"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
store
|
||||
.commit("ns1", "msg_2", vec![thinking("other thought", "sig_other")])
|
||||
.await;
|
||||
|
||||
// Replay claims to be msg_1's turn (tool_use toolu_1) but carries an
|
||||
// intruder thinking block whose signature belongs to msg_2.
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("other thought", "sig_other"),
|
||||
thinking("owner thought", "sig_owner"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("dropped 1"));
|
||||
let blocks = last_blocks(&req);
|
||||
assert_eq!(blocks.len(), 2);
|
||||
assert!(
|
||||
matches!(&blocks[0], ContentBlock::Thinking { signature: Some(s), .. } if s == "sig_owner")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_signature_with_known_owner_is_dropped() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("owner thought", "sig_owner"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("garbage, no matching record", "sig_unknown"),
|
||||
thinking("owner thought", "sig_owner"),
|
||||
tool_use("toolu_1"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(result.unwrap().contains("dropped 1"));
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_signature_with_no_owner_fails_open() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
// No tool_use in this turn, so no owner is resolvable.
|
||||
let mut req = req_with_last_assistant(vec![
|
||||
thinking("standalone thought", "sig_unknown"),
|
||||
text("done"),
|
||||
]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"no owner evidence -> fail open, pass through"
|
||||
);
|
||||
assert_eq!(last_blocks(&req).len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reinserts_lost_redacted_thinking_block() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
let redacted = ContentBlock::RedactedThinking {
|
||||
data: "encrypted-blob".to_string(),
|
||||
};
|
||||
store
|
||||
.commit("ns1", "msg_1", vec![redacted.clone(), tool_use("toolu_1")])
|
||||
.await;
|
||||
|
||||
// Client replay is missing the redacted_thinking block entirely
|
||||
// (Claude Code doesn't persist it to JSONL).
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
let msg = result.unwrap();
|
||||
assert!(msg.contains("reinserted"), "got: {msg}");
|
||||
let blocks = last_blocks(&req);
|
||||
assert_eq!(blocks.len(), 2);
|
||||
assert!(
|
||||
matches!(&blocks[0], ContentBlock::RedactedThinking { data } if data == "encrypted-blob")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn surplus_current_nonthinking_block_disables_reinsert() {
|
||||
// rec has [thinking, tool_use_1]; current turn replays tool_use_1
|
||||
// AND adds an extra tool_use_2 that was never recorded, while also
|
||||
// dropping the thinking block. Non-thinking counts don't line up
|
||||
// (2 vs. 1), so positional reinsertion would either drop toolu_2 or
|
||||
// misalign it against rec -- must fail open instead.
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![thinking("hmm", "sig_1"), tool_use("toolu_1")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1"), tool_use("toolu_2")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
let blocks = last_blocks(&req);
|
||||
|
||||
assert!(result.is_none(), "count mismatch should fail open");
|
||||
assert_eq!(
|
||||
blocks.len(),
|
||||
2,
|
||||
"current turn's blocks must be preserved as-is"
|
||||
);
|
||||
assert!(
|
||||
blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { id, .. } if id == "toolu_2")),
|
||||
"toolu_2 must not be dropped"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_current_nonthinking_block_disables_reinsert() {
|
||||
// rec has [thinking, tool_use_1, tool_use_2]; current turn only
|
||||
// replays tool_use_1 (dropped tool_use_2 AND the thinking block).
|
||||
// Non-thinking counts don't line up (1 vs. 2), so reinserting
|
||||
// positionally would resurrect stale toolu_2 -- must fail open.
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit(
|
||||
"ns1",
|
||||
"msg_1",
|
||||
vec![
|
||||
thinking("hmm", "sig_1"),
|
||||
tool_use("toolu_1"),
|
||||
tool_use("toolu_2"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut req = req_with_last_assistant(vec![tool_use("toolu_1")]);
|
||||
let result = repair_request(&store, "ns1", &mut req).await;
|
||||
let blocks = last_blocks(&req);
|
||||
|
||||
assert!(result.is_none(), "count mismatch should fail open");
|
||||
assert_eq!(
|
||||
blocks.len(),
|
||||
1,
|
||||
"current turn's blocks must be preserved as-is"
|
||||
);
|
||||
assert!(
|
||||
!blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { id, .. } if id == "toolu_2")),
|
||||
"stale recorded toolu_2 must not be reinserted"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_before_last_assistant_are_never_touched() {
|
||||
let store = ThinkingRepairStore::new();
|
||||
store
|
||||
.commit("ns1", "msg_2", vec![thinking("turn 2", "sig_2")])
|
||||
.await;
|
||||
|
||||
let mut req = MessageCreateRequest {
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
max_tokens: 1024,
|
||||
messages: vec![
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("first".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
// Deliberately mutated/corrupt — must be left alone
|
||||
// because it is NOT the last assistant message.
|
||||
content: Content::Blocks(vec![thinking("mutated turn 1", "sig_1_unknown")]),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::User,
|
||||
content: Content::Text("second".to_string()),
|
||||
},
|
||||
InputMessage {
|
||||
role: Role::Assistant,
|
||||
content: Content::Blocks(vec![thinking("turn 2", "sig_2")]),
|
||||
},
|
||||
],
|
||||
system: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
stop_sequences: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
metadata: None,
|
||||
thinking: None,
|
||||
stream: None,
|
||||
extra: serde_json::Map::new(),
|
||||
};
|
||||
|
||||
repair_request(&store, "ns1", &mut req).await;
|
||||
|
||||
match &req.messages[1].content {
|
||||
Content::Blocks(b) => match &b[0] {
|
||||
ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
} => {
|
||||
assert_eq!(thinking, "mutated turn 1");
|
||||
assert_eq!(signature.as_deref(), Some("sig_1_unknown"));
|
||||
}
|
||||
other => panic!("unexpected block: {other:?}"),
|
||||
},
|
||||
Content::Text(_) => panic!("expected blocks"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
//! Tool guardrail logic: partition tool calls into auto-execute / pass-through /
|
||||
//! denied buckets, build denial + nudge results, and the combined
|
||||
//! `partition_and_nudge` shared by both tool loops.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use super::{ToolCall, ToolResult};
|
||||
use crate::tools::policy::{PolicyAction, ToolExecutionPolicy};
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
use crate::tools::trace::ToolOutcome;
|
||||
use crate::tools::{evaluate_tool_guardrails, ToolGuardrailNudge, ToolGuardrailRequestState};
|
||||
|
||||
/// Partition tool calls into three categories.
|
||||
///
|
||||
/// A tool call is only eligible for server-side policy evaluation when the
|
||||
/// proxy advertised that exact tool name for this request. This prevents
|
||||
/// client-supplied tool schemas from reusing privileged server tool names.
|
||||
///
|
||||
/// - `auto_execute`: proxy-advertised AND in registry AND policy says Allow
|
||||
/// - `pass_through`: not proxy-advertised, not in registry, OR policy says PassThrough
|
||||
/// - `denied`: proxy-advertised AND policy says Deny
|
||||
pub fn partition_tool_calls<'a>(
|
||||
tool_calls: &'a [ToolCall],
|
||||
registry: &ToolRegistry,
|
||||
policy: &ToolExecutionPolicy,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
) -> (Vec<&'a ToolCall>, Vec<&'a ToolCall>, Vec<&'a ToolCall>) {
|
||||
let mut auto_execute = Vec::new();
|
||||
let mut pass_through = Vec::new();
|
||||
let mut denied = Vec::new();
|
||||
|
||||
for call in tool_calls {
|
||||
if !server_advertised_tool_names.contains(&call.name) {
|
||||
pass_through.push(call);
|
||||
continue;
|
||||
}
|
||||
|
||||
match policy.resolve(&call.name) {
|
||||
PolicyAction::Deny => denied.push(call),
|
||||
PolicyAction::Allow if registry.contains(&call.name) => auto_execute.push(call),
|
||||
// Allow but not in registry, or PassThrough
|
||||
_ => pass_through.push(call),
|
||||
}
|
||||
}
|
||||
|
||||
(auto_execute, pass_through, denied)
|
||||
}
|
||||
|
||||
/// Build error ToolResults for denied tool calls.
|
||||
pub fn denied_tool_results(denied: &[&ToolCall]) -> Vec<ToolResult> {
|
||||
denied
|
||||
.iter()
|
||||
.map(|call| ToolResult {
|
||||
tool_use_id: call.id.clone(),
|
||||
tool_name: call.name.clone(),
|
||||
outcome: ToolOutcome::Error {
|
||||
message: format!("Tool '{}' is denied by policy", call.name),
|
||||
retryable: false,
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Build retryable ToolResults for guardrail nudges.
|
||||
///
|
||||
/// Only the tool calls a nudge targets get a result; every other call in the
|
||||
/// batch is left for normal execution/pass-through. `calls` supplies the tool
|
||||
/// name for each nudged `call_id`.
|
||||
pub fn guardrail_nudge_results(
|
||||
calls: &[ToolCall],
|
||||
nudges: &[ToolGuardrailNudge],
|
||||
) -> Vec<ToolResult> {
|
||||
nudges
|
||||
.iter()
|
||||
.filter_map(|nudge| {
|
||||
let call = calls.iter().find(|c| c.id == nudge.call_id)?;
|
||||
Some(ToolResult {
|
||||
tool_use_id: call.id.clone(),
|
||||
tool_name: call.name.clone(),
|
||||
outcome: ToolOutcome::Error {
|
||||
message: format!("[ToolCallPolicyNudge] {}", nudge.content),
|
||||
retryable: true,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Partition `tool_calls`, then evaluate guardrail nudges only against the
|
||||
/// calls that would otherwise be auto-executed by this proxy.
|
||||
///
|
||||
/// A nudge target must be something the proxy actually owns: guardrails are
|
||||
/// an advisory retry mechanism that answers the model on the proxy's behalf
|
||||
/// (a synthetic tool_result + a follow-up backend call), so applying it to a
|
||||
/// pass-through call (not in the registry / not proxy-advertised, meaning the
|
||||
/// *caller* owns and expects to execute it -- e.g. a Claude-Code-style client's
|
||||
/// own Bash/Grep/Edit/Write tool) would silently swallow that call's real
|
||||
/// tool_use turn instead of returning it to the caller. Restricting
|
||||
/// `evaluate_tool_guardrails` to the post-partition `auto_exec` set keeps
|
||||
/// nudges impossible to apply to anything the proxy doesn't own, and
|
||||
/// `denied` never needs nudge-filtering since nudges only ever originate
|
||||
/// from `auto_exec` now (the two sets are disjoint by construction).
|
||||
///
|
||||
/// Shared by the streaming and non-streaming tool loops so a fix here can't
|
||||
/// drift between the two copies.
|
||||
pub fn partition_and_nudge<'a>(
|
||||
tool_calls: &'a [ToolCall],
|
||||
tool_specs: &[anyllm_translate::anthropic::Tool],
|
||||
registry: &ToolRegistry,
|
||||
policy: &ToolExecutionPolicy,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
guardrails: &crate::tools::ToolGuardrailConfig,
|
||||
guardrail_state: &mut ToolGuardrailRequestState,
|
||||
) -> (Vec<&'a ToolCall>, Vec<ToolResult>, Vec<ToolResult>) {
|
||||
let (auto_exec, _pass_through, denied) =
|
||||
partition_tool_calls(tool_calls, registry, policy, server_advertised_tool_names);
|
||||
|
||||
let auto_exec_owned: Vec<ToolCall> = auto_exec.iter().map(|c| (*c).clone()).collect();
|
||||
let nudges =
|
||||
evaluate_tool_guardrails(&auto_exec_owned, tool_specs, guardrails, guardrail_state);
|
||||
let nudged_ids: HashSet<&str> = nudges.iter().map(|n| n.call_id.as_str()).collect();
|
||||
let nudge_results = guardrail_nudge_results(&auto_exec_owned, &nudges);
|
||||
|
||||
let auto_exec: Vec<&ToolCall> = auto_exec
|
||||
.into_iter()
|
||||
.filter(|c| !nudged_ids.contains(c.id.as_str()))
|
||||
.collect();
|
||||
let denied_results = denied_tool_results(&denied);
|
||||
|
||||
(auto_exec, nudge_results, denied_results)
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
//! Parsers/converters between Anthropic message types and `ToolCall`/`ToolResult`.
|
||||
|
||||
use super::{ToolCall, ToolResult};
|
||||
use crate::tools::trace::ToolOutcome;
|
||||
|
||||
/// Extract ToolCall structs from an Anthropic MessageResponse.
|
||||
pub fn extract_tool_calls(
|
||||
response: &anyllm_translate::anthropic::MessageResponse,
|
||||
) -> Vec<ToolCall> {
|
||||
response
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| {
|
||||
if let anyllm_translate::anthropic::ContentBlock::ToolUse { id, name, input } = block {
|
||||
Some(ToolCall {
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
input: input.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Convert tool execution results to an Anthropic user message with ToolResult blocks.
|
||||
pub fn tool_results_to_user_message(
|
||||
results: &[ToolResult],
|
||||
) -> anyllm_translate::anthropic::InputMessage {
|
||||
let blocks: Vec<anyllm_translate::anthropic::ContentBlock> = results
|
||||
.iter()
|
||||
.map(|r| {
|
||||
let (content_text, is_error) = match &r.outcome {
|
||||
ToolOutcome::Success(v) => (serde_json::to_string(v).unwrap_or_default(), false),
|
||||
ToolOutcome::Error { message, .. } => (message.clone(), true),
|
||||
ToolOutcome::Timeout => ("Tool execution timed out".to_string(), true),
|
||||
};
|
||||
anyllm_translate::anthropic::ContentBlock::ToolResult {
|
||||
tool_use_id: r.tool_use_id.clone(),
|
||||
content: Some(anyllm_translate::anthropic::ToolResultContent::Text(
|
||||
content_text,
|
||||
)),
|
||||
is_error: Some(is_error),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
anyllm_translate::anthropic::InputMessage {
|
||||
role: anyllm_translate::anthropic::Role::User,
|
||||
content: anyllm_translate::anthropic::Content::Blocks(blocks),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a MessageResponse's content into an assistant InputMessage for conversation history.
|
||||
pub fn response_to_assistant_message(
|
||||
response: &anyllm_translate::anthropic::MessageResponse,
|
||||
) -> anyllm_translate::anthropic::InputMessage {
|
||||
anyllm_translate::anthropic::InputMessage {
|
||||
role: anyllm_translate::anthropic::Role::Assistant,
|
||||
content: anyllm_translate::anthropic::Content::Blocks(response.content.clone()),
|
||||
}
|
||||
}
|
||||
@@ -1,14 +1,28 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
//! Tool execution: the shared types plus the guardrail, runner, and helper
|
||||
//! submodules that implement partitioning, parallel execution, and the bounded
|
||||
//! non-streaming tool loop.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::Value;
|
||||
use tokio::task::JoinSet;
|
||||
|
||||
use crate::tools::policy::{PolicyAction, ToolExecutionPolicy};
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
use crate::tools::trace::ToolOutcome;
|
||||
use crate::tools::{evaluate_tool_guardrails, ToolGuardrailNudge, ToolGuardrailRequestState};
|
||||
|
||||
mod guardrail;
|
||||
mod helpers;
|
||||
mod runner;
|
||||
|
||||
pub use guardrail::{
|
||||
denied_tool_results, guardrail_nudge_results, partition_and_nudge, partition_tool_calls,
|
||||
};
|
||||
pub use helpers::{
|
||||
extract_tool_calls, response_to_assistant_message, tool_results_to_user_message,
|
||||
};
|
||||
pub use runner::{execute_tool_calls, execute_tool_calls_timed, is_duplicate, maybe_execute_tools};
|
||||
|
||||
/// Engine state needed by `maybe_execute_tools`. Re-exported alias so callers
|
||||
/// do not need to reach into `server::state`.
|
||||
pub use crate::server::state::ToolEngineState;
|
||||
|
||||
/// A tool call extracted from an LLM response.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -46,555 +60,5 @@ impl Default for LoopConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Partition tool calls into three categories.
|
||||
///
|
||||
/// A tool call is only eligible for server-side policy evaluation when the
|
||||
/// proxy advertised that exact tool name for this request. This prevents
|
||||
/// client-supplied tool schemas from reusing privileged server tool names.
|
||||
///
|
||||
/// - `auto_execute`: proxy-advertised AND in registry AND policy says Allow
|
||||
/// - `pass_through`: not proxy-advertised, not in registry, OR policy says PassThrough
|
||||
/// - `denied`: proxy-advertised AND policy says Deny
|
||||
pub fn partition_tool_calls<'a>(
|
||||
tool_calls: &'a [ToolCall],
|
||||
registry: &ToolRegistry,
|
||||
policy: &ToolExecutionPolicy,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
) -> (Vec<&'a ToolCall>, Vec<&'a ToolCall>, Vec<&'a ToolCall>) {
|
||||
let mut auto_execute = Vec::new();
|
||||
let mut pass_through = Vec::new();
|
||||
let mut denied = Vec::new();
|
||||
|
||||
for call in tool_calls {
|
||||
if !server_advertised_tool_names.contains(&call.name) {
|
||||
pass_through.push(call);
|
||||
continue;
|
||||
}
|
||||
|
||||
match policy.resolve(&call.name) {
|
||||
PolicyAction::Deny => denied.push(call),
|
||||
PolicyAction::Allow if registry.contains(&call.name) => auto_execute.push(call),
|
||||
// Allow but not in registry, or PassThrough
|
||||
_ => pass_through.push(call),
|
||||
}
|
||||
}
|
||||
|
||||
(auto_execute, pass_through, denied)
|
||||
}
|
||||
|
||||
/// Build error ToolResults for denied tool calls.
|
||||
pub fn denied_tool_results(denied: &[&ToolCall]) -> Vec<ToolResult> {
|
||||
denied
|
||||
.iter()
|
||||
.map(|call| ToolResult {
|
||||
tool_use_id: call.id.clone(),
|
||||
tool_name: call.name.clone(),
|
||||
outcome: ToolOutcome::Error {
|
||||
message: format!("Tool '{}' is denied by policy", call.name),
|
||||
retryable: false,
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Build retryable ToolResults for guardrail nudges.
|
||||
///
|
||||
/// Only the tool calls a nudge targets get a result; every other call in the
|
||||
/// batch is left for normal execution/pass-through. `calls` supplies the tool
|
||||
/// name for each nudged `call_id`.
|
||||
pub fn guardrail_nudge_results(
|
||||
calls: &[ToolCall],
|
||||
nudges: &[ToolGuardrailNudge],
|
||||
) -> Vec<ToolResult> {
|
||||
nudges
|
||||
.iter()
|
||||
.filter_map(|nudge| {
|
||||
let call = calls.iter().find(|c| c.id == nudge.call_id)?;
|
||||
Some(ToolResult {
|
||||
tool_use_id: call.id.clone(),
|
||||
tool_name: call.name.clone(),
|
||||
outcome: ToolOutcome::Error {
|
||||
message: format!("[ToolCallPolicyNudge] {}", nudge.content),
|
||||
retryable: true,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Partition `tool_calls`, then evaluate guardrail nudges only against the
|
||||
/// calls that would otherwise be auto-executed by this proxy.
|
||||
///
|
||||
/// A nudge target must be something the proxy actually owns: guardrails are
|
||||
/// an advisory retry mechanism that answers the model on the proxy's behalf
|
||||
/// (a synthetic tool_result + a follow-up backend call), so applying it to a
|
||||
/// pass-through call (not in the registry / not proxy-advertised, meaning the
|
||||
/// *caller* owns and expects to execute it — e.g. a Claude-Code-style client's
|
||||
/// own Bash/Grep/Edit/Write tool) would silently swallow that call's real
|
||||
/// tool_use turn instead of returning it to the caller. Restricting
|
||||
/// `evaluate_tool_guardrails` to the post-partition `auto_exec` set keeps
|
||||
/// nudges impossible to apply to anything the proxy doesn't own, and
|
||||
/// `denied` never needs nudge-filtering since nudges only ever originate
|
||||
/// from `auto_exec` now (the two sets are disjoint by construction).
|
||||
///
|
||||
/// Shared by the streaming and non-streaming tool loops so a fix here can't
|
||||
/// drift between the two copies.
|
||||
pub fn partition_and_nudge<'a>(
|
||||
tool_calls: &'a [ToolCall],
|
||||
tool_specs: &[anyllm_translate::anthropic::Tool],
|
||||
registry: &ToolRegistry,
|
||||
policy: &ToolExecutionPolicy,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
guardrails: &crate::tools::ToolGuardrailConfig,
|
||||
guardrail_state: &mut ToolGuardrailRequestState,
|
||||
) -> (Vec<&'a ToolCall>, Vec<ToolResult>, Vec<ToolResult>) {
|
||||
let (auto_exec, _pass_through, denied) =
|
||||
partition_tool_calls(tool_calls, registry, policy, server_advertised_tool_names);
|
||||
|
||||
let auto_exec_owned: Vec<ToolCall> = auto_exec.iter().map(|c| (*c).clone()).collect();
|
||||
let nudges =
|
||||
evaluate_tool_guardrails(&auto_exec_owned, tool_specs, guardrails, guardrail_state);
|
||||
let nudged_ids: HashSet<&str> = nudges.iter().map(|n| n.call_id.as_str()).collect();
|
||||
let nudge_results = guardrail_nudge_results(&auto_exec_owned, &nudges);
|
||||
|
||||
let auto_exec: Vec<&ToolCall> = auto_exec
|
||||
.into_iter()
|
||||
.filter(|c| !nudged_ids.contains(c.id.as_str()))
|
||||
.collect();
|
||||
let denied_results = denied_tool_results(&denied);
|
||||
|
||||
(auto_exec, nudge_results, denied_results)
|
||||
}
|
||||
|
||||
/// Execute tool calls in parallel, respecting per-tool timeouts.
|
||||
///
|
||||
/// Results are returned in the same order as `calls`.
|
||||
pub async fn execute_tool_calls(
|
||||
calls: &[&ToolCall],
|
||||
registry: Arc<ToolRegistry>,
|
||||
policy: &ToolExecutionPolicy,
|
||||
config: &LoopConfig,
|
||||
) -> Vec<ToolResult> {
|
||||
let capped = &calls[..calls.len().min(config.max_tool_calls_per_turn)];
|
||||
|
||||
// Collect (original_index, ToolCall) to restore order after parallel execution.
|
||||
let indexed: Vec<(usize, &ToolCall)> = capped.iter().copied().enumerate().collect();
|
||||
|
||||
let mut join_set: JoinSet<(usize, ToolResult)> = JoinSet::new();
|
||||
|
||||
for (idx, call) in indexed {
|
||||
let timeout = policy
|
||||
.find_rule(&call.name)
|
||||
.and_then(|r| r.timeout)
|
||||
.unwrap_or(config.tool_timeout);
|
||||
|
||||
let registry = Arc::clone(®istry);
|
||||
let id = call.id.clone();
|
||||
let name = call.name.clone();
|
||||
let input = call.input.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
let result =
|
||||
tokio::time::timeout(timeout, execute_single(®istry, &name, input)).await;
|
||||
|
||||
let outcome = match result {
|
||||
Ok(Ok(value)) => ToolOutcome::Success(value),
|
||||
Ok(Err(msg)) => ToolOutcome::Error {
|
||||
message: msg,
|
||||
retryable: false,
|
||||
},
|
||||
Err(_elapsed) => ToolOutcome::Timeout,
|
||||
};
|
||||
|
||||
(
|
||||
idx,
|
||||
ToolResult {
|
||||
tool_use_id: id,
|
||||
tool_name: name,
|
||||
outcome,
|
||||
},
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
let mut collected: Vec<(usize, ToolResult)> = Vec::with_capacity(capped.len());
|
||||
while let Some(res) = join_set.join_next().await {
|
||||
match res {
|
||||
Ok(pair) => collected.push(pair),
|
||||
Err(e) => {
|
||||
// JoinError means the task panicked; treat as an error outcome.
|
||||
// We don't have the index here, so we skip (shouldn't happen in practice).
|
||||
tracing::error!("tool execution task panicked: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
collected.sort_by_key(|(idx, _)| *idx);
|
||||
collected.into_iter().map(|(_, r)| r).collect()
|
||||
}
|
||||
|
||||
/// Check whether two slices of ToolCall represent the same logical calls.
|
||||
///
|
||||
/// Same length, same multiset of (name, input) pairs. IDs are ignored.
|
||||
pub fn is_duplicate(a: &[ToolCall], b: &[ToolCall]) -> bool {
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut a_pairs: Vec<(&str, &Value)> = a.iter().map(|c| (c.name.as_str(), &c.input)).collect();
|
||||
let mut b_pairs: Vec<(&str, &Value)> = b.iter().map(|c| (c.name.as_str(), &c.input)).collect();
|
||||
|
||||
// Sort by name so comparison is order-independent.
|
||||
a_pairs.sort_by_key(|(name, _)| *name);
|
||||
b_pairs.sort_by_key(|(name, _)| *name);
|
||||
|
||||
a_pairs == b_pairs
|
||||
}
|
||||
|
||||
/// Execute a single tool by name, looking it up in the registry.
|
||||
async fn execute_single(
|
||||
registry: &ToolRegistry,
|
||||
tool_name: &str,
|
||||
input: Value,
|
||||
) -> Result<Value, String> {
|
||||
match registry.get(tool_name) {
|
||||
Some(tool) => tool.execute(input).await,
|
||||
None => Err(format!("tool '{}' not found in registry", tool_name)),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper functions for extracting tool calls and building follow-up messages
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Extract ToolCall structs from an Anthropic MessageResponse.
|
||||
pub fn extract_tool_calls(
|
||||
response: &anyllm_translate::anthropic::MessageResponse,
|
||||
) -> Vec<ToolCall> {
|
||||
response
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| {
|
||||
if let anyllm_translate::anthropic::ContentBlock::ToolUse { id, name, input } = block {
|
||||
Some(ToolCall {
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
input: input.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Convert tool execution results to an Anthropic user message with ToolResult blocks.
|
||||
pub fn tool_results_to_user_message(
|
||||
results: &[ToolResult],
|
||||
) -> anyllm_translate::anthropic::InputMessage {
|
||||
let blocks: Vec<anyllm_translate::anthropic::ContentBlock> = results
|
||||
.iter()
|
||||
.map(|r| {
|
||||
let (content_text, is_error) = match &r.outcome {
|
||||
ToolOutcome::Success(v) => (serde_json::to_string(v).unwrap_or_default(), false),
|
||||
ToolOutcome::Error { message, .. } => (message.clone(), true),
|
||||
ToolOutcome::Timeout => ("Tool execution timed out".to_string(), true),
|
||||
};
|
||||
anyllm_translate::anthropic::ContentBlock::ToolResult {
|
||||
tool_use_id: r.tool_use_id.clone(),
|
||||
content: Some(anyllm_translate::anthropic::ToolResultContent::Text(
|
||||
content_text,
|
||||
)),
|
||||
is_error: Some(is_error),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
anyllm_translate::anthropic::InputMessage {
|
||||
role: anyllm_translate::anthropic::Role::User,
|
||||
content: anyllm_translate::anthropic::Content::Blocks(blocks),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a MessageResponse's content into an assistant InputMessage for conversation history.
|
||||
pub fn response_to_assistant_message(
|
||||
response: &anyllm_translate::anthropic::MessageResponse,
|
||||
) -> anyllm_translate::anthropic::InputMessage {
|
||||
anyllm_translate::anthropic::InputMessage {
|
||||
role: anyllm_translate::anthropic::Role::Assistant,
|
||||
content: anyllm_translate::anthropic::Content::Blocks(response.content.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Timing wrapper used in the execution loop (available to callers)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Run `execute_tool_calls` and record wall-clock duration per call.
|
||||
/// Returns (results, elapsed_per_call). Exposed for loop-level tracing.
|
||||
pub async fn execute_tool_calls_timed(
|
||||
calls: &[&ToolCall],
|
||||
registry: Arc<ToolRegistry>,
|
||||
policy: &ToolExecutionPolicy,
|
||||
config: &LoopConfig,
|
||||
) -> (Vec<ToolResult>, Duration) {
|
||||
let start = Instant::now();
|
||||
let results = execute_tool_calls(calls, registry, policy, config).await;
|
||||
(results, start.elapsed())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Centralized bounded tool-execution loop for non-streaming requests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
use crate::tools::trace::{IterationTrace, LoopTrace, TerminationReason, ToolCallTrace};
|
||||
|
||||
/// Engine state needed by `maybe_execute_tools`. Re-exported alias so callers
|
||||
/// do not need to reach into `server::routes`.
|
||||
pub use crate::server::state::ToolEngineState;
|
||||
|
||||
/// Process an LLM response for tool execution. If auto-executable tool calls
|
||||
/// are found, executes them and makes follow-up backend calls in a bounded loop.
|
||||
///
|
||||
/// Returns the final response (original if no tools were auto-executed) and a
|
||||
/// `LoopTrace` recording what happened.
|
||||
///
|
||||
/// `backend_call` is a closure the caller provides. It takes a
|
||||
/// `MessageCreateRequest` and returns the translated `MessageResponse`.
|
||||
/// This keeps the loop backend-agnostic: the handler knows how to translate
|
||||
/// and call its specific backend; this function only knows about Anthropic types.
|
||||
///
|
||||
/// `guardrails` is the effective guardrail config for this request (the
|
||||
/// runtime-tunable override applied on top of `engine.guardrails`, the
|
||||
/// static per-process preset -- see `crate::tools::resolve_runtime_guardrails`
|
||||
/// and `AppState::effective_tool_guardrails`). Callers that don't need the
|
||||
/// runtime override can pass `&engine.guardrails` directly.
|
||||
pub async fn maybe_execute_tools<F, Fut>(
|
||||
engine: &ToolEngineState,
|
||||
original_req: &anyllm_translate::anthropic::MessageCreateRequest,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
initial_response: anyllm_translate::anthropic::MessageResponse,
|
||||
guardrails: &crate::tools::ToolGuardrailConfig,
|
||||
backend_call: F,
|
||||
) -> (anyllm_translate::anthropic::MessageResponse, LoopTrace)
|
||||
where
|
||||
F: Fn(anyllm_translate::anthropic::MessageCreateRequest) -> Fut,
|
||||
Fut: std::future::Future<Output = Result<anyllm_translate::anthropic::MessageResponse, String>>,
|
||||
{
|
||||
let loop_start = Instant::now();
|
||||
let mut iterations: Vec<IterationTrace> = Vec::new();
|
||||
let mut current_response = initial_response;
|
||||
let mut current_messages = original_req.messages.clone();
|
||||
let mut prev_tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut guardrail_state = ToolGuardrailRequestState::new();
|
||||
|
||||
for _iteration in 0..engine.loop_config.max_iterations {
|
||||
// Guard: total timeout
|
||||
if loop_start.elapsed() > engine.loop_config.total_timeout {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::Timeout,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let tool_calls = extract_tool_calls(¤t_response);
|
||||
|
||||
// Advisory guardrails: each nudge targets one offending call and comes
|
||||
// back as a retryable error result. Nudged calls are skipped this turn;
|
||||
// every other call is partitioned and executed as usual. Each distinct
|
||||
// decision nudges only once per request (across iterations), so a
|
||||
// model that ignores a nudge and repeats the call lets it proceed
|
||||
// instead of spinning the loop to max_iterations. Nudges only ever
|
||||
// apply to calls this proxy would actually execute -- see
|
||||
// `partition_and_nudge`'s doc comment.
|
||||
let (auto_exec, nudge_results, denied_results) = partition_and_nudge(
|
||||
&tool_calls,
|
||||
original_req.tools.as_deref().unwrap_or(&[]),
|
||||
&engine.registry,
|
||||
&engine.policy,
|
||||
server_advertised_tool_names,
|
||||
guardrails,
|
||||
&mut guardrail_state,
|
||||
);
|
||||
|
||||
if auto_exec.is_empty() && denied_results.is_empty() && nudge_results.is_empty() {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// If there is nothing to execute (only nudges and/or denials), send the
|
||||
// advisory/error results back to the LLM immediately without running a tool.
|
||||
if auto_exec.is_empty() {
|
||||
let mut results = nudge_results;
|
||||
results.extend(denied_results);
|
||||
current_messages.push(response_to_assistant_message(¤t_response));
|
||||
current_messages.push(tool_results_to_user_message(&results));
|
||||
let mut follow_up_req = original_req.clone();
|
||||
follow_up_req.messages = current_messages.clone();
|
||||
let llm_start = Instant::now();
|
||||
let traces: Vec<ToolCallTrace> = results
|
||||
.iter()
|
||||
.map(|r| ToolCallTrace {
|
||||
tool_name: r.tool_name.clone(),
|
||||
duration: Duration::ZERO,
|
||||
outcome: r.outcome.clone(),
|
||||
})
|
||||
.collect();
|
||||
iterations.push(IterationTrace {
|
||||
tool_calls: traces,
|
||||
llm_latency: Duration::ZERO,
|
||||
});
|
||||
match backend_call(follow_up_req).await {
|
||||
Ok(resp) => {
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
prev_tool_calls = None;
|
||||
current_response = resp;
|
||||
continue;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "follow-up backend call failed after nudge/deny");
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Guard: duplicate detection (same tool calls as previous iteration)
|
||||
let auto_calls: Vec<ToolCall> = auto_exec.iter().map(|c| (*c).clone()).collect();
|
||||
if let Some(ref prev) = prev_tool_calls {
|
||||
if is_duplicate(prev, &auto_calls) {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::DuplicateDetected,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Execute auto-allowed tools in parallel
|
||||
let exec_start = Instant::now();
|
||||
let mut results = execute_tool_calls(
|
||||
&auto_exec,
|
||||
engine.registry.clone(),
|
||||
&engine.policy,
|
||||
&engine.loop_config,
|
||||
)
|
||||
.await;
|
||||
let exec_duration = exec_start.elapsed();
|
||||
|
||||
// Append nudge + denied-tool error results so the LLM sees all outcomes.
|
||||
results.extend(nudge_results);
|
||||
results.extend(denied_results);
|
||||
|
||||
// Build per-tool traces
|
||||
let tool_traces: Vec<ToolCallTrace> = results
|
||||
.iter()
|
||||
.map(|r| ToolCallTrace {
|
||||
tool_name: r.tool_name.clone(),
|
||||
duration: exec_duration,
|
||||
outcome: r.outcome.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Guard: all tools failed (includes deny errors, which are non-retryable)
|
||||
let all_failed = results
|
||||
.iter()
|
||||
.all(|r| !matches!(r.outcome, ToolOutcome::Success(_)));
|
||||
|
||||
iterations.push(IterationTrace {
|
||||
tool_calls: tool_traces,
|
||||
llm_latency: Duration::ZERO, // filled after backend call below
|
||||
});
|
||||
|
||||
if all_failed {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::AllToolsFailed,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Build follow-up: append assistant response + tool results to conversation
|
||||
current_messages.push(response_to_assistant_message(¤t_response));
|
||||
current_messages.push(tool_results_to_user_message(&results));
|
||||
|
||||
let mut follow_up_req = original_req.clone();
|
||||
follow_up_req.messages = current_messages.clone();
|
||||
|
||||
// Call backend via caller-provided closure
|
||||
let llm_start = Instant::now();
|
||||
match backend_call(follow_up_req).await {
|
||||
Ok(resp) => {
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
tracing::info!(
|
||||
tools_executed = results.len(),
|
||||
iteration = _iteration + 1,
|
||||
"tool execution loop: iteration complete"
|
||||
);
|
||||
prev_tool_calls = Some(auto_calls);
|
||||
current_response = resp;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"follow-up backend call failed, returning last good response"
|
||||
);
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
// Backend error is not a clean termination; closest match is NoToolCalls
|
||||
// since we are stopping the loop and returning what we have.
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Exhausted max_iterations
|
||||
(
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::MaxIterations,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,373 @@
|
||||
//! Core async tool-execution runners and the centralized bounded tool loop
|
||||
//! for non-streaming requests.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::Value;
|
||||
use tokio::task::JoinSet;
|
||||
|
||||
use super::{
|
||||
extract_tool_calls, partition_and_nudge, response_to_assistant_message,
|
||||
tool_results_to_user_message, LoopConfig, ToolCall, ToolEngineState, ToolResult,
|
||||
};
|
||||
use crate::tools::policy::ToolExecutionPolicy;
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
use crate::tools::trace::{
|
||||
IterationTrace, LoopTrace, TerminationReason, ToolCallTrace, ToolOutcome,
|
||||
};
|
||||
use crate::tools::ToolGuardrailRequestState;
|
||||
|
||||
/// Execute tool calls in parallel, respecting per-tool timeouts.
|
||||
///
|
||||
/// Results are returned in the same order as `calls`.
|
||||
pub async fn execute_tool_calls(
|
||||
calls: &[&ToolCall],
|
||||
registry: Arc<ToolRegistry>,
|
||||
policy: &ToolExecutionPolicy,
|
||||
config: &LoopConfig,
|
||||
) -> Vec<ToolResult> {
|
||||
let capped = &calls[..calls.len().min(config.max_tool_calls_per_turn)];
|
||||
|
||||
// Collect (original_index, ToolCall) to restore order after parallel execution.
|
||||
let indexed: Vec<(usize, &ToolCall)> = capped.iter().copied().enumerate().collect();
|
||||
|
||||
let mut join_set: JoinSet<(usize, ToolResult)> = JoinSet::new();
|
||||
|
||||
for (idx, call) in indexed {
|
||||
let timeout = policy
|
||||
.find_rule(&call.name)
|
||||
.and_then(|r| r.timeout)
|
||||
.unwrap_or(config.tool_timeout);
|
||||
|
||||
let registry = Arc::clone(®istry);
|
||||
let id = call.id.clone();
|
||||
let name = call.name.clone();
|
||||
let input = call.input.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
let result =
|
||||
tokio::time::timeout(timeout, execute_single(®istry, &name, input)).await;
|
||||
|
||||
let outcome = match result {
|
||||
Ok(Ok(value)) => ToolOutcome::Success(value),
|
||||
Ok(Err(msg)) => ToolOutcome::Error {
|
||||
message: msg,
|
||||
retryable: false,
|
||||
},
|
||||
Err(_elapsed) => ToolOutcome::Timeout,
|
||||
};
|
||||
|
||||
(
|
||||
idx,
|
||||
ToolResult {
|
||||
tool_use_id: id,
|
||||
tool_name: name,
|
||||
outcome,
|
||||
},
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
let mut collected: Vec<(usize, ToolResult)> = Vec::with_capacity(capped.len());
|
||||
while let Some(res) = join_set.join_next().await {
|
||||
match res {
|
||||
Ok(pair) => collected.push(pair),
|
||||
Err(e) => {
|
||||
// JoinError means the task panicked; treat as an error outcome.
|
||||
// We don't have the index here, so we skip (shouldn't happen in practice).
|
||||
tracing::error!("tool execution task panicked: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
collected.sort_by_key(|(idx, _)| *idx);
|
||||
collected.into_iter().map(|(_, r)| r).collect()
|
||||
}
|
||||
|
||||
/// Check whether two slices of ToolCall represent the same logical calls.
|
||||
///
|
||||
/// Same length, same multiset of (name, input) pairs. IDs are ignored.
|
||||
pub fn is_duplicate(a: &[ToolCall], b: &[ToolCall]) -> bool {
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut a_pairs: Vec<(&str, &Value)> = a.iter().map(|c| (c.name.as_str(), &c.input)).collect();
|
||||
let mut b_pairs: Vec<(&str, &Value)> = b.iter().map(|c| (c.name.as_str(), &c.input)).collect();
|
||||
|
||||
// Sort by name so comparison is order-independent.
|
||||
a_pairs.sort_by_key(|(name, _)| *name);
|
||||
b_pairs.sort_by_key(|(name, _)| *name);
|
||||
|
||||
a_pairs == b_pairs
|
||||
}
|
||||
|
||||
/// Execute a single tool by name, looking it up in the registry.
|
||||
async fn execute_single(
|
||||
registry: &ToolRegistry,
|
||||
tool_name: &str,
|
||||
input: Value,
|
||||
) -> Result<Value, String> {
|
||||
match registry.get(tool_name) {
|
||||
Some(tool) => tool.execute(input).await,
|
||||
None => Err(format!("tool '{}' not found in registry", tool_name)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Run `execute_tool_calls` and record wall-clock duration per call.
|
||||
/// Returns (results, elapsed_per_call). Exposed for loop-level tracing.
|
||||
pub async fn execute_tool_calls_timed(
|
||||
calls: &[&ToolCall],
|
||||
registry: Arc<ToolRegistry>,
|
||||
policy: &ToolExecutionPolicy,
|
||||
config: &LoopConfig,
|
||||
) -> (Vec<ToolResult>, Duration) {
|
||||
let start = Instant::now();
|
||||
let results = execute_tool_calls(calls, registry, policy, config).await;
|
||||
(results, start.elapsed())
|
||||
}
|
||||
|
||||
/// Process an LLM response for tool execution. If auto-executable tool calls
|
||||
/// are found, executes them and makes follow-up backend calls in a bounded loop.
|
||||
///
|
||||
/// Returns the final response (original if no tools were auto-executed) and a
|
||||
/// `LoopTrace` recording what happened.
|
||||
///
|
||||
/// `backend_call` is a closure the caller provides. It takes a
|
||||
/// `MessageCreateRequest` and returns the translated `MessageResponse`.
|
||||
/// This keeps the loop backend-agnostic: the handler knows how to translate
|
||||
/// and call its specific backend; this function only knows about Anthropic types.
|
||||
///
|
||||
/// `guardrails` is the effective guardrail config for this request (the
|
||||
/// runtime-tunable override applied on top of `engine.guardrails`, the
|
||||
/// static per-process preset -- see `crate::tools::resolve_runtime_guardrails`
|
||||
/// and `AppState::effective_tool_guardrails`). Callers that don't need the
|
||||
/// runtime override can pass `&engine.guardrails` directly.
|
||||
pub async fn maybe_execute_tools<F, Fut>(
|
||||
engine: &ToolEngineState,
|
||||
original_req: &anyllm_translate::anthropic::MessageCreateRequest,
|
||||
server_advertised_tool_names: &HashSet<String>,
|
||||
initial_response: anyllm_translate::anthropic::MessageResponse,
|
||||
guardrails: &crate::tools::ToolGuardrailConfig,
|
||||
backend_call: F,
|
||||
) -> (anyllm_translate::anthropic::MessageResponse, LoopTrace)
|
||||
where
|
||||
F: Fn(anyllm_translate::anthropic::MessageCreateRequest) -> Fut,
|
||||
Fut: std::future::Future<Output = Result<anyllm_translate::anthropic::MessageResponse, String>>,
|
||||
{
|
||||
let loop_start = Instant::now();
|
||||
let mut iterations: Vec<IterationTrace> = Vec::new();
|
||||
let mut current_response = initial_response;
|
||||
let mut current_messages = original_req.messages.clone();
|
||||
let mut prev_tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut guardrail_state = ToolGuardrailRequestState::new();
|
||||
|
||||
for _iteration in 0..engine.loop_config.max_iterations {
|
||||
// Guard: total timeout
|
||||
if loop_start.elapsed() > engine.loop_config.total_timeout {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::Timeout,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let tool_calls = extract_tool_calls(¤t_response);
|
||||
|
||||
// Advisory guardrails: each nudge targets one offending call and comes
|
||||
// back as a retryable error result. Nudged calls are skipped this turn;
|
||||
// every other call is partitioned and executed as usual. Each distinct
|
||||
// decision nudges only once per request (across iterations), so a
|
||||
// model that ignores a nudge and repeats the call lets it proceed
|
||||
// instead of spinning the loop to max_iterations. Nudges only ever
|
||||
// apply to calls this proxy would actually execute -- see
|
||||
// `partition_and_nudge`'s doc comment.
|
||||
let (auto_exec, nudge_results, denied_results) = partition_and_nudge(
|
||||
&tool_calls,
|
||||
original_req.tools.as_deref().unwrap_or(&[]),
|
||||
&engine.registry,
|
||||
&engine.policy,
|
||||
server_advertised_tool_names,
|
||||
guardrails,
|
||||
&mut guardrail_state,
|
||||
);
|
||||
|
||||
if auto_exec.is_empty() && denied_results.is_empty() && nudge_results.is_empty() {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// If there is nothing to execute (only nudges and/or denials), send the
|
||||
// advisory/error results back to the LLM immediately without running a tool.
|
||||
if auto_exec.is_empty() {
|
||||
let mut results = nudge_results;
|
||||
results.extend(denied_results);
|
||||
current_messages.push(response_to_assistant_message(¤t_response));
|
||||
current_messages.push(tool_results_to_user_message(&results));
|
||||
let mut follow_up_req = original_req.clone();
|
||||
follow_up_req.messages = current_messages.clone();
|
||||
let llm_start = Instant::now();
|
||||
let traces: Vec<ToolCallTrace> = results
|
||||
.iter()
|
||||
.map(|r| ToolCallTrace {
|
||||
tool_name: r.tool_name.clone(),
|
||||
duration: Duration::ZERO,
|
||||
outcome: r.outcome.clone(),
|
||||
})
|
||||
.collect();
|
||||
iterations.push(IterationTrace {
|
||||
tool_calls: traces,
|
||||
llm_latency: Duration::ZERO,
|
||||
});
|
||||
match backend_call(follow_up_req).await {
|
||||
Ok(resp) => {
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
prev_tool_calls = None;
|
||||
current_response = resp;
|
||||
continue;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "follow-up backend call failed after nudge/deny");
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Guard: duplicate detection (same tool calls as previous iteration)
|
||||
let auto_calls: Vec<ToolCall> = auto_exec.iter().map(|c| (*c).clone()).collect();
|
||||
if let Some(ref prev) = prev_tool_calls {
|
||||
if is_duplicate(prev, &auto_calls) {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::DuplicateDetected,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Execute auto-allowed tools in parallel
|
||||
let exec_start = Instant::now();
|
||||
let mut results = execute_tool_calls(
|
||||
&auto_exec,
|
||||
engine.registry.clone(),
|
||||
&engine.policy,
|
||||
&engine.loop_config,
|
||||
)
|
||||
.await;
|
||||
let exec_duration = exec_start.elapsed();
|
||||
|
||||
// Append nudge + denied-tool error results so the LLM sees all outcomes.
|
||||
results.extend(nudge_results);
|
||||
results.extend(denied_results);
|
||||
|
||||
// Build per-tool traces
|
||||
let tool_traces: Vec<ToolCallTrace> = results
|
||||
.iter()
|
||||
.map(|r| ToolCallTrace {
|
||||
tool_name: r.tool_name.clone(),
|
||||
duration: exec_duration,
|
||||
outcome: r.outcome.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Guard: all tools failed (includes deny errors, which are non-retryable)
|
||||
let all_failed = results
|
||||
.iter()
|
||||
.all(|r| !matches!(r.outcome, ToolOutcome::Success(_)));
|
||||
|
||||
iterations.push(IterationTrace {
|
||||
tool_calls: tool_traces,
|
||||
llm_latency: Duration::ZERO, // filled after backend call below
|
||||
});
|
||||
|
||||
if all_failed {
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::AllToolsFailed,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Build follow-up: append assistant response + tool results to conversation
|
||||
current_messages.push(response_to_assistant_message(¤t_response));
|
||||
current_messages.push(tool_results_to_user_message(&results));
|
||||
|
||||
let mut follow_up_req = original_req.clone();
|
||||
follow_up_req.messages = current_messages.clone();
|
||||
|
||||
// Call backend via caller-provided closure
|
||||
let llm_start = Instant::now();
|
||||
match backend_call(follow_up_req).await {
|
||||
Ok(resp) => {
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
tracing::info!(
|
||||
tools_executed = results.len(),
|
||||
iteration = _iteration + 1,
|
||||
"tool execution loop: iteration complete"
|
||||
);
|
||||
prev_tool_calls = Some(auto_calls);
|
||||
current_response = resp;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"follow-up backend call failed, returning last good response"
|
||||
);
|
||||
if let Some(last) = iterations.last_mut() {
|
||||
last.llm_latency = llm_start.elapsed();
|
||||
}
|
||||
return (
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
// Backend error is not a clean termination; closest match is NoToolCalls
|
||||
// since we are stopping the loop and returning what we have.
|
||||
termination_reason: TerminationReason::NoToolCalls,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Exhausted max_iterations
|
||||
(
|
||||
current_response,
|
||||
LoopTrace {
|
||||
iterations,
|
||||
total_duration: loop_start.elapsed(),
|
||||
termination_reason: TerminationReason::MaxIterations,
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -2,6 +2,7 @@ use super::*;
|
||||
use crate::tools::policy::{PolicyAction, PolicyRule, ToolExecutionPolicy};
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
use serde_json::json;
|
||||
use std::collections::HashSet;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -6,9 +6,15 @@
|
||||
|
||||
use crate::tools::execution::ToolCall;
|
||||
use anyllm_translate::anthropic::Tool;
|
||||
use serde_json::{Map, Value};
|
||||
use std::str::FromStr;
|
||||
|
||||
mod lsp;
|
||||
mod quiet;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
mod utils;
|
||||
mod write_cap;
|
||||
|
||||
/// Default maximum string payload size for write/edit policy nudges.
|
||||
pub const DEFAULT_MAX_WRITE_PAYLOAD_BYTES: usize = 64 * 1024;
|
||||
|
||||
@@ -207,7 +213,7 @@ pub fn evaluate_tool_guardrails(
|
||||
}
|
||||
|
||||
let lsp_tools = if config.lsp_first {
|
||||
available_lsp_tools(tool_specs)
|
||||
lsp::available_lsp_tools(tool_specs)
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
@@ -216,18 +222,20 @@ pub fn evaluate_tool_guardrails(
|
||||
let mut nudged_fingerprints_this_batch: indexmap::IndexSet<String> = indexmap::IndexSet::new();
|
||||
for call in tool_calls {
|
||||
let candidate = (!lsp_tools.is_empty())
|
||||
.then(|| lsp_first_nudge(call, &lsp_tools))
|
||||
.then(|| lsp::lsp_first_nudge(call, &lsp_tools))
|
||||
.flatten()
|
||||
.or_else(|| {
|
||||
config
|
||||
.write_payload_caps
|
||||
.then(|| write_payload_cap_nudge(call, config.max_write_payload_bytes))
|
||||
.then(|| {
|
||||
write_cap::write_payload_cap_nudge(call, config.max_write_payload_bytes)
|
||||
})
|
||||
.flatten()
|
||||
})
|
||||
.or_else(|| {
|
||||
config
|
||||
.quiet_commands
|
||||
.then(|| quiet_command_nudge(call))
|
||||
.then(|| quiet::quiet_command_nudge(call))
|
||||
.flatten()
|
||||
});
|
||||
|
||||
@@ -246,515 +254,3 @@ pub fn evaluate_tool_guardrails(
|
||||
|
||||
nudges
|
||||
}
|
||||
|
||||
fn available_lsp_tools(tool_specs: &[Tool]) -> Vec<String> {
|
||||
let supported = [
|
||||
"find_definition",
|
||||
"find_references",
|
||||
"get_hover",
|
||||
"document_symbols",
|
||||
"workspace_symbols",
|
||||
];
|
||||
tool_specs
|
||||
.iter()
|
||||
.filter(|tool| supported.contains(&tool.name.as_str()))
|
||||
.map(|tool| tool.name.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn lsp_first_nudge(call: &ToolCall, lsp_tools: &[String]) -> Option<ToolGuardrailNudge> {
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
let args = object_args(&call.input);
|
||||
let symbol = if is_shell_tool(&tool_name) {
|
||||
shell_grep_symbol(command_arg(args?)?)?
|
||||
} else if is_grep_tool(&tool_name) {
|
||||
string_arg(
|
||||
args?,
|
||||
&["symbol", "name", "pattern", "query", "regex", "needle"],
|
||||
)
|
||||
.and_then(symbol_from_search_value)?
|
||||
} else if is_glob_tool(&tool_name) {
|
||||
string_arg(args?, &["pattern", "query", "glob"]).and_then(symbol_from_search_value)?
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let tools = lsp_tools.join(", ");
|
||||
let fingerprint = format!("lsp_first:{}:{symbol}", call.name);
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "lsp_first",
|
||||
content: format!(
|
||||
"Use available LSP tools for symbol lookup instead of grep/glob/shell search. Available LSP tools: {tools}. Retry with the best matching LSP tool for `{symbol}`."
|
||||
),
|
||||
fingerprint,
|
||||
})
|
||||
}
|
||||
|
||||
fn quiet_command_nudge(call: &ToolCall) -> Option<ToolGuardrailNudge> {
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
if !is_shell_tool(&tool_name) {
|
||||
return None;
|
||||
}
|
||||
let command = command_arg(object_args(&call.input)?)?.trim();
|
||||
let suggestion = quiet_command_suggestion(command)?;
|
||||
let fingerprint = format!("quiet:{}:{command}:{suggestion}", call.name);
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "quiet_command",
|
||||
content: format!(
|
||||
"The requested shell command is likely to produce noisy output. Prefer `{suggestion}`. Repeat the original command only if verbose output is required."
|
||||
),
|
||||
fingerprint,
|
||||
})
|
||||
}
|
||||
|
||||
fn write_payload_cap_nudge(call: &ToolCall, max_bytes: usize) -> Option<ToolGuardrailNudge> {
|
||||
// A cap of 0 disables the nudge rather than firing on every non-empty write.
|
||||
if max_bytes == 0 {
|
||||
return None;
|
||||
}
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
if !is_write_or_edit_tool(&tool_name) {
|
||||
return None;
|
||||
}
|
||||
let bytes = object_args(&call.input)
|
||||
.map(write_payload_bytes)
|
||||
.unwrap_or(0);
|
||||
if bytes <= max_bytes {
|
||||
return None;
|
||||
}
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "write_payload_cap",
|
||||
content: format!(
|
||||
"The requested write/edit payload is too large for this proxy policy ({bytes} bytes > {max_bytes} bytes). Retry with a smaller targeted edit or split the change."
|
||||
),
|
||||
// Include a hash of the actual payload, not just its byte length, so
|
||||
// two unrelated oversized writes that happen to be the same size
|
||||
// (e.g. same-size templated files) don't collide and suppress each
|
||||
// other's nudge. Identical repeated calls (same target, same
|
||||
// content) still dedupe as intended.
|
||||
fingerprint: format!(
|
||||
"write_payload_cap:{}:{:x}:{bytes}:{max_bytes}",
|
||||
call.name,
|
||||
payload_content_hash(&call.input)
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn payload_content_hash(value: &Value) -> u64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
value.to_string().hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
fn is_shell_tool(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
"bash" | "shell" | "run_command" | "execute_command" | "terminal" | "exec" | "execute_bash"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_grep_tool(name: &str) -> bool {
|
||||
matches!(name, "grep" | "rg" | "ripgrep" | "st")
|
||||
}
|
||||
|
||||
fn is_glob_tool(name: &str) -> bool {
|
||||
matches!(name, "glob" | "find_files" | "file_glob")
|
||||
}
|
||||
|
||||
fn is_write_or_edit_tool(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
"write"
|
||||
| "write_file"
|
||||
| "edit"
|
||||
| "edit_file"
|
||||
| "replace"
|
||||
| "apply_patch"
|
||||
| "create_file"
|
||||
| "update_file"
|
||||
)
|
||||
}
|
||||
|
||||
fn object_args(value: &Value) -> Option<&Map<String, Value>> {
|
||||
value.as_object()
|
||||
}
|
||||
|
||||
fn command_arg(args: &Map<String, Value>) -> Option<&str> {
|
||||
string_arg(args, &["command", "cmd", "shell_command", "input"])
|
||||
}
|
||||
|
||||
fn string_arg<'a>(args: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
|
||||
// Only the named keys, in order. Do NOT fall back to an arbitrary string
|
||||
// value: that would let an unrelated field be misread as the search symbol
|
||||
// or command for an advisory nudge.
|
||||
keys.iter()
|
||||
.find_map(|key| args.get(*key).and_then(Value::as_str))
|
||||
}
|
||||
|
||||
fn shell_grep_symbol(command: &str) -> Option<String> {
|
||||
let mut parts = command.split_whitespace();
|
||||
let binary = strip_shell_quotes(parts.next()?)
|
||||
.rsplit('/')
|
||||
.next()?
|
||||
.to_string();
|
||||
if !matches!(binary.as_str(), "rg" | "ripgrep" | "grep" | "st") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut previous_took_value = false;
|
||||
for part in parts {
|
||||
let token = strip_shell_quotes(part);
|
||||
if previous_took_value {
|
||||
previous_took_value = false;
|
||||
continue;
|
||||
}
|
||||
if token.starts_with("--") {
|
||||
previous_took_value = matches!(
|
||||
token.as_str(),
|
||||
"--glob" | "--type" | "--context" | "--after-context" | "--before-context"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if token.starts_with('-') {
|
||||
continue;
|
||||
}
|
||||
return symbol_from_search_value(&token);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn symbol_from_search_value(value: &str) -> Option<String> {
|
||||
let trimmed = strip_shell_quotes(value)
|
||||
.trim_matches('/')
|
||||
.replace("\\b", "")
|
||||
.replace(['^', '$'], "");
|
||||
for token in trimmed.split(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '_')) {
|
||||
if looks_like_symbol_token(token) {
|
||||
return Some(token.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn looks_like_symbol_token(token: &str) -> bool {
|
||||
if token.len() < 3 {
|
||||
return false;
|
||||
}
|
||||
let mut chars = token.chars();
|
||||
let Some(first) = chars.next() else {
|
||||
return false;
|
||||
};
|
||||
if !(first.is_ascii_alphabetic() || first == '_') {
|
||||
return false;
|
||||
}
|
||||
if !chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_') {
|
||||
return false;
|
||||
}
|
||||
|
||||
let lower = token.to_ascii_lowercase();
|
||||
if matches!(
|
||||
lower.as_str(),
|
||||
"todo" | "fixme" | "error" | "warning" | "debug" | "test" | "src" | "main"
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
token.contains('_') || token.chars().any(|ch| ch.is_ascii_uppercase())
|
||||
}
|
||||
|
||||
fn strip_shell_quotes(value: &str) -> String {
|
||||
value
|
||||
.trim()
|
||||
.trim_matches('"')
|
||||
.trim_matches('\'')
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn quiet_command_suggestion(command: &str) -> Option<String> {
|
||||
if command_has_prefix(command, "git log") && !contains_word(command, "--oneline") {
|
||||
return Some(insert_after_prefix(command, "git log", "--oneline"));
|
||||
}
|
||||
for prefix in ["cargo build", "cargo check", "cargo clippy", "cargo test"] {
|
||||
if command_has_prefix(command, prefix) && !contains_word(command, "--quiet") {
|
||||
return Some(insert_after_prefix(command, prefix, "--quiet"));
|
||||
}
|
||||
}
|
||||
if command_has_prefix(command, "pytest") && !contains_word(command, "-q") {
|
||||
return Some(insert_after_prefix(command, "pytest", "-q"));
|
||||
}
|
||||
if command_has_prefix(command, "npm install") && !contains_word(command, "--silent") {
|
||||
return Some(insert_after_prefix(command, "npm install", "--silent"));
|
||||
}
|
||||
if command_has_prefix(command, "pip install") && !contains_word(command, "--quiet") {
|
||||
return Some(insert_after_prefix(command, "pip install", "--quiet"));
|
||||
}
|
||||
if command_has_prefix(command, "docker build") && !contains_word(command, "--progress=quiet") {
|
||||
return Some(insert_after_prefix(
|
||||
command,
|
||||
"docker build",
|
||||
"--progress=quiet",
|
||||
));
|
||||
}
|
||||
if command_has_prefix(command, "curl") && !contains_word(command, "-s") {
|
||||
return Some(insert_after_prefix(command, "curl", "-s"));
|
||||
}
|
||||
if command_has_prefix(command, "make") && !contains_word(command, "-s") {
|
||||
return Some(insert_after_prefix(command, "make", "-s"));
|
||||
}
|
||||
if command_has_prefix(command, "tree") && !contains_word(command, "-I") {
|
||||
return Some(insert_after_prefix(
|
||||
command,
|
||||
"tree",
|
||||
"-I \"node_modules|.git|target|dist|build\"",
|
||||
));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn command_has_prefix(command: &str, prefix: &str) -> bool {
|
||||
command == prefix || command.starts_with(&format!("{prefix} "))
|
||||
}
|
||||
|
||||
fn contains_word(command: &str, word: &str) -> bool {
|
||||
command.split_whitespace().any(|part| part == word)
|
||||
}
|
||||
|
||||
fn insert_after_prefix(command: &str, prefix: &str, insertion: &str) -> String {
|
||||
let rest = command[prefix.len()..].trim_start();
|
||||
if rest.is_empty() {
|
||||
format!("{prefix} {insertion}")
|
||||
} else {
|
||||
format!("{prefix} {insertion} {rest}")
|
||||
}
|
||||
}
|
||||
|
||||
/// Depth cap for `value_payload_bytes`'s recursion. Model-controlled JSON has
|
||||
/// no natural nesting limit; a real write/edit payload never needs anywhere
|
||||
/// near this depth, so hitting it is itself a signal the payload is oversized.
|
||||
const MAX_PAYLOAD_DEPTH: usize = 32;
|
||||
|
||||
fn write_payload_bytes(args: &Map<String, Value>) -> usize {
|
||||
let payload_keys = [
|
||||
"content",
|
||||
"text",
|
||||
"new_content",
|
||||
"patch",
|
||||
"diff",
|
||||
"replacement",
|
||||
"data",
|
||||
"old_string",
|
||||
"new_string",
|
||||
];
|
||||
let known = args
|
||||
.iter()
|
||||
.filter(|(key, _)| payload_keys.contains(&key.as_str()))
|
||||
.fold(0usize, |acc, (_, value)| {
|
||||
acc.saturating_add(value_payload_bytes(value, 0))
|
||||
});
|
||||
if known > 0 {
|
||||
return known;
|
||||
}
|
||||
// None of the known field names matched (a tool with a non-standard
|
||||
// schema, e.g. `body`/`value`) -- fall back to summing every value in
|
||||
// the payload rather than silently reporting zero and never capping it.
|
||||
args.values().fold(0usize, |acc, value| {
|
||||
acc.saturating_add(value_payload_bytes(value, 0))
|
||||
})
|
||||
}
|
||||
|
||||
fn value_payload_bytes(value: &Value, depth: usize) -> usize {
|
||||
if depth >= MAX_PAYLOAD_DEPTH {
|
||||
// Treat implausibly deep nesting as oversized rather than recursing
|
||||
// further; a fixed large sentinel (not usize::MAX) keeps the
|
||||
// saturating sum below safe from overflowing back to a small value.
|
||||
return 1_000_000_000;
|
||||
}
|
||||
match value {
|
||||
Value::String(value) => value.len(),
|
||||
Value::Array(values) => values.iter().fold(0usize, |acc, v| {
|
||||
acc.saturating_add(value_payload_bytes(v, depth + 1))
|
||||
}),
|
||||
Value::Object(values) => values.values().fold(0usize, |acc, v| {
|
||||
acc.saturating_add(value_payload_bytes(v, depth + 1))
|
||||
}),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn tool_spec(name: &str) -> Tool {
|
||||
Tool {
|
||||
name: name.to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type": "object"}),
|
||||
}
|
||||
}
|
||||
|
||||
fn call(name: &str, input: Value) -> ToolCall {
|
||||
ToolCall {
|
||||
id: "toolu_1".to_string(),
|
||||
name: name.to_string(),
|
||||
input,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lsp_nudge_requires_replacement_tool() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
lsp_first: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("grep", json!({"pattern": "UserService"}))];
|
||||
|
||||
assert!(
|
||||
evaluate_tool_guardrails(&calls, &[tool_spec("grep")], &config, &mut state).is_empty()
|
||||
);
|
||||
|
||||
let nudges = evaluate_tool_guardrails(
|
||||
&calls,
|
||||
&[tool_spec("grep"), tool_spec("find_definition")],
|
||||
&config,
|
||||
&mut state,
|
||||
);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
let nudge = &nudges[0];
|
||||
assert_eq!(nudge.kind, "lsp_first");
|
||||
assert_eq!(nudge.call_id, "toolu_1");
|
||||
assert!(nudge.content.contains("find_definition"));
|
||||
assert!(nudge.content.contains("UserService"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quiet_command_nudges_once() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
quiet_commands: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("bash", json!({"command": "cargo test"}))];
|
||||
let first = evaluate_tool_guardrails(&calls, &[], &config, &mut state);
|
||||
assert_eq!(first.len(), 1);
|
||||
assert_eq!(first[0].kind, "quiet_command");
|
||||
assert!(first[0].content.contains("cargo test --quiet"));
|
||||
assert!(evaluate_tool_guardrails(&calls, &[], &config, &mut state).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_offending_call_is_nudged() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
lsp_first: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![
|
||||
ToolCall {
|
||||
id: "toolu_grep".into(),
|
||||
name: "grep".into(),
|
||||
input: json!({"pattern": "UserService"}),
|
||||
},
|
||||
ToolCall {
|
||||
id: "toolu_write".into(),
|
||||
name: "write_file".into(),
|
||||
input: json!({"content": "ok"}),
|
||||
},
|
||||
];
|
||||
let nudges = evaluate_tool_guardrails(
|
||||
&calls,
|
||||
&[tool_spec("grep"), tool_spec("find_definition")],
|
||||
&config,
|
||||
&mut state,
|
||||
);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
assert_eq!(nudges[0].call_id, "toolu_grep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_detects_oversized_payload() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 4,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("write_file", json!({"content": "12345"}))];
|
||||
let nudges = evaluate_tool_guardrails(&calls, &[], &config, &mut state);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
assert_eq!(nudges[0].kind, "write_payload_cap");
|
||||
assert!(nudges[0].content.contains("5 bytes > 4 bytes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_zero_disables_nudge() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 0,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("write_file", json!({"content": "12345"}))];
|
||||
assert!(evaluate_tool_guardrails(&calls, &[], &config, &mut state).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_does_not_collide_on_equal_length_different_content() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 4,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
// Same byte length (5), different content/target -- both must nudge.
|
||||
let first = vec![call("write_file", json!({"content": "aaaaa"}))];
|
||||
let second = vec![call("write_file", json!({"content": "bbbbb"}))];
|
||||
assert_eq!(
|
||||
evaluate_tool_guardrails(&first, &[], &config, &mut state).len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
evaluate_tool_guardrails(&second, &[], &config, &mut state).len(),
|
||||
1,
|
||||
"an unrelated oversized write of equal byte length must still be nudged"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_bytes_falls_back_to_unrecognized_field_names() {
|
||||
let args = json!({"body": "12345"});
|
||||
assert_eq!(write_payload_bytes(args.as_object().unwrap()), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_prefers_runtime_override() {
|
||||
let static_config = ToolGuardrailConfig::disabled();
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "standard");
|
||||
assert_eq!(resolved, ToolGuardrailConfig::standard());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_keeps_static_when_modes_match() {
|
||||
let static_config = ToolGuardrailConfig {
|
||||
max_write_payload_bytes: 123,
|
||||
..ToolGuardrailConfig::standard()
|
||||
};
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "standard");
|
||||
// Same mode as the static preset -- the static config (with its
|
||||
// custom max_write_payload_bytes) must be preserved, not rebuilt
|
||||
// from the bare preset.
|
||||
assert_eq!(resolved, static_config);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_falls_back_on_unparseable_value() {
|
||||
let static_config = ToolGuardrailConfig::standard();
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "not-a-real-mode");
|
||||
assert_eq!(resolved, static_config);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
use super::utils::*;
|
||||
use super::ToolGuardrailNudge;
|
||||
use crate::tools::execution::ToolCall;
|
||||
use anyllm_translate::anthropic::Tool;
|
||||
|
||||
pub(super) fn available_lsp_tools(tool_specs: &[Tool]) -> Vec<String> {
|
||||
let supported = [
|
||||
"find_definition",
|
||||
"find_references",
|
||||
"get_hover",
|
||||
"document_symbols",
|
||||
"workspace_symbols",
|
||||
];
|
||||
tool_specs
|
||||
.iter()
|
||||
.filter(|tool| supported.contains(&tool.name.as_str()))
|
||||
.map(|tool| tool.name.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn lsp_first_nudge(call: &ToolCall, lsp_tools: &[String]) -> Option<ToolGuardrailNudge> {
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
let args = object_args(&call.input);
|
||||
let symbol = if is_shell_tool(&tool_name) {
|
||||
shell_grep_symbol(command_arg(args?)?)?
|
||||
} else if is_grep_tool(&tool_name) {
|
||||
string_arg(
|
||||
args?,
|
||||
&["symbol", "name", "pattern", "query", "regex", "needle"],
|
||||
)
|
||||
.and_then(symbol_from_search_value)?
|
||||
} else if is_glob_tool(&tool_name) {
|
||||
string_arg(args?, &["pattern", "query", "glob"]).and_then(symbol_from_search_value)?
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let tools = lsp_tools.join(", ");
|
||||
let fingerprint = format!("lsp_first:{}:{symbol}", call.name);
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "lsp_first",
|
||||
content: format!(
|
||||
"Use available LSP tools for symbol lookup instead of grep/glob/shell search. Available LSP tools: {tools}. Retry with the best matching LSP tool for `{symbol}`."
|
||||
),
|
||||
fingerprint,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
use super::utils::*;
|
||||
use super::ToolGuardrailNudge;
|
||||
use crate::tools::execution::ToolCall;
|
||||
|
||||
pub(super) fn quiet_command_nudge(call: &ToolCall) -> Option<ToolGuardrailNudge> {
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
if !is_shell_tool(&tool_name) {
|
||||
return None;
|
||||
}
|
||||
let command = command_arg(object_args(&call.input)?)?.trim();
|
||||
let suggestion = quiet_command_suggestion(command)?;
|
||||
let fingerprint = format!("quiet:{}:{command}:{suggestion}", call.name);
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "quiet_command",
|
||||
content: format!(
|
||||
"The requested shell command is likely to produce noisy output. Prefer `{suggestion}`. Repeat the original command only if verbose output is required."
|
||||
),
|
||||
fingerprint,
|
||||
})
|
||||
}
|
||||
|
||||
fn quiet_command_suggestion(command: &str) -> Option<String> {
|
||||
if command_has_prefix(command, "git log") && !contains_word(command, "--oneline") {
|
||||
return Some(insert_after_prefix(command, "git log", "--oneline"));
|
||||
}
|
||||
for prefix in ["cargo build", "cargo check", "cargo clippy", "cargo test"] {
|
||||
if command_has_prefix(command, prefix) && !contains_word(command, "--quiet") {
|
||||
return Some(insert_after_prefix(command, prefix, "--quiet"));
|
||||
}
|
||||
}
|
||||
if command_has_prefix(command, "pytest") && !contains_word(command, "-q") {
|
||||
return Some(insert_after_prefix(command, "pytest", "-q"));
|
||||
}
|
||||
if command_has_prefix(command, "npm install") && !contains_word(command, "--silent") {
|
||||
return Some(insert_after_prefix(command, "npm install", "--silent"));
|
||||
}
|
||||
if command_has_prefix(command, "pip install") && !contains_word(command, "--quiet") {
|
||||
return Some(insert_after_prefix(command, "pip install", "--quiet"));
|
||||
}
|
||||
if command_has_prefix(command, "docker build") && !contains_word(command, "--progress=quiet") {
|
||||
return Some(insert_after_prefix(
|
||||
command,
|
||||
"docker build",
|
||||
"--progress=quiet",
|
||||
));
|
||||
}
|
||||
if command_has_prefix(command, "curl") && !contains_word(command, "-s") {
|
||||
return Some(insert_after_prefix(command, "curl", "-s"));
|
||||
}
|
||||
if command_has_prefix(command, "make") && !contains_word(command, "-s") {
|
||||
return Some(insert_after_prefix(command, "make", "-s"));
|
||||
}
|
||||
if command_has_prefix(command, "tree") && !contains_word(command, "-I") {
|
||||
return Some(insert_after_prefix(
|
||||
command,
|
||||
"tree",
|
||||
"-I \"node_modules|.git|target|dist|build\"",
|
||||
));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn command_has_prefix(command: &str, prefix: &str) -> bool {
|
||||
command == prefix || command.starts_with(&format!("{prefix} "))
|
||||
}
|
||||
|
||||
fn contains_word(command: &str, word: &str) -> bool {
|
||||
command.split_whitespace().any(|part| part == word)
|
||||
}
|
||||
|
||||
fn insert_after_prefix(command: &str, prefix: &str, insertion: &str) -> String {
|
||||
let rest = command[prefix.len()..].trim_start();
|
||||
if rest.is_empty() {
|
||||
format!("{prefix} {insertion}")
|
||||
} else {
|
||||
format!("{prefix} {insertion} {rest}")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
use super::*;
|
||||
use crate::tools::execution::ToolCall;
|
||||
use anyllm_translate::anthropic::Tool;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn tool_spec(name: &str) -> Tool {
|
||||
Tool {
|
||||
name: name.to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type": "object"}),
|
||||
}
|
||||
}
|
||||
|
||||
fn call(name: &str, input: Value) -> ToolCall {
|
||||
ToolCall {
|
||||
id: "toolu_1".to_string(),
|
||||
name: name.to_string(),
|
||||
input,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lsp_nudge_requires_replacement_tool() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
lsp_first: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("grep", json!({"pattern": "UserService"}))];
|
||||
|
||||
assert!(evaluate_tool_guardrails(&calls, &[tool_spec("grep")], &config, &mut state).is_empty());
|
||||
|
||||
let nudges = evaluate_tool_guardrails(
|
||||
&calls,
|
||||
&[tool_spec("grep"), tool_spec("find_definition")],
|
||||
&config,
|
||||
&mut state,
|
||||
);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
let nudge = &nudges[0];
|
||||
assert_eq!(nudge.kind, "lsp_first");
|
||||
assert_eq!(nudge.call_id, "toolu_1");
|
||||
assert!(nudge.content.contains("find_definition"));
|
||||
assert!(nudge.content.contains("UserService"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quiet_command_nudges_once() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
quiet_commands: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("bash", json!({"command": "cargo test"}))];
|
||||
let first = evaluate_tool_guardrails(&calls, &[], &config, &mut state);
|
||||
assert_eq!(first.len(), 1);
|
||||
assert_eq!(first[0].kind, "quiet_command");
|
||||
assert!(first[0].content.contains("cargo test --quiet"));
|
||||
assert!(evaluate_tool_guardrails(&calls, &[], &config, &mut state).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_offending_call_is_nudged() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
lsp_first: true,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![
|
||||
ToolCall {
|
||||
id: "toolu_grep".into(),
|
||||
name: "grep".into(),
|
||||
input: json!({"pattern": "UserService"}),
|
||||
},
|
||||
ToolCall {
|
||||
id: "toolu_write".into(),
|
||||
name: "write_file".into(),
|
||||
input: json!({"content": "ok"}),
|
||||
},
|
||||
];
|
||||
let nudges = evaluate_tool_guardrails(
|
||||
&calls,
|
||||
&[tool_spec("grep"), tool_spec("find_definition")],
|
||||
&config,
|
||||
&mut state,
|
||||
);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
assert_eq!(nudges[0].call_id, "toolu_grep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_detects_oversized_payload() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 4,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("write_file", json!({"content": "12345"}))];
|
||||
let nudges = evaluate_tool_guardrails(&calls, &[], &config, &mut state);
|
||||
assert_eq!(nudges.len(), 1);
|
||||
assert_eq!(nudges[0].kind, "write_payload_cap");
|
||||
assert!(nudges[0].content.contains("5 bytes > 4 bytes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_zero_disables_nudge() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 0,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
let calls = vec![call("write_file", json!({"content": "12345"}))];
|
||||
assert!(evaluate_tool_guardrails(&calls, &[], &config, &mut state).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_cap_does_not_collide_on_equal_length_different_content() {
|
||||
let mut state = ToolGuardrailRequestState::new();
|
||||
let config = ToolGuardrailConfig {
|
||||
write_payload_caps: true,
|
||||
max_write_payload_bytes: 4,
|
||||
..ToolGuardrailConfig::disabled()
|
||||
};
|
||||
// Same byte length (5), different content/target -- both must nudge.
|
||||
let first = vec![call("write_file", json!({"content": "aaaaa"}))];
|
||||
let second = vec![call("write_file", json!({"content": "bbbbb"}))];
|
||||
assert_eq!(
|
||||
evaluate_tool_guardrails(&first, &[], &config, &mut state).len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
evaluate_tool_guardrails(&second, &[], &config, &mut state).len(),
|
||||
1,
|
||||
"an unrelated oversized write of equal byte length must still be nudged"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_payload_bytes_falls_back_to_unrecognized_field_names() {
|
||||
let args = json!({"body": "12345"});
|
||||
assert_eq!(write_cap::write_payload_bytes(args.as_object().unwrap()), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_prefers_runtime_override() {
|
||||
let static_config = ToolGuardrailConfig::disabled();
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "standard");
|
||||
assert_eq!(resolved, ToolGuardrailConfig::standard());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_keeps_static_when_modes_match() {
|
||||
let static_config = ToolGuardrailConfig {
|
||||
max_write_payload_bytes: 123,
|
||||
..ToolGuardrailConfig::standard()
|
||||
};
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "standard");
|
||||
// Same mode as the static preset -- the static config (with its
|
||||
// custom max_write_payload_bytes) must be preserved, not rebuilt
|
||||
// from the bare preset.
|
||||
assert_eq!(resolved, static_config);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_runtime_guardrails_falls_back_on_unparseable_value() {
|
||||
let static_config = ToolGuardrailConfig::standard();
|
||||
let resolved = resolve_runtime_guardrails(&static_config, "not-a-real-mode");
|
||||
assert_eq!(resolved, static_config);
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) fn is_shell_tool(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
"bash" | "shell" | "run_command" | "execute_command" | "terminal" | "exec" | "execute_bash"
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn is_grep_tool(name: &str) -> bool {
|
||||
matches!(name, "grep" | "rg" | "ripgrep" | "st")
|
||||
}
|
||||
|
||||
pub(super) fn is_glob_tool(name: &str) -> bool {
|
||||
matches!(name, "glob" | "find_files" | "file_glob")
|
||||
}
|
||||
|
||||
pub(super) fn is_write_or_edit_tool(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
"write"
|
||||
| "write_file"
|
||||
| "edit"
|
||||
| "edit_file"
|
||||
| "replace"
|
||||
| "apply_patch"
|
||||
| "create_file"
|
||||
| "update_file"
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn object_args(value: &Value) -> Option<&Map<String, Value>> {
|
||||
value.as_object()
|
||||
}
|
||||
|
||||
pub(super) fn command_arg(args: &Map<String, Value>) -> Option<&str> {
|
||||
string_arg(args, &["command", "cmd", "shell_command", "input"])
|
||||
}
|
||||
|
||||
pub(super) fn string_arg<'a>(args: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
|
||||
// Only the named keys, in order. Do NOT fall back to an arbitrary string
|
||||
// value: that would let an unrelated field be misread as the search symbol
|
||||
// or command for an advisory nudge.
|
||||
keys.iter()
|
||||
.find_map(|key| args.get(*key).and_then(Value::as_str))
|
||||
}
|
||||
|
||||
pub(super) fn shell_grep_symbol(command: &str) -> Option<String> {
|
||||
let mut parts = command.split_whitespace();
|
||||
let binary = strip_shell_quotes(parts.next()?)
|
||||
.rsplit('/')
|
||||
.next()?
|
||||
.to_string();
|
||||
if !matches!(binary.as_str(), "rg" | "ripgrep" | "grep" | "st") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut previous_took_value = false;
|
||||
for part in parts {
|
||||
let token = strip_shell_quotes(part);
|
||||
if previous_took_value {
|
||||
previous_took_value = false;
|
||||
continue;
|
||||
}
|
||||
if token.starts_with("--") {
|
||||
previous_took_value = matches!(
|
||||
token.as_str(),
|
||||
"--glob" | "--type" | "--context" | "--after-context" | "--before-context"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if token.starts_with('-') {
|
||||
continue;
|
||||
}
|
||||
return symbol_from_search_value(&token);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub(super) fn symbol_from_search_value(value: &str) -> Option<String> {
|
||||
let trimmed = strip_shell_quotes(value)
|
||||
.trim_matches('/')
|
||||
.replace("\\b", "")
|
||||
.replace(['^', '$'], "");
|
||||
for token in trimmed.split(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '_')) {
|
||||
if looks_like_symbol_token(token) {
|
||||
return Some(token.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn looks_like_symbol_token(token: &str) -> bool {
|
||||
if token.len() < 3 {
|
||||
return false;
|
||||
}
|
||||
let mut chars = token.chars();
|
||||
let Some(first) = chars.next() else {
|
||||
return false;
|
||||
};
|
||||
if !(first.is_ascii_alphabetic() || first == '_') {
|
||||
return false;
|
||||
}
|
||||
if !chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_') {
|
||||
return false;
|
||||
}
|
||||
|
||||
let lower = token.to_ascii_lowercase();
|
||||
if matches!(
|
||||
lower.as_str(),
|
||||
"todo" | "fixme" | "error" | "warning" | "debug" | "test" | "src" | "main"
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
token.contains('_') || token.chars().any(|ch| ch.is_ascii_uppercase())
|
||||
}
|
||||
|
||||
pub(super) fn strip_shell_quotes(value: &str) -> String {
|
||||
value
|
||||
.trim()
|
||||
.trim_matches('"')
|
||||
.trim_matches('\'')
|
||||
.to_string()
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
use super::utils::*;
|
||||
use super::ToolGuardrailNudge;
|
||||
use crate::tools::execution::ToolCall;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// Depth cap for `value_payload_bytes`'s recursion. Model-controlled JSON has
|
||||
/// no natural nesting limit; a real write/edit payload never needs anywhere
|
||||
/// near this depth, so hitting it is itself a signal the payload is oversized.
|
||||
const MAX_PAYLOAD_DEPTH: usize = 32;
|
||||
|
||||
pub(super) fn write_payload_cap_nudge(
|
||||
call: &ToolCall,
|
||||
max_bytes: usize,
|
||||
) -> Option<ToolGuardrailNudge> {
|
||||
// A cap of 0 disables the nudge rather than firing on every non-empty write.
|
||||
if max_bytes == 0 {
|
||||
return None;
|
||||
}
|
||||
let tool_name = call.name.to_ascii_lowercase();
|
||||
if !is_write_or_edit_tool(&tool_name) {
|
||||
return None;
|
||||
}
|
||||
let bytes = object_args(&call.input)
|
||||
.map(write_payload_bytes)
|
||||
.unwrap_or(0);
|
||||
if bytes <= max_bytes {
|
||||
return None;
|
||||
}
|
||||
Some(ToolGuardrailNudge {
|
||||
call_id: call.id.clone(),
|
||||
kind: "write_payload_cap",
|
||||
content: format!(
|
||||
"The requested write/edit payload is too large for this proxy policy ({bytes} bytes > {max_bytes} bytes). Retry with a smaller targeted edit or split the change."
|
||||
),
|
||||
// Include a hash of the actual payload, not just its byte length, so
|
||||
// two unrelated oversized writes that happen to be the same size
|
||||
// (e.g. same-size templated files) don't collide and suppress each
|
||||
// other's nudge. Identical repeated calls (same target, same
|
||||
// content) still dedupe as intended.
|
||||
fingerprint: format!(
|
||||
"write_payload_cap:{}:{:x}:{bytes}:{max_bytes}",
|
||||
call.name,
|
||||
payload_content_hash(&call.input)
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn payload_content_hash(value: &Value) -> u64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
value.to_string().hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
pub(super) fn write_payload_bytes(args: &Map<String, Value>) -> usize {
|
||||
let payload_keys = [
|
||||
"content",
|
||||
"text",
|
||||
"new_content",
|
||||
"patch",
|
||||
"diff",
|
||||
"replacement",
|
||||
"data",
|
||||
"old_string",
|
||||
"new_string",
|
||||
];
|
||||
let known = args
|
||||
.iter()
|
||||
.filter(|(key, _)| payload_keys.contains(&key.as_str()))
|
||||
.fold(0usize, |acc, (_, value)| {
|
||||
acc.saturating_add(value_payload_bytes(value, 0))
|
||||
});
|
||||
if known > 0 {
|
||||
return known;
|
||||
}
|
||||
// None of the known field names matched (a tool with a non-standard
|
||||
// schema, e.g. `body`/`value`) -- fall back to summing every value in
|
||||
// the payload rather than silently reporting zero and never capping it.
|
||||
args.values().fold(0usize, |acc, value| {
|
||||
acc.saturating_add(value_payload_bytes(value, 0))
|
||||
})
|
||||
}
|
||||
|
||||
fn value_payload_bytes(value: &Value, depth: usize) -> usize {
|
||||
if depth >= MAX_PAYLOAD_DEPTH {
|
||||
// Treat implausibly deep nesting as oversized rather than recursing
|
||||
// further; a fixed large sentinel (not usize::MAX) keeps the
|
||||
// saturating sum below safe from overflowing back to a small value.
|
||||
return 1_000_000_000;
|
||||
}
|
||||
match value {
|
||||
Value::String(value) => value.len(),
|
||||
Value::Array(values) => values.iter().fold(0usize, |acc, v| {
|
||||
acc.saturating_add(value_payload_bytes(v, depth + 1))
|
||||
}),
|
||||
Value::Object(values) => values.values().fold(0usize, |acc, v| {
|
||||
acc.saturating_add(value_payload_bytes(v, depth + 1))
|
||||
}),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,340 @@
|
||||
use crate::render::{render_text, RenderOpts};
|
||||
use crate::transform::factsheet;
|
||||
use crate::transform::gate;
|
||||
use base64::Engine;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct AnthropicOpts {
|
||||
pub cols: usize,
|
||||
pub max_height_px: usize,
|
||||
pub chars_per_token: f64,
|
||||
/// Skip entirely below this many slab chars (per-image cost dominates).
|
||||
pub min_compress_chars: usize,
|
||||
/// Move tool descriptions/schema annotations into the imaged Tool Reference.
|
||||
pub compress_tools: bool,
|
||||
/// Image large `<system-reminder>` text blocks in the first user message.
|
||||
pub compress_reminders: bool,
|
||||
/// Image large `tool_result` text content across all user messages.
|
||||
pub compress_tool_results: bool,
|
||||
/// Per-block char floor for reminder/tool_result imaging.
|
||||
pub min_live_block_chars: usize,
|
||||
/// Cap on images per tool_result; source is truncated (with a paging marker)
|
||||
/// above this so one giant result can't blow Anthropic's 100-image/request cap.
|
||||
pub max_images_per_tool_result: usize,
|
||||
/// Aggregate ceiling on TOTAL image blocks in the outgoing request
|
||||
/// (client-supplied + every pxpipe pass combined). Anthropic rejects requests
|
||||
/// past ~100 images, so imaging stops once this budget is exhausted and the
|
||||
/// remaining regions pass through as text. The per-tool_result cap only bounds
|
||||
/// one result; this bounds the sum.
|
||||
pub max_total_images: usize,
|
||||
/// Collapse the OLD closed-tool-call conversation prefix into history image(s),
|
||||
/// keeping the recent tail as text. Default OFF — highest cache-stability risk
|
||||
/// (see `apply_history`); gate on live validation before enabling by default.
|
||||
pub compress_history: bool,
|
||||
/// Trailing messages kept as live text (never collapsed).
|
||||
pub keep_tail_messages: usize,
|
||||
/// Minimum collapsible messages; below this the cache-amortization math doesn't pay.
|
||||
pub min_collapse_prefix_messages: usize,
|
||||
/// Snap the collapse boundary to this message-grid so the rendered history PNG
|
||||
/// stays byte-identical across turns and keeps hitting the prompt cache.
|
||||
pub history_collapse_chunk: usize,
|
||||
}
|
||||
|
||||
impl Default for AnthropicOpts {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cols: crate::render::RenderOpts::default().cols,
|
||||
max_height_px: crate::render::RenderOpts::default().max_height_px,
|
||||
// Slab is dense; 2.0 is conservative vs the real ~1.9 (see pxpipe
|
||||
// SLAB_CHARS_PER_TOKEN) so the gate biases toward pass-through.
|
||||
chars_per_token: 2.0,
|
||||
min_compress_chars: 6_000,
|
||||
compress_tools: true,
|
||||
compress_reminders: true,
|
||||
compress_tool_results: true,
|
||||
min_live_block_chars: 2_000,
|
||||
max_images_per_tool_result: 10,
|
||||
// Leave headroom under Anthropic's 100-image/request cap for a few
|
||||
// client-supplied images the counter may miss on unusual shapes.
|
||||
max_total_images: 95,
|
||||
compress_history: false,
|
||||
keep_tail_messages: 8,
|
||||
min_collapse_prefix_messages: 20,
|
||||
history_collapse_chunk: 20,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Neutral framing header co-rendered into the slab image. Deliberately avoids
|
||||
/// "system prompt"/"authoritative" wording — pxpipe found that phrasing trips
|
||||
/// Anthropic's reasoning_extraction refusal (reads as a replayed/extracted
|
||||
/// prompt). First-party, matter-of-fact framing keeps the model reading it as
|
||||
/// this session's own reference material.
|
||||
pub(crate) const SLAB_HEADER: &str =
|
||||
"Reference context for this session, rendered as an image. Read it as text.\n\n";
|
||||
|
||||
/// Chars-per-token for the live-region gate (reminders, tool_results). Higher
|
||||
/// than the slab's 2.0 because that content is prose/log, not dense config —
|
||||
/// pxpipe uses 4 here, which is conservative (biases toward pass-through).
|
||||
pub(crate) const LIVE_CHARS_PER_TOKEN: f64 = 4.0;
|
||||
|
||||
/// Banner text bracketing the collapsed-history image(s). Constant (byte-stable).
|
||||
pub(crate) const HISTORY_INTRO: &str = "Transcript of EARLIER conversation turns, rendered as images below. \
|
||||
Attribute each turn strictly by its <user>/<assistant> tag. This is PAST context, not the live request.\n";
|
||||
pub(crate) const HISTORY_OUTRO: &str =
|
||||
"\n[End of earlier transcript. The live request follows in the messages below.]";
|
||||
|
||||
/// Concat the text of a system field (string or block array) and return
|
||||
/// `(text, last_cache_control)`. The cache_control is the caller's prefix
|
||||
/// breakpoint we will relocate onto the image.
|
||||
pub(crate) fn read_system(system: Option<&Value>) -> (String, Option<Value>) {
|
||||
match system {
|
||||
Some(Value::String(s)) => (s.clone(), None),
|
||||
Some(Value::Array(blocks)) => {
|
||||
let mut text = String::new();
|
||||
let mut cc = None;
|
||||
for b in blocks {
|
||||
if let Some(t) = b.get("text").and_then(|t| t.as_str()) {
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(t);
|
||||
}
|
||||
if let Some(c) = b.get("cache_control") {
|
||||
cc = Some(c.clone()); // last one wins (latest in prefix)
|
||||
}
|
||||
}
|
||||
(text, cc)
|
||||
}
|
||||
_ => (String::new(), None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Render one tool's full doc (prose + compact schema) for the imaged reference.
|
||||
pub(crate) fn render_tool_doc(tool: &Value) -> String {
|
||||
let name = tool.get("name").and_then(|n| n.as_str()).unwrap_or("");
|
||||
let desc = tool
|
||||
.get("description")
|
||||
.and_then(|d| d.as_str())
|
||||
.unwrap_or("");
|
||||
let schema = tool
|
||||
.get("input_schema")
|
||||
.map(|s| serde_json::to_string(s).unwrap_or_default())
|
||||
.unwrap_or_default();
|
||||
format!("## Tool: {name}\n{desc}\n```json\n{schema}\n```\n")
|
||||
}
|
||||
|
||||
/// Build a base64 image content block from PNG bytes.
|
||||
pub(crate) fn image_block(png: &[u8]) -> Value {
|
||||
let data = base64::engine::general_purpose::STANDARD.encode(png);
|
||||
json!({
|
||||
"type": "image",
|
||||
"source": { "type": "base64", "media_type": "image/png", "data": data }
|
||||
})
|
||||
}
|
||||
|
||||
/// Count image blocks already present in the request (client-supplied, incl.
|
||||
/// images nested inside `tool_result` content), so the added-image budget is
|
||||
/// measured against the true outgoing total under Anthropic's ~100-image cap.
|
||||
pub(crate) fn count_existing_images(root: &Value) -> usize {
|
||||
let Some(msgs) = root.get("messages").and_then(|m| m.as_array()) else {
|
||||
return 0;
|
||||
};
|
||||
let is_image = |b: &Value| b.get("type").and_then(|t| t.as_str()) == Some("image");
|
||||
let mut n = 0;
|
||||
for m in msgs {
|
||||
let Some(blocks) = m.get("content").and_then(|c| c.as_array()) else {
|
||||
continue;
|
||||
};
|
||||
for b in blocks {
|
||||
match b.get("type").and_then(|t| t.as_str()) {
|
||||
Some("image") => n += 1,
|
||||
Some("tool_result") => {
|
||||
if let Some(inner) = b.get("content").and_then(|c| c.as_array()) {
|
||||
n += inner.iter().filter(|x| is_image(x)).count();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
n
|
||||
}
|
||||
|
||||
/// Prepend `images` to the first user message's content. Converts string
|
||||
/// content to a text block first; inserts a fresh user message if none exists.
|
||||
pub(crate) fn prepend_images_to_first_user(obj: &mut Map<String, Value>, images: &mut Vec<Value>) {
|
||||
let msgs = obj
|
||||
.entry("messages")
|
||||
.or_insert_with(|| Value::Array(vec![]));
|
||||
let Some(arr) = msgs.as_array_mut() else {
|
||||
return;
|
||||
};
|
||||
let first_user = arr
|
||||
.iter()
|
||||
.position(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"));
|
||||
|
||||
match first_user {
|
||||
Some(idx) => {
|
||||
let content = arr[idx].as_object_mut().and_then(|m| m.get_mut("content"));
|
||||
match content {
|
||||
Some(Value::String(s)) => {
|
||||
let text = std::mem::take(s);
|
||||
let mut blocks = vec![json!({ "type": "text", "text": text })];
|
||||
let mut new = std::mem::take(images);
|
||||
new.append(&mut blocks);
|
||||
arr[idx]
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.insert("content".into(), Value::Array(new));
|
||||
}
|
||||
Some(Value::Array(existing)) => {
|
||||
let mut new = std::mem::take(images);
|
||||
new.append(existing);
|
||||
*existing = new;
|
||||
}
|
||||
_ => {
|
||||
arr[idx]
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.insert("content".into(), Value::Array(std::mem::take(images)));
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
arr.insert(
|
||||
0,
|
||||
json!({ "role": "user", "content": std::mem::take(images) }),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Rendered-live-block accumulators, folded into `TransformInfo` by the caller.
|
||||
pub(crate) struct LiveRender {
|
||||
pub(crate) blocks: Vec<Value>,
|
||||
pub(crate) bytes: usize,
|
||||
pub(crate) pixels: usize,
|
||||
pub(crate) dropped: usize,
|
||||
pub(crate) img_count: usize,
|
||||
}
|
||||
|
||||
/// Render `text` to image blocks for a live region (reminder / tool_result),
|
||||
/// gated on profitability AND the remaining `budget` of new images. Moves `cc`
|
||||
/// (the block's cache_control, if any) onto the last image and optionally
|
||||
/// appends a verbatim fact-sheet text block. Returns `None` (leave the block as
|
||||
/// text) when empty, not profitable, or over budget.
|
||||
pub(crate) fn render_live_block(
|
||||
text: &str,
|
||||
opts: &AnthropicOpts,
|
||||
cc: Option<Value>,
|
||||
append_factsheet: bool,
|
||||
budget: usize,
|
||||
) -> Option<LiveRender> {
|
||||
let images = render_text(
|
||||
text,
|
||||
RenderOpts {
|
||||
cols: opts.cols,
|
||||
max_height_px: opts.max_height_px,
|
||||
},
|
||||
);
|
||||
let chars = text.chars().count();
|
||||
if images.len() > budget {
|
||||
return None;
|
||||
}
|
||||
if !gate::is_profitable(&images, chars, LIVE_CHARS_PER_TOKEN) {
|
||||
return None;
|
||||
}
|
||||
let mut blocks: Vec<Value> = images.iter().map(|im| image_block(&im.png)).collect();
|
||||
if let (Some(cc), Some(last)) = (cc, blocks.last_mut()) {
|
||||
if let Some(o) = last.as_object_mut() {
|
||||
o.insert("cache_control".into(), cc);
|
||||
}
|
||||
}
|
||||
if append_factsheet {
|
||||
let fact = factsheet::fact_sheet_text(text);
|
||||
if !fact.is_empty() {
|
||||
blocks.push(json!({ "type": "text", "text": fact }));
|
||||
}
|
||||
}
|
||||
Some(LiveRender {
|
||||
blocks,
|
||||
bytes: images.iter().map(|im| im.png.len()).sum(),
|
||||
pixels: images
|
||||
.iter()
|
||||
.map(|im| im.width as usize * im.height as usize)
|
||||
.sum(),
|
||||
dropped: images.iter().map(|im| im.dropped).sum(),
|
||||
img_count: images.len(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Flatten a tool_result `content` (string, or an array of text/other blocks)
|
||||
/// to a single text string for rendering.
|
||||
pub(crate) fn tool_result_text(content: Option<&Value>) -> String {
|
||||
match content {
|
||||
Some(Value::String(s)) => s.clone(),
|
||||
Some(Value::Array(blocks)) => {
|
||||
let mut out = String::new();
|
||||
for b in blocks {
|
||||
if let Some(t) = b.get("text").and_then(|t| t.as_str()) {
|
||||
if !out.is_empty() {
|
||||
out.push('\n');
|
||||
}
|
||||
out.push_str(t);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// True when a tool_result `content` array already holds an image block. Imaging
|
||||
/// such a result would drop that image (only text sub-blocks are flattened), so
|
||||
/// the caller skips it.
|
||||
pub(crate) fn tool_result_has_image(content: Option<&Value>) -> bool {
|
||||
matches!(content, Some(Value::Array(blocks))
|
||||
if blocks
|
||||
.iter()
|
||||
.any(|b| b.get("type").and_then(|t| t.as_str()) == Some("image")))
|
||||
}
|
||||
|
||||
/// Truncate `text` so it renders to at most `max_images` pages, keeping a 60/40
|
||||
/// head/tail split with a paging marker in the middle. Returns the (possibly
|
||||
/// truncated) text and the count of chars elided. Renders once to measure — the
|
||||
/// oversized path is rare, so the extra render is acceptable.
|
||||
pub(crate) fn truncate_for_budget(
|
||||
text: &str,
|
||||
max_images: usize,
|
||||
opts: &AnthropicOpts,
|
||||
) -> (String, usize) {
|
||||
let render_opts = RenderOpts {
|
||||
cols: opts.cols,
|
||||
max_height_px: opts.max_height_px,
|
||||
};
|
||||
let full = render_text(text, render_opts);
|
||||
let cap = max_images.max(1);
|
||||
if full.len() <= cap {
|
||||
return (text.to_string(), 0);
|
||||
}
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
// Start from the linear estimate, then shrink and RE-RENDER until it truly
|
||||
// fits `cap` pages — the paging marker and heavy line-wrapping can push a
|
||||
// linear estimate over, so the cap must be verified, not assumed.
|
||||
let mut ratio = cap as f64 / full.len() as f64;
|
||||
loop {
|
||||
let keep = (((chars.len() as f64) * ratio) as usize).min(chars.len());
|
||||
let head_len = keep * 6 / 10;
|
||||
let tail_len = keep - head_len;
|
||||
let head: String = chars[..head_len].iter().collect();
|
||||
let tail: String = chars[chars.len() - tail_len..].iter().collect();
|
||||
let omitted = chars.len() - head_len - tail_len;
|
||||
let out = format!("{head}\n\n[... {omitted} chars omitted for length ...]\n\n{tail}");
|
||||
if render_text(&out, render_opts).len() <= cap || ratio < 0.05 {
|
||||
return (out, omitted);
|
||||
}
|
||||
ratio *= 0.8;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
use super::common::{image_block, tool_result_text, AnthropicOpts, HISTORY_INTRO, HISTORY_OUTRO};
|
||||
use crate::render::{render_text, RenderOpts};
|
||||
use crate::transform::gate;
|
||||
use crate::transform::info::TransformInfo;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// Collapse the OLD closed-tool-call message prefix into ONE synthetic user
|
||||
/// message holding history image(s); keep the recent tail as text.
|
||||
///
|
||||
/// **Cache stability** is the whole risk here. Two guarantees keep the rendered
|
||||
/// PNG byte-identical across turns so Anthropic prompt-caches it instead of
|
||||
/// re-creating it every turn: (1) the collapse boundary is snapped DOWN to a
|
||||
/// `history_collapse_chunk` message grid, so it only advances in steps and the
|
||||
/// serialized text is stable for a whole window; (2) the serializer and renderer
|
||||
/// are pure functions of the message bytes (no timestamps/rng, thinking blocks
|
||||
/// dropped deterministically). If either breaks, this NET-LOSES money — hence
|
||||
/// default-off until validated live.
|
||||
///
|
||||
/// **Correctness**: only a tool-CLOSED prefix is collapsed (every `tool_use` has
|
||||
/// its matching `tool_result` within the range), so no tool call is ever
|
||||
/// orphaned. The first user message (which carries the slab images) is protected.
|
||||
///
|
||||
/// NOTE: the synthetic message is role `user`, which can place it adjacent to the
|
||||
/// protected first user message. The Anthropic Messages API accepts consecutive
|
||||
/// same-role messages (pxpipe relies on this in production); images require the
|
||||
/// user role, so this is unavoidable for a history-image message.
|
||||
pub(crate) fn apply_history(
|
||||
root: &mut Value,
|
||||
opts: &AnthropicOpts,
|
||||
info: &mut TransformInfo,
|
||||
budget: &mut usize,
|
||||
) {
|
||||
let Some(messages) = root.get("messages").and_then(|m| m.as_array()) else {
|
||||
return;
|
||||
};
|
||||
let len = messages.len();
|
||||
// Protect the slab-bearing first user message: collapse starts after it.
|
||||
let Some(first_user) = messages
|
||||
.iter()
|
||||
.position(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let protected = first_user + 1;
|
||||
let cutoff = len.saturating_sub(opts.keep_tail_messages);
|
||||
if cutoff <= protected {
|
||||
return;
|
||||
}
|
||||
let Some(boundary) = find_closed_boundary(messages, cutoff, protected) else {
|
||||
return;
|
||||
};
|
||||
// Snap DOWN to the chunk grid (relative to `protected`) for byte-stability.
|
||||
let chunk = opts.history_collapse_chunk.max(1);
|
||||
let grid = protected + ((boundary - protected) / chunk) * chunk;
|
||||
// The grid line is NOT guaranteed tool-closed (parity shifts from text-only
|
||||
// turns or parallel tool spans can leave it mid-open-span). Re-snap to the
|
||||
// largest CLOSED boundary <= the grid line so no tool_use is orphaned into
|
||||
// the history image; correctness beats the grid's cache-stability here.
|
||||
let Some(snapped) = find_closed_boundary(messages, grid, protected) else {
|
||||
return;
|
||||
};
|
||||
if snapped.saturating_sub(protected) < opts.min_collapse_prefix_messages {
|
||||
return;
|
||||
}
|
||||
|
||||
let text = messages_to_history_text(messages, protected, snapped);
|
||||
if text.trim().is_empty() {
|
||||
return;
|
||||
}
|
||||
let images = render_text(
|
||||
&text,
|
||||
RenderOpts {
|
||||
cols: opts.cols,
|
||||
max_height_px: opts.max_height_px,
|
||||
},
|
||||
);
|
||||
if !gate::is_profitable(&images, text.chars().count(), opts.chars_per_token) {
|
||||
return;
|
||||
}
|
||||
if images.len() > *budget {
|
||||
return;
|
||||
}
|
||||
*budget -= images.len();
|
||||
|
||||
// Build the synthetic message content: intro, images, outro.
|
||||
let mut content: Vec<Value> = Vec::with_capacity(images.len() + 2);
|
||||
content.push(json!({ "type": "text", "text": HISTORY_INTRO }));
|
||||
content.extend(images.iter().map(|im| image_block(&im.png)));
|
||||
content.push(json!({ "type": "text", "text": HISTORY_OUTRO }));
|
||||
let synthetic = json!({ "role": "user", "content": content });
|
||||
|
||||
let collapsed_turns = snapped - protected;
|
||||
info.collapsed_turns = collapsed_turns;
|
||||
info.collapsed_chars = text.len();
|
||||
info.collapsed_images = images.len();
|
||||
info.image_count += images.len();
|
||||
info.image_bytes += images.iter().map(|im| im.png.len()).sum::<usize>();
|
||||
info.image_pixels += images
|
||||
.iter()
|
||||
.map(|im| im.width as usize * im.height as usize)
|
||||
.sum::<usize>();
|
||||
info.dropped_chars += images.iter().map(|im| im.dropped).sum::<usize>();
|
||||
info.compressed_chars += text.len();
|
||||
|
||||
// Splice: [0..protected] + synthetic + [snapped..].
|
||||
let arr = root
|
||||
.get_mut("messages")
|
||||
.and_then(|m| m.as_array_mut())
|
||||
.expect("messages was an array above");
|
||||
arr.splice(protected..snapped, std::iter::once(synthetic));
|
||||
}
|
||||
|
||||
/// Largest exclusive end `e` in `(from, cutoff]` where messages `[from..e)` open
|
||||
/// no tool call they don't also close. Returns `None` if none exists. Robust to
|
||||
/// interleaved/parallel tool calls via the open-id set.
|
||||
fn find_closed_boundary(messages: &[Value], cutoff: usize, from: usize) -> Option<usize> {
|
||||
let mut open: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||
let mut last_closed = None;
|
||||
for (i, m) in messages.iter().enumerate().take(cutoff).skip(from) {
|
||||
if let Some(blocks) = m.get("content").and_then(|c| c.as_array()) {
|
||||
for b in blocks {
|
||||
match b.get("type").and_then(|t| t.as_str()) {
|
||||
Some("tool_use") => {
|
||||
if let Some(id) = b.get("id").and_then(|i| i.as_str()) {
|
||||
open.insert(id.to_string());
|
||||
}
|
||||
}
|
||||
Some("tool_result") => {
|
||||
if let Some(id) = b.get("tool_use_id").and_then(|i| i.as_str()) {
|
||||
open.remove(id);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
if open.is_empty() {
|
||||
last_closed = Some(i + 1);
|
||||
}
|
||||
}
|
||||
last_closed
|
||||
}
|
||||
|
||||
/// Serialize messages `[from..to)` to `<role>…</role>` XML text. thinking blocks
|
||||
/// dropped; tool_use/tool_result flattened; inline images become `[image]`.
|
||||
fn messages_to_history_text(messages: &[Value], from: usize, to: usize) -> String {
|
||||
let mut out = String::new();
|
||||
for m in &messages[from..to] {
|
||||
let role = m.get("role").and_then(|r| r.as_str()).unwrap_or("user");
|
||||
let body = flatten_content(m.get("content"));
|
||||
out.push('<');
|
||||
out.push_str(role);
|
||||
out.push_str(">\n");
|
||||
out.push_str(&body);
|
||||
out.push_str("\n</");
|
||||
out.push_str(role);
|
||||
out.push_str(">\n");
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Flatten one message's content to text: text verbatim, tool_use/tool_result to
|
||||
/// a compact marker, thinking dropped, inline images to `[image]`.
|
||||
fn flatten_content(content: Option<&Value>) -> String {
|
||||
match content {
|
||||
Some(Value::String(s)) => s.clone(),
|
||||
Some(Value::Array(blocks)) => {
|
||||
let mut parts: Vec<String> = Vec::new();
|
||||
for b in blocks {
|
||||
match b.get("type").and_then(|t| t.as_str()) {
|
||||
Some("text") => {
|
||||
if let Some(t) = b.get("text").and_then(|t| t.as_str()) {
|
||||
parts.push(t.to_string());
|
||||
}
|
||||
}
|
||||
Some("tool_use") => {
|
||||
let name = b.get("name").and_then(|n| n.as_str()).unwrap_or("");
|
||||
let input = b
|
||||
.get("input")
|
||||
.map(|v| serde_json::to_string(v).unwrap_or_default())
|
||||
.unwrap_or_default();
|
||||
parts.push(format!("[tool_use {name} {input}]"));
|
||||
}
|
||||
Some("tool_result") => {
|
||||
parts.push(format!(
|
||||
"[tool_result {}]",
|
||||
tool_result_text(b.get("content"))
|
||||
));
|
||||
}
|
||||
Some("image") => parts.push("[image]".to_string()),
|
||||
// thinking / redacted_thinking / unknown: dropped.
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
parts.join("\n")
|
||||
}
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
pub mod common;
|
||||
pub mod history;
|
||||
pub mod reminders;
|
||||
pub mod slab;
|
||||
pub mod tool_results;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use crate::transform::info::TransformInfo;
|
||||
use serde_json::Value;
|
||||
|
||||
pub use common::AnthropicOpts;
|
||||
|
||||
/// Transform `root` (a parsed Anthropic Messages body) in place. Returns info;
|
||||
/// on any full skip `root` is left byte-identical to the input. Three
|
||||
/// independent passes — static slab, `<system-reminder>` blocks, and
|
||||
/// `tool_result` content — so a request below the slab floor can still get its
|
||||
/// live regions imaged (and vice-versa).
|
||||
pub fn transform(root: &mut Value, opts: &AnthropicOpts) -> TransformInfo {
|
||||
if !root.is_object() {
|
||||
return TransformInfo::skipped("parse_error");
|
||||
}
|
||||
let mut info = TransformInfo::default();
|
||||
// NEW-image budget: total ceiling minus images the client already sent, so
|
||||
// the combined passes never push the request past Anthropic's ~100-image cap.
|
||||
let mut budget = opts
|
||||
.max_total_images
|
||||
.saturating_sub(common::count_existing_images(root));
|
||||
let slab_reason = slab::apply_slab(root, opts, &mut info, &mut budget);
|
||||
// History runs BEFORE the live-region passes: it serializes the OLD message
|
||||
// prefix to text, so tool_result imaging must not have already replaced that
|
||||
// content with `[image]` placeholders. Reminders/tool_results then image only
|
||||
// what survives (the protected first message + the live tail).
|
||||
if opts.compress_history {
|
||||
history::apply_history(root, opts, &mut info, &mut budget);
|
||||
}
|
||||
if opts.compress_reminders {
|
||||
reminders::apply_reminders(root, opts, &mut info, &mut budget);
|
||||
}
|
||||
if opts.compress_tool_results {
|
||||
tool_results::apply_tool_results(root, opts, &mut info, &mut budget);
|
||||
}
|
||||
info.compressed = info.image_count > 0;
|
||||
// "applied" when anything imaged; otherwise report why the slab (the primary
|
||||
// path) declined.
|
||||
info.reason = if info.compressed {
|
||||
"applied"
|
||||
} else {
|
||||
slab_reason
|
||||
};
|
||||
info
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
use super::common::{render_live_block, AnthropicOpts};
|
||||
use crate::transform::info::TransformInfo;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Image large `<system-reminder>` text blocks in the first user message. These
|
||||
/// are per-turn injected context (env, hints) that Claude Code ships as separate
|
||||
/// text blocks; the big ones are pure token cost the model rarely needs to quote.
|
||||
pub(crate) fn apply_reminders(
|
||||
root: &mut Value,
|
||||
opts: &AnthropicOpts,
|
||||
info: &mut TransformInfo,
|
||||
budget: &mut usize,
|
||||
) {
|
||||
let Some(arr) = root.get_mut("messages").and_then(|m| m.as_array_mut()) else {
|
||||
return;
|
||||
};
|
||||
let Some(msg) = arr
|
||||
.iter_mut()
|
||||
.find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) else {
|
||||
return; // string content: reminders arrive only as array blocks
|
||||
};
|
||||
|
||||
let mut out: Vec<Value> = Vec::with_capacity(content.len());
|
||||
for block in std::mem::take(content) {
|
||||
let is_reminder_text = block.get("type").and_then(|t| t.as_str()) == Some("text")
|
||||
&& block.get("text").and_then(|t| t.as_str()).is_some_and(|t| {
|
||||
t.contains("<system-reminder>") && t.len() >= opts.min_live_block_chars
|
||||
});
|
||||
if !is_reminder_text {
|
||||
out.push(block);
|
||||
continue;
|
||||
}
|
||||
let text = block
|
||||
.get("text")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let cc = block.get("cache_control").cloned();
|
||||
match render_live_block(&text, opts, cc, false, *budget) {
|
||||
Some(r) => {
|
||||
*budget -= r.img_count;
|
||||
info.reminder_imgs += r.img_count;
|
||||
info.image_count += r.img_count;
|
||||
info.image_bytes += r.bytes;
|
||||
info.image_pixels += r.pixels;
|
||||
info.dropped_chars += r.dropped;
|
||||
info.compressed_chars += text.len();
|
||||
out.extend(r.blocks);
|
||||
}
|
||||
None => out.push(block),
|
||||
}
|
||||
}
|
||||
*content = out;
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
use super::common::{
|
||||
image_block, prepend_images_to_first_user, read_system, render_tool_doc, AnthropicOpts,
|
||||
SLAB_HEADER,
|
||||
};
|
||||
use crate::render::{render_text, RenderOpts};
|
||||
use crate::transform::factsheet;
|
||||
use crate::transform::gate;
|
||||
use crate::transform::info::TransformInfo;
|
||||
use crate::transform::schema_strip;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Static system+tools slab pass. Returns the skip reason (or "applied"); only
|
||||
/// mutates `root` / fills `info` when profitable.
|
||||
pub(crate) fn apply_slab(
|
||||
root: &mut Value,
|
||||
opts: &AnthropicOpts,
|
||||
info: &mut TransformInfo,
|
||||
budget: &mut usize,
|
||||
) -> &'static str {
|
||||
let (system_text, system_cc) = read_system(root.get("system"));
|
||||
|
||||
// Capture a tool cache_control as a fallback anchor (tools sit before system
|
||||
// in the cache prefix, so system's wins when both exist).
|
||||
let tools_present = opts.compress_tools
|
||||
&& root
|
||||
.get("tools")
|
||||
.and_then(|t| t.as_array())
|
||||
.is_some_and(|a| !a.is_empty());
|
||||
let mut tool_cc = None;
|
||||
let mut tool_ref = String::new();
|
||||
if tools_present {
|
||||
if let Some(arr) = root.get("tools").and_then(|t| t.as_array()) {
|
||||
for t in arr {
|
||||
tool_ref.push_str(&render_tool_doc(t));
|
||||
if let Some(c) = t.get("cache_control") {
|
||||
tool_cc = Some(c.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Assemble the slab: header + system + tool reference.
|
||||
let mut slab = String::with_capacity(system_text.len() + tool_ref.len() + 64);
|
||||
slab.push_str(SLAB_HEADER);
|
||||
slab.push_str(&system_text);
|
||||
if !tool_ref.is_empty() {
|
||||
slab.push_str("\n\n# Tool Reference\n");
|
||||
slab.push_str(&tool_ref);
|
||||
}
|
||||
|
||||
let slab_chars = slab.chars().count();
|
||||
// Nothing meaningful to image (no system text and no tools).
|
||||
if system_text.trim().is_empty() && tool_ref.is_empty() {
|
||||
return "no_slab";
|
||||
}
|
||||
if slab_chars < opts.min_compress_chars {
|
||||
return "below_min_chars";
|
||||
}
|
||||
|
||||
let images = render_text(
|
||||
&slab,
|
||||
RenderOpts {
|
||||
cols: opts.cols,
|
||||
max_height_px: opts.max_height_px,
|
||||
},
|
||||
);
|
||||
if !gate::is_profitable(&images, slab_chars, opts.chars_per_token) {
|
||||
return "not_profitable";
|
||||
}
|
||||
if images.len() > *budget {
|
||||
return "image_budget";
|
||||
}
|
||||
|
||||
// ---- commit: mutate root ------------------------------------------------
|
||||
let anchor = system_cc.or(tool_cc); // relocate this onto the last image
|
||||
let fact = factsheet::fact_sheet_text(&slab);
|
||||
|
||||
let mut image_blocks: Vec<Value> = images.iter().map(|im| image_block(&im.png)).collect();
|
||||
let image_bytes: usize = images.iter().map(|im| im.png.len()).sum();
|
||||
let image_pixels: usize = images
|
||||
.iter()
|
||||
.map(|im| im.width as usize * im.height as usize)
|
||||
.sum();
|
||||
let dropped: usize = images.iter().map(|im| im.dropped).sum();
|
||||
|
||||
// Relocate the caller's cache breakpoint onto the LAST image so the whole
|
||||
// imaged prefix caches as one stable segment. pxpipe never *adds* a marker.
|
||||
let relocated = if let (Some(cc), Some(last)) = (anchor, image_blocks.last_mut()) {
|
||||
if let Some(obj) = last.as_object_mut() {
|
||||
obj.insert("cache_control".into(), cc);
|
||||
}
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
let obj = root.as_object_mut().expect("checked is_object above");
|
||||
|
||||
// 1. system -> short pointer (+ factsheet). Removes original cache_control
|
||||
// (relocated onto the image) since we're replacing the whole field.
|
||||
let mut pointer = String::from(
|
||||
"[Your reference context for this session is provided as an image in the first user message. Read it there.]",
|
||||
);
|
||||
if !fact.is_empty() {
|
||||
pointer.push('\n');
|
||||
pointer.push_str(&fact);
|
||||
}
|
||||
obj.insert("system".into(), Value::String(pointer));
|
||||
|
||||
// 2. tools -> stub description + annotation-stripped schema, drop the
|
||||
// relocated cache_control. Keep name + structural schema so Anthropic's
|
||||
// tool-use validator still accepts calls.
|
||||
if tools_present {
|
||||
if let Some(arr) = obj.get_mut("tools").and_then(|t| t.as_array_mut()) {
|
||||
for t in arr.iter_mut() {
|
||||
let Some(tm) = t.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
let name = tm
|
||||
.get("name")
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
tm.insert(
|
||||
"description".into(),
|
||||
Value::String(format!("See \"## Tool: {name}\" in the reference image.")),
|
||||
);
|
||||
if let Some(schema) = tm.get("input_schema") {
|
||||
let stripped = schema_strip::strip(schema);
|
||||
if schema_strip::has_structure(&stripped) {
|
||||
tm.insert("input_schema".into(), stripped);
|
||||
}
|
||||
// else keep original: a bare {type:object} stub causes 400s.
|
||||
}
|
||||
tm.remove("cache_control");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Prepend image blocks to the first user message (system rejects images).
|
||||
prepend_images_to_first_user(obj, &mut image_blocks);
|
||||
|
||||
info.compressed_chars += slab.len();
|
||||
info.image_count += images.len();
|
||||
info.image_bytes += image_bytes;
|
||||
info.image_pixels += image_pixels;
|
||||
info.dropped_chars += dropped;
|
||||
info.relocated_cache_anchor = relocated;
|
||||
*budget -= images.len();
|
||||
"applied"
|
||||
}
|
||||
@@ -0,0 +1,409 @@
|
||||
use super::common::AnthropicOpts;
|
||||
use super::transform;
|
||||
use serde_json::json;
|
||||
|
||||
fn big_system() -> String {
|
||||
// ~14k chars with some notable identifiers for the factsheet.
|
||||
let mut s = String::from("Operating rules. See src/main.rs and CONFIG_PATH env.\n");
|
||||
s.push_str(&"Do the thing carefully and precisely. ".repeat(320));
|
||||
s.push_str(" commit deadbeef1 flag --verbose");
|
||||
s
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn images_system_and_relocates_anchor() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": [
|
||||
{ "type": "text", "text": big_system(), "cache_control": { "type": "ephemeral" } }
|
||||
],
|
||||
"messages": [ { "role": "user", "content": "hi there" } ]
|
||||
});
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(info.compressed, "reason={}", info.reason);
|
||||
assert!(info.relocated_cache_anchor);
|
||||
|
||||
// system is now a short pointer string.
|
||||
assert!(root["system"].is_string());
|
||||
assert!(root["system"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("first user message"));
|
||||
|
||||
// first user message leads with an image block carrying the anchor.
|
||||
let content = root["messages"][0]["content"].as_array().unwrap();
|
||||
assert_eq!(content[0]["type"], "image");
|
||||
assert_eq!(content.last().unwrap()["type"], "text"); // original "hi there"
|
||||
assert_eq!(content.last().unwrap()["text"], "hi there");
|
||||
// anchor is on exactly the last image, nowhere in system.
|
||||
let last_img = content.iter().rev().find(|b| b["type"] == "image").unwrap();
|
||||
assert!(last_img.get("cache_control").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_system_passes_through_unchanged() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "be helpful",
|
||||
"messages": [ { "role": "user", "content": "hi" } ]
|
||||
});
|
||||
let before = root.clone();
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(!info.compressed);
|
||||
assert_eq!(info.reason, "below_min_chars");
|
||||
assert_eq!(root, before, "skip path must not mutate the body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_input_same_output_cache_stable() {
|
||||
let make = || {
|
||||
json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": big_system(),
|
||||
"messages": [ { "role": "user", "content": "go" } ]
|
||||
})
|
||||
};
|
||||
let mut a = make();
|
||||
let mut b = make();
|
||||
transform(&mut a, &AnthropicOpts::default());
|
||||
transform(&mut b, &AnthropicOpts::default());
|
||||
assert_eq!(a, b, "transform must be deterministic (cache stability)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tools_stubbed_and_schema_stripped() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": big_system(),
|
||||
"tools": [{
|
||||
"name": "edit_file",
|
||||
"description": "very long tool description ".repeat(50),
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"description": "annotation",
|
||||
"properties": { "path": { "type": "string", "description": "p" } },
|
||||
"required": ["path"]
|
||||
},
|
||||
"cache_control": { "type": "ephemeral" }
|
||||
}],
|
||||
"messages": [ { "role": "user", "content": "go" } ]
|
||||
});
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(info.compressed);
|
||||
let tool = &root["tools"][0];
|
||||
assert!(tool["description"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("## Tool: edit_file"));
|
||||
assert!(tool["input_schema"].get("description").is_none());
|
||||
assert_eq!(tool["input_schema"]["required"], json!(["path"]));
|
||||
assert!(tool.get("cache_control").is_none(), "tool anchor relocated");
|
||||
}
|
||||
|
||||
/// Dense filler text large enough to clear the live-block gate.
|
||||
fn big_text() -> String {
|
||||
"log line with detail /var/run/app.sock port 8080 CODE_X 12345 ".repeat(300)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn images_large_reminder_block() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "small",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{ "type": "text", "text": format!("<system-reminder>{}</system-reminder>", big_text()) },
|
||||
{ "type": "text", "text": "the actual question" }
|
||||
]
|
||||
}]
|
||||
});
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(info.compressed);
|
||||
assert!(info.reminder_imgs >= 1);
|
||||
let content = root["messages"][0]["content"].as_array().unwrap();
|
||||
assert_eq!(content[0]["type"], "image", "reminder replaced by image");
|
||||
// The real question text survives untouched.
|
||||
assert!(content.iter().any(|b| b["text"] == "the actual question"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn images_large_tool_result_with_factsheet() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "small",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"content": big_text(),
|
||||
"cache_control": { "type": "ephemeral" }
|
||||
}]
|
||||
}]
|
||||
});
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(info.compressed);
|
||||
assert!(info.tool_result_imgs >= 1);
|
||||
let tr = &root["messages"][0]["content"][0];
|
||||
assert!(
|
||||
tr.get("cache_control").is_none(),
|
||||
"anchor relocated onto image"
|
||||
);
|
||||
let inner = tr["content"].as_array().unwrap();
|
||||
assert_eq!(inner[0]["type"], "image");
|
||||
// fact-sheet text block rides alongside the image (paths/ids survive OCR).
|
||||
assert!(inner
|
||||
.iter()
|
||||
.any(|b| b["type"] == "text" && b["text"].as_str().unwrap().contains("/var/run/app.sock")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_error_tool_result() {
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "small",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"is_error": true,
|
||||
"content": big_text()
|
||||
}]
|
||||
}]
|
||||
});
|
||||
let before = root.clone();
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(!info.compressed, "error tool_results must not be imaged");
|
||||
assert_eq!(root, before);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pages_oversized_tool_result() {
|
||||
// Force a tiny image cap so a modest result trips the paging path.
|
||||
let opts = AnthropicOpts {
|
||||
max_images_per_tool_result: 1,
|
||||
max_height_px: 80, // ~9 rows/page → many pages before truncation
|
||||
..AnthropicOpts::default()
|
||||
};
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "small",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"content": big_text()
|
||||
}]
|
||||
}]
|
||||
});
|
||||
let info = transform(&mut root, &opts);
|
||||
assert!(info.compressed);
|
||||
assert_eq!(info.truncated_tool_results, 1);
|
||||
assert!(info.omitted_chars > 0);
|
||||
}
|
||||
|
||||
/// A conversation with `pairs` closed assistant(tool_use)/user(tool_result) turns.
|
||||
fn convo(pairs: usize) -> serde_json::Value {
|
||||
let mut msgs = vec![json!({ "role": "user", "content": "start the task" })];
|
||||
for i in 0..pairs {
|
||||
msgs.push(json!({
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{ "type": "text", "text": format!("step {i}") },
|
||||
{ "type": "tool_use", "id": format!("tu_{i}"), "name": "run", "input": { "cmd": format!("do {i}") } }
|
||||
]
|
||||
}));
|
||||
msgs.push(json!({
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": format!("tu_{i}"),
|
||||
"content": format!("output /path/file{i}.rs {}", "detail ".repeat(60))
|
||||
}]
|
||||
}));
|
||||
}
|
||||
json!({ "model": "claude-fable-5", "system": "small", "messages": msgs })
|
||||
}
|
||||
|
||||
fn history_opts() -> AnthropicOpts {
|
||||
AnthropicOpts {
|
||||
compress_history: true,
|
||||
// Isolate the history pass so tail tool_results don't also image.
|
||||
compress_tool_results: false,
|
||||
compress_reminders: false,
|
||||
..AnthropicOpts::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collapses_closed_history_prefix() {
|
||||
let mut root = convo(18); // 37 messages
|
||||
let before_len = root["messages"].as_array().unwrap().len();
|
||||
let info = transform(&mut root, &history_opts());
|
||||
assert!(info.compressed, "reason={}", info.reason);
|
||||
assert_eq!(info.collapsed_turns, 20, "snapped to the 20-message grid");
|
||||
assert!(info.collapsed_images >= 1);
|
||||
|
||||
let msgs = root["messages"].as_array().unwrap();
|
||||
assert!(msgs.len() < before_len, "prefix collapsed");
|
||||
// First user message is protected (untouched).
|
||||
assert_eq!(msgs[0]["content"], "start the task");
|
||||
// Synthetic history message: intro text, then an image.
|
||||
let syn = msgs[1]["content"].as_array().unwrap();
|
||||
assert!(syn[0]["text"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("EARLIER conversation"));
|
||||
assert!(syn.iter().any(|b| b["type"] == "image"));
|
||||
assert!(syn.last().unwrap()["text"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("live request follows"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn history_render_is_cache_stable() {
|
||||
let mut a = convo(18);
|
||||
let mut b = convo(18);
|
||||
transform(&mut a, &history_opts());
|
||||
transform(&mut b, &history_opts());
|
||||
assert_eq!(a, b, "same conversation must collapse to identical bytes");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_history_not_collapsed() {
|
||||
// Below min_collapse_prefix_messages (20) → left as text.
|
||||
let mut root = convo(5); // 11 messages
|
||||
let before = root.clone();
|
||||
let info = transform(&mut root, &history_opts());
|
||||
assert!(!info.compressed);
|
||||
assert_eq!(root, before);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_image_budget_is_enforced() {
|
||||
// Many large tool_results but a tiny total budget: imaging stops at the
|
||||
// cap so the request can't exceed Anthropic's per-request image limit.
|
||||
let mut msgs = vec![json!({ "role": "user", "content": "start" })];
|
||||
for i in 0..10 {
|
||||
msgs.push(json!({
|
||||
"role": "user",
|
||||
"content": [{ "type": "tool_result", "tool_use_id": format!("tu_{i}"), "content": big_text() }]
|
||||
}));
|
||||
}
|
||||
let mut root = json!({ "model": "claude-fable-5", "system": "small", "messages": msgs });
|
||||
let opts = AnthropicOpts {
|
||||
max_total_images: 3,
|
||||
..AnthropicOpts::default()
|
||||
};
|
||||
let info = transform(&mut root, &opts);
|
||||
assert!(info.compressed);
|
||||
assert!(
|
||||
info.image_count <= 3,
|
||||
"aggregate image budget exceeded: {}",
|
||||
info.image_count
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_result_carrying_an_image_is_left_alone() {
|
||||
// A screenshot + long log in one tool_result: imaging the log would drop
|
||||
// the screenshot, so the whole result must pass through untouched.
|
||||
let mut root = json!({
|
||||
"model": "claude-fable-5",
|
||||
"system": "small",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"content": [
|
||||
{ "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": "AAAA" } },
|
||||
{ "type": "text", "text": big_text() }
|
||||
]
|
||||
}]
|
||||
}]
|
||||
});
|
||||
let before = root.clone();
|
||||
let info = transform(&mut root, &AnthropicOpts::default());
|
||||
assert!(
|
||||
!info.compressed,
|
||||
"tool_result with an image must not be imaged"
|
||||
);
|
||||
assert_eq!(root, before);
|
||||
}
|
||||
|
||||
/// True unless a tool_result references a tool_use id not present as a real
|
||||
/// block earlier in the message list (an orphan the API rejects).
|
||||
fn no_orphaned_tool_results(root: &serde_json::Value) -> bool {
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
for m in root["messages"].as_array().unwrap() {
|
||||
let Some(blocks) = m.get("content").and_then(|c| c.as_array()) else {
|
||||
continue;
|
||||
};
|
||||
for b in blocks {
|
||||
match b.get("type").and_then(|t| t.as_str()) {
|
||||
Some("tool_use") => {
|
||||
if let Some(id) = b.get("id").and_then(|i| i.as_str()) {
|
||||
seen.insert(id.to_string());
|
||||
}
|
||||
}
|
||||
Some("tool_result") => {
|
||||
if let Some(id) = b.get("tool_use_id").and_then(|i| i.as_str()) {
|
||||
if !seen.contains(id) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn history_snap_never_orphans_a_tool_result() {
|
||||
// A text-only assistant turn right after the first user message shifts
|
||||
// parity so the chunk-grid line (61) lands mid-pair on an OPEN tool_use;
|
||||
// the snap must back off to the nearest closed boundary (60) instead of
|
||||
// orphaning tu_29's tool_result into the tail.
|
||||
let mut msgs = vec![json!({ "role": "user", "content": "start" })];
|
||||
msgs.push(json!({ "role": "assistant", "content": [{ "type": "text", "text": "thinking out loud" }] }));
|
||||
for i in 0..40 {
|
||||
msgs.push(json!({
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{ "type": "text", "text": format!("step {i}") },
|
||||
{ "type": "tool_use", "id": format!("tu_{i}"), "name": "run", "input": { "cmd": format!("do {i}") } }
|
||||
]
|
||||
}));
|
||||
msgs.push(json!({
|
||||
"role": "user",
|
||||
"content": [{ "type": "tool_result", "tool_use_id": format!("tu_{i}"), "content": format!("out /p/f{i}.rs {}", "detail ".repeat(60)) }]
|
||||
}));
|
||||
}
|
||||
let mut root = json!({ "model": "claude-fable-5", "system": "small", "messages": msgs });
|
||||
let info = transform(&mut root, &history_opts());
|
||||
assert!(info.compressed, "reason={}", info.reason);
|
||||
assert!(
|
||||
no_orphaned_tool_results(&root),
|
||||
"history collapse orphaned a tool_result"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn open_tool_call_not_crossed() {
|
||||
// Last prefix turn leaves a tool_use unmatched (open); the boundary must
|
||||
// stop before it so no tool call is orphaned into the image.
|
||||
let mut root = convo(18);
|
||||
// Drop the tool_result of an early pair to open a call at message 4.
|
||||
root["messages"][4] = json!({ "role": "user", "content": "no tool result here" });
|
||||
let info = transform(&mut root, &history_opts());
|
||||
// With an open call at msg 3 (tu_1) never closed, the closed boundary
|
||||
// can't advance past msg 2, so the 20-message grid step never fills.
|
||||
assert!(!info.compressed || info.collapsed_turns == 0);
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
use super::common::{
|
||||
render_live_block, tool_result_has_image, tool_result_text, truncate_for_budget, AnthropicOpts,
|
||||
};
|
||||
use crate::transform::info::TransformInfo;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Image large `tool_result` text content across every user message. tool_result
|
||||
/// output (find trees, file dumps, logs) is the bulk of agentic input. Skips
|
||||
/// `is_error` results (Anthropic rejects images inside those) and pages oversized
|
||||
/// content down to the image cap.
|
||||
pub(crate) fn apply_tool_results(
|
||||
root: &mut Value,
|
||||
opts: &AnthropicOpts,
|
||||
info: &mut TransformInfo,
|
||||
budget: &mut usize,
|
||||
) {
|
||||
let Some(msgs) = root.get_mut("messages").and_then(|m| m.as_array_mut()) else {
|
||||
return;
|
||||
};
|
||||
for msg in msgs.iter_mut() {
|
||||
if *budget == 0 {
|
||||
break;
|
||||
}
|
||||
let Some(blocks) = msg.get_mut("content").and_then(|c| c.as_array_mut()) else {
|
||||
continue;
|
||||
};
|
||||
for block in blocks.iter_mut() {
|
||||
let Some(bm) = block.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
if bm.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
|
||||
continue;
|
||||
}
|
||||
if bm.get("is_error").and_then(|e| e.as_bool()) == Some(true) {
|
||||
continue; // images forbidden in error tool_results
|
||||
}
|
||||
// Don't clobber a tool_result that already carries an image block
|
||||
// (e.g. a screenshot the tool returned) — tool_result_text only reads
|
||||
// the text sub-blocks, so imaging here would silently drop that image.
|
||||
if tool_result_has_image(bm.get("content")) {
|
||||
continue;
|
||||
}
|
||||
let text = tool_result_text(bm.get("content"));
|
||||
if text.len() < opts.min_live_block_chars {
|
||||
continue;
|
||||
}
|
||||
// Cap each result at the per-result limit AND whatever total budget
|
||||
// is left, so the aggregate can't exceed max_total_images.
|
||||
let per_result = opts.max_images_per_tool_result.min(*budget);
|
||||
if per_result == 0 {
|
||||
break;
|
||||
}
|
||||
let (rendered, omitted) = truncate_for_budget(&text, per_result, opts);
|
||||
let cc = bm.get("cache_control").cloned();
|
||||
if let Some(r) = render_live_block(&rendered, opts, cc, true, *budget) {
|
||||
*budget -= r.img_count;
|
||||
if omitted > 0 {
|
||||
info.truncated_tool_results += 1;
|
||||
info.omitted_chars += omitted;
|
||||
}
|
||||
info.tool_result_imgs += r.img_count;
|
||||
info.image_count += r.img_count;
|
||||
info.image_bytes += r.bytes;
|
||||
info.image_pixels += r.pixels;
|
||||
info.dropped_chars += r.dropped;
|
||||
info.compressed_chars += text.len();
|
||||
bm.remove("cache_control"); // relocated onto the last image
|
||||
bm.insert("content".into(), Value::Array(r.blocks));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+10
-475
@@ -17,16 +17,16 @@ pub struct CommandDetection {
|
||||
pub category: String,
|
||||
}
|
||||
|
||||
struct Detector {
|
||||
r#type: &'static str,
|
||||
category: &'static str,
|
||||
command_patterns: Vec<Regex>,
|
||||
content_patterns: Vec<Regex>,
|
||||
pub(super) struct Detector {
|
||||
pub(super) r#type: &'static str,
|
||||
pub(super) category: &'static str,
|
||||
pub(super) command_patterns: Vec<Regex>,
|
||||
pub(super) content_patterns: Vec<Regex>,
|
||||
}
|
||||
|
||||
/// Compile with explicit flags: `i` = case-insensitive, `m` = multiline.
|
||||
/// Fixed, in-tree patterns — a bad transcription panics (caught by `table_builds`).
|
||||
fn rx(pattern: &str, flags: &str) -> Regex {
|
||||
pub(super) fn rx(pattern: &str, flags: &str) -> Regex {
|
||||
RegexBuilder::new(pattern)
|
||||
.case_insensitive(flags.contains('i'))
|
||||
.multi_line(flags.contains('m'))
|
||||
@@ -96,486 +96,21 @@ fn command_prefix_re() -> &'static Regex {
|
||||
RE.get_or_init(|| Regex::new(&format!(r"^(?:{})\b", COMMAND_PREFIXES.join("|"))).unwrap())
|
||||
}
|
||||
|
||||
mod table;
|
||||
|
||||
fn detectors() -> &'static Vec<Detector> {
|
||||
static D: OnceLock<Vec<Detector>> = OnceLock::new();
|
||||
D.get_or_init(build_detectors)
|
||||
D.get_or_init(table::build_detectors)
|
||||
}
|
||||
|
||||
/// (type, category, command patterns (all "i"), (content pattern, flags) pairs).
|
||||
type DetectorRow = (
|
||||
pub(super) type DetectorRow = (
|
||||
&'static str,
|
||||
&'static str,
|
||||
&'static [&'static str],
|
||||
&'static [(&'static str, &'static str)],
|
||||
);
|
||||
|
||||
fn build_detectors() -> Vec<Detector> {
|
||||
let table: &[DetectorRow] = &[
|
||||
(
|
||||
"git-status",
|
||||
"git",
|
||||
&[r"^git\s+status\b"],
|
||||
&[
|
||||
(r"^On branch ", "m"),
|
||||
(r"^Changes (?:not staged|to be committed)", "m"),
|
||||
(r"^Untracked files:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-branch",
|
||||
"git",
|
||||
&[r"^git\s+branch\b", r"^git\s+checkout\b", r"^git\s+switch\b"],
|
||||
&[
|
||||
(r"^\*\s+\S+", "m"),
|
||||
(r"Switched to (?:a new )?branch", "i"),
|
||||
(r#"Already on ['"][^'"]+['"]"#, "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-diff",
|
||||
"git",
|
||||
&[r"^git\s+diff\b", r"^git\s+show\b"],
|
||||
&[
|
||||
(r"^diff --git ", "m"),
|
||||
(r"^@@\s+-\d+,\d+\s+\+\d+,\d+\s+@@", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-log",
|
||||
"git",
|
||||
&[r"^git\s+log\b"],
|
||||
&[(r"^commit [0-9a-f]{7,40}", "m"), (r"^Author: ", "m")],
|
||||
),
|
||||
(
|
||||
"make",
|
||||
"build",
|
||||
&[r"^make\b"],
|
||||
&[
|
||||
(r"^make\[\d+\]: (?:Entering|Leaving) directory", "m"),
|
||||
(r"make: \*\*\* ", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"gradle",
|
||||
"build",
|
||||
&[r"^(?:gradle|gradlew|\./gradlew)\b"],
|
||||
&[
|
||||
(r"^> Task :", "m"),
|
||||
(r"^BUILD (?:SUCCESSFUL|FAILED)\b", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"dotnet",
|
||||
"build",
|
||||
&[
|
||||
r"^dotnet\s+(?:build|test|run|restore|publish|pack|msbuild)\b",
|
||||
r"^dotnet\b",
|
||||
],
|
||||
&[
|
||||
(r"^Build (?:succeeded|FAILED)\b", "m"),
|
||||
(r"\b(?:error|warning) CS\d+\b", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"terraform-plan",
|
||||
"infra",
|
||||
&[r"^terraform\s+plan\b"],
|
||||
&[
|
||||
(r"Terraform will perform the following actions:", ""),
|
||||
(r"Plan: \d+ to add", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"tofu-plan",
|
||||
"infra",
|
||||
&[r"^(?:tofu|opentofu)\s+plan\b"],
|
||||
&[
|
||||
(r"OpenTofu will perform the following actions:", ""),
|
||||
(r"Plan: \d+ to add", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"systemctl-status",
|
||||
"infra",
|
||||
&[r"^systemctl\s+status\b"],
|
||||
&[
|
||||
(r"^\s*Loaded:\s+", "m"),
|
||||
(r"^\s*Active:\s+", "m"),
|
||||
(r"^●\s+\S+\.service", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-vitest",
|
||||
"test",
|
||||
&[r"^vitest\b", r"^npm\s+(?:run\s+)?test:vitest\b"],
|
||||
&[
|
||||
(r"\bvitest\b", "i"),
|
||||
(r"^ ✓ ", "m"),
|
||||
(r"^ ❯ ", "m"),
|
||||
(r"Test Files\s+\d+\s+(?:passed|failed)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-jest",
|
||||
"test",
|
||||
&[r"^jest\b", r"^npm\s+(?:run\s+)?test\b"],
|
||||
&[
|
||||
(r"Test Suites:\s+\d+", "i"),
|
||||
(r"Tests:\s+\d+", "i"),
|
||||
(r"^PASS\s+", "m"),
|
||||
(r"^FAIL\s+", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-pytest",
|
||||
"test",
|
||||
&[r"^pytest\b", r"^python\s+-m\s+pytest\b"],
|
||||
&[
|
||||
(r"=+\s+(?:\d+\s+)?(?:passed|failed|errors?)", "i"),
|
||||
(r"^E\s+", "m"),
|
||||
(r"^FAILED ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-cargo",
|
||||
"test",
|
||||
&[r"^cargo\s+test\b", r"^cargo\s+nextest\b"],
|
||||
&[
|
||||
(r"^running \d+ tests?", "m"),
|
||||
(r"^test\s+[\w:.-]+\s+\.\.\.\s+(?:ok|FAILED|ignored)", "m"),
|
||||
(r"test result:\s+(?:ok|FAILED)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-go",
|
||||
"test",
|
||||
&[r"^go\s+test\b"],
|
||||
&[
|
||||
(r"^(?:ok|FAIL)\s+[\w./-]+\s+[\d.]+s", "m"),
|
||||
(r"^--- FAIL: ", "m"),
|
||||
(r"^panic: ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-typescript",
|
||||
"build",
|
||||
&[r"^tsc\b", r"^npm\s+run\s+typecheck\b"],
|
||||
&[(r"TS\d{4}:", ""), (r"error TS\d{4}", "i")],
|
||||
),
|
||||
(
|
||||
"build-eslint",
|
||||
"build",
|
||||
&[r"^eslint\b", r"^npm\s+run\s+lint\b"],
|
||||
&[
|
||||
(r"\s+\d+:\d+\s+(?:error|warning)\s+", ""),
|
||||
(r"✖\s+\d+\s+problems?", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-webpack",
|
||||
"build",
|
||||
&[
|
||||
r"^webpack\b",
|
||||
r"^npx\s+webpack\b",
|
||||
r"^npm\s+run\s+build:webpack\b",
|
||||
],
|
||||
&[
|
||||
(r"webpack\s+\d", "i"),
|
||||
(r"compiled (?:successfully|with \d+ errors?)", "i"),
|
||||
(r"asset .+\.js", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-vite",
|
||||
"build",
|
||||
&[
|
||||
r"^vite\s+build\b",
|
||||
r"^npm\s+run\s+build\b",
|
||||
r"^pnpm\s+build\b",
|
||||
],
|
||||
&[
|
||||
(r"vite v[\d.]+", "i"),
|
||||
(r"✓ built in", "i"),
|
||||
(r"transforming \(\d+\)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"biome",
|
||||
"build",
|
||||
&[r"^biome\b", r"^npx\s+biome\b"],
|
||||
&[
|
||||
(r"lint/[A-Za-z0-9/.-]+", ""),
|
||||
(r"Checked \d+ files? in", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"prettier",
|
||||
"build",
|
||||
&[r"^prettier\b", r"^npx\s+prettier\b"],
|
||||
&[
|
||||
(r"^Checking formatting\.\.\.", "m"),
|
||||
(r"Code style issues found", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"turbo",
|
||||
"build",
|
||||
&[r"^turbo\b", r"^npx\s+turbo\b"],
|
||||
&[
|
||||
(r"^• Packages in scope:", "m"),
|
||||
(r"^Tasks:\s+\d+\s+successful", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"nx",
|
||||
"build",
|
||||
&[r"^nx\b", r"^npx\s+nx\b"],
|
||||
&[(r"^NX\s+", "m"), (r"^> nx run ", "m")],
|
||||
),
|
||||
(
|
||||
"playwright",
|
||||
"test",
|
||||
&[r"^playwright\s+test\b", r"^npx\s+playwright\s+test\b"],
|
||||
&[
|
||||
(r"Running \d+ tests? using \d+ workers?", "i"),
|
||||
(r"^\s+\d+ failed", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"npm-install",
|
||||
"package",
|
||||
&[r"^(?:npm|pnpm|yarn)\s+(?:install|add|update)\b"],
|
||||
&[
|
||||
(r"added \d+ packages", "i"),
|
||||
(r"packages are looking for funding", "i"),
|
||||
(r"audited \d+ packages", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"npm-audit",
|
||||
"package",
|
||||
&[r"^(?:npm|pnpm|yarn)\s+audit\b"],
|
||||
&[
|
||||
(r"found \d+ vulnerabilities", "i"),
|
||||
(r"\b(?:low|moderate|high|critical)\b", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"ruff",
|
||||
"build",
|
||||
&[r"^ruff\b", r"^uv\s+run\s+ruff\b"],
|
||||
&[
|
||||
(r"^[\w./-]+\.py:\d+:\d+:\s+[A-Z]\d+", "m"),
|
||||
(r"Found \d+ errors?\.", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"mypy",
|
||||
"build",
|
||||
&[r"^mypy\b", r"^python\s+-m\s+mypy\b"],
|
||||
&[
|
||||
(r"^[\w./-]+\.py:\d+:\s+error:", "m"),
|
||||
(r"Found \d+ errors? in \d+ files?", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"pip",
|
||||
"package",
|
||||
&[
|
||||
r"^pip\s+(?:install|download|uninstall)\b",
|
||||
r"^python\s+-m\s+pip\b",
|
||||
],
|
||||
&[(r"^Collecting ", "m"), (r"^Successfully installed ", "m")],
|
||||
),
|
||||
(
|
||||
"uv-sync",
|
||||
"package",
|
||||
&[r"^uv\s+sync\b", r"^uv\s+pip\s+install\b"],
|
||||
&[
|
||||
(r"^Resolved \d+ packages?", "m"),
|
||||
(r"^Installed \d+ packages?", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"poetry-install",
|
||||
"package",
|
||||
&[r"^poetry\s+install\b"],
|
||||
&[
|
||||
(r"^Installing dependencies from lock file", "m"),
|
||||
(r"^Package operations:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"golangci-lint",
|
||||
"build",
|
||||
&[r"^golangci-lint\b"],
|
||||
&[(r"^[\w./-]+\.go:\d+:\d+:", "m"), (r"^\d+ issues?:", "m")],
|
||||
),
|
||||
(
|
||||
"bundle-install",
|
||||
"package",
|
||||
&[r"^bundle\s+install\b"],
|
||||
&[
|
||||
(r"^Fetching gem metadata from ", "m"),
|
||||
(r"^Bundle complete!", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"rubocop",
|
||||
"build",
|
||||
&[r"^rubocop\b", r"^bundle\s+exec\s+rubocop\b"],
|
||||
&[
|
||||
(r"^Inspecting \d+ files", "m"),
|
||||
(r"^[\w./-]+\.rb:\d+:\d+:\s+[A-Z]:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"docker-ps",
|
||||
"docker",
|
||||
&[r"^docker\s+ps\b"],
|
||||
&[(r"^CONTAINER ID\s+IMAGE\s+COMMAND", "m")],
|
||||
),
|
||||
(
|
||||
"docker-logs",
|
||||
"docker",
|
||||
&[r"^docker\s+logs\b", r"^docker\s+compose\s+logs\b"],
|
||||
&[
|
||||
(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}", "m"),
|
||||
(r"\b(?:ERROR|WARN|INFO)\b", ""),
|
||||
(r"^Attaching to ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"aws",
|
||||
"cloud",
|
||||
&[r"^aws\b"],
|
||||
&[
|
||||
(r"An error occurred \([A-Za-z0-9]+\) when calling", ""),
|
||||
(r"^(?:upload|download): ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"gcloud",
|
||||
"cloud",
|
||||
&[r"^gcloud\b"],
|
||||
&[(r"^ERROR: \(gcloud\.", "m"), (r"^Updated property \[", "m")],
|
||||
),
|
||||
(
|
||||
"ssh",
|
||||
"cloud",
|
||||
&[r"^ssh\b"],
|
||||
&[
|
||||
(r"Permission denied \(", ""),
|
||||
(r"Host key verification failed", ""),
|
||||
(r"Connection timed out", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"rsync",
|
||||
"cloud",
|
||||
&[r"^rsync\b"],
|
||||
&[
|
||||
(r"^sending incremental file list", "m"),
|
||||
(r"^rsync error:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"curl",
|
||||
"cloud",
|
||||
&[r"^curl\b"],
|
||||
&[(r"curl: \(\d+\)", ""), (r"^HTTP/\d(?:\.\d)? \d{3}", "m")],
|
||||
),
|
||||
(
|
||||
"wget",
|
||||
"cloud",
|
||||
&[r"^wget\b"],
|
||||
&[(r"^--\d{4}-\d{2}-\d{2}", "m"), (r"^ERROR \d{3}:", "m")],
|
||||
),
|
||||
(
|
||||
"json-output",
|
||||
"generic",
|
||||
&[r"^jq\b", r"^cat\s+.*\.json\b"],
|
||||
&[(r"^\s*[\[{][\s\S]*[\]}]\s*$", "")],
|
||||
),
|
||||
(
|
||||
"shell-ls",
|
||||
"shell",
|
||||
&[r"^ls(?:\s+-[A-Za-z]+)?\b"],
|
||||
&[
|
||||
(r"^total \d+", "m"),
|
||||
(r"^\S+\s+\S+\s+\d+\s+\w+\s+\d{1,2}\s+", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"shell-find",
|
||||
"shell",
|
||||
&[r"^find\b"],
|
||||
&[(r"^(?:\.{1,2}|/|[\w.-]+/).+", "m")],
|
||||
),
|
||||
(
|
||||
"shell-grep",
|
||||
"shell",
|
||||
&[r"^(?:grep|rg|ag)\b"],
|
||||
&[
|
||||
(
|
||||
r"^[\w./-]+\.(?:ts|tsx|js|jsx|py|go|rs|java|rb|md|json|ya?ml|txt):\d*:",
|
||||
"m",
|
||||
),
|
||||
(r"^[\w./-]+/[\w./-]+:\d*:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"shell-ps",
|
||||
"shell",
|
||||
&[r"^ps\b"],
|
||||
&[(r"^(?:USER\s+PID|\s*PID\s+)", "m")],
|
||||
),
|
||||
(
|
||||
"shell-df",
|
||||
"shell",
|
||||
&[r"^df\b"],
|
||||
&[(r"^Filesystem\s+.*Use%", "m")],
|
||||
),
|
||||
(
|
||||
"shell-du",
|
||||
"shell",
|
||||
&[r"^du\b"],
|
||||
&[(r"^\d+(?:\.\d+)?[KMGTP]?\s+\S+", "m")],
|
||||
),
|
||||
(
|
||||
"error-stacktrace",
|
||||
"generic",
|
||||
&[],
|
||||
&[
|
||||
(r"Traceback \(most recent call last\):", ""),
|
||||
(r"^\s+at\s+\S+\s+\(.+:\d+:\d+\)", "m"),
|
||||
(r"^panic: ", "m"),
|
||||
(r"^thread '[^']+' panicked at", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"generic-error",
|
||||
"generic",
|
||||
&[],
|
||||
&[
|
||||
(r"Error:", ""),
|
||||
(r"Exception:", ""),
|
||||
(r"Traceback \(most recent call last\):", ""),
|
||||
],
|
||||
),
|
||||
];
|
||||
|
||||
table
|
||||
.iter()
|
||||
.map(|(t, cat, cmds, contents)| Detector {
|
||||
r#type: t,
|
||||
category: cat,
|
||||
command_patterns: cmds.iter().map(|p| rx(p, "i")).collect(),
|
||||
content_patterns: contents.iter().map(|(p, f)| rx(p, f)).collect(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Quote-aware composite-command splitter — port of `lastCommandSegment`.
|
||||
/// Returns the last top-level `&&`/`||`/`;`-separated segment (trimmed).
|
||||
pub fn last_command_segment(command: &str) -> String {
|
||||
|
||||
@@ -0,0 +1,468 @@
|
||||
use super::{rx, Detector, DetectorRow};
|
||||
|
||||
pub(super) fn build_detectors() -> Vec<Detector> {
|
||||
let table: &[DetectorRow] = &[
|
||||
(
|
||||
"git-status",
|
||||
"git",
|
||||
&[r"^git\s+status\b"],
|
||||
&[
|
||||
(r"^On branch ", "m"),
|
||||
(r"^Changes (?:not staged|to be committed)", "m"),
|
||||
(r"^Untracked files:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-branch",
|
||||
"git",
|
||||
&[r"^git\s+branch\b", r"^git\s+checkout\b", r"^git\s+switch\b"],
|
||||
&[
|
||||
(r"^\*\s+\S+", "m"),
|
||||
(r"Switched to (?:a new )?branch", "i"),
|
||||
(r#"Already on ['"][^'"]+['"]"#, "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-diff",
|
||||
"git",
|
||||
&[r"^git\s+diff\b", r"^git\s+show\b"],
|
||||
&[
|
||||
(r"^diff --git ", "m"),
|
||||
(r"^@@\s+-\d+,\d+\s+\+\d+,\d+\s+@@", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"git-log",
|
||||
"git",
|
||||
&[r"^git\s+log\b"],
|
||||
&[(r"^commit [0-9a-f]{7,40}", "m"), (r"^Author: ", "m")],
|
||||
),
|
||||
(
|
||||
"make",
|
||||
"build",
|
||||
&[r"^make\b"],
|
||||
&[
|
||||
(r"^make\[\d+\]: (?:Entering|Leaving) directory", "m"),
|
||||
(r"make: \*\*\* ", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"gradle",
|
||||
"build",
|
||||
&[r"^(?:gradle|gradlew|\./gradlew)\b"],
|
||||
&[
|
||||
(r"^> Task :", "m"),
|
||||
(r"^BUILD (?:SUCCESSFUL|FAILED)\b", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"dotnet",
|
||||
"build",
|
||||
&[
|
||||
r"^dotnet\s+(?:build|test|run|restore|publish|pack|msbuild)\b",
|
||||
r"^dotnet\b",
|
||||
],
|
||||
&[
|
||||
(r"^Build (?:succeeded|FAILED)\b", "m"),
|
||||
(r"\b(?:error|warning) CS\d+\b", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"terraform-plan",
|
||||
"infra",
|
||||
&[r"^terraform\s+plan\b"],
|
||||
&[
|
||||
(r"Terraform will perform the following actions:", ""),
|
||||
(r"Plan: \d+ to add", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"tofu-plan",
|
||||
"infra",
|
||||
&[r"^(?:tofu|opentofu)\s+plan\b"],
|
||||
&[
|
||||
(r"OpenTofu will perform the following actions:", ""),
|
||||
(r"Plan: \d+ to add", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"systemctl-status",
|
||||
"infra",
|
||||
&[r"^systemctl\s+status\b"],
|
||||
&[
|
||||
(r"^\s*Loaded:\s+", "m"),
|
||||
(r"^\s*Active:\s+", "m"),
|
||||
(r"^●\s+\S+\.service", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-vitest",
|
||||
"test",
|
||||
&[r"^vitest\b", r"^npm\s+(?:run\s+)?test:vitest\b"],
|
||||
&[
|
||||
(r"\bvitest\b", "i"),
|
||||
(r"^ ✓ ", "m"),
|
||||
(r"^ ❯ ", "m"),
|
||||
(r"Test Files\s+\d+\s+(?:passed|failed)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-jest",
|
||||
"test",
|
||||
&[r"^jest\b", r"^npm\s+(?:run\s+)?test\b"],
|
||||
&[
|
||||
(r"Test Suites:\s+\d+", "i"),
|
||||
(r"Tests:\s+\d+", "i"),
|
||||
(r"^PASS\s+", "m"),
|
||||
(r"^FAIL\s+", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-pytest",
|
||||
"test",
|
||||
&[r"^pytest\b", r"^python\s+-m\s+pytest\b"],
|
||||
&[
|
||||
(r"=+\s+(?:\d+\s+)?(?:passed|failed|errors?)", "i"),
|
||||
(r"^E\s+", "m"),
|
||||
(r"^FAILED ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-cargo",
|
||||
"test",
|
||||
&[r"^cargo\s+test\b", r"^cargo\s+nextest\b"],
|
||||
&[
|
||||
(r"^running \d+ tests?", "m"),
|
||||
(r"^test\s+[\w:.-]+\s+\.\.\.\s+(?:ok|FAILED|ignored)", "m"),
|
||||
(r"test result:\s+(?:ok|FAILED)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"test-go",
|
||||
"test",
|
||||
&[r"^go\s+test\b"],
|
||||
&[
|
||||
(r"^(?:ok|FAIL)\s+[\w./-]+\s+[\d.]+s", "m"),
|
||||
(r"^--- FAIL: ", "m"),
|
||||
(r"^panic: ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-typescript",
|
||||
"build",
|
||||
&[r"^tsc\b", r"^npm\s+run\s+typecheck\b"],
|
||||
&[(r"TS\d{4}:", ""), (r"error TS\d{4}", "i")],
|
||||
),
|
||||
(
|
||||
"build-eslint",
|
||||
"build",
|
||||
&[r"^eslint\b", r"^npm\s+run\s+lint\b"],
|
||||
&[
|
||||
(r"\s+\d+:\d+\s+(?:error|warning)\s+", ""),
|
||||
(r"✖\s+\d+\s+problems?", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-webpack",
|
||||
"build",
|
||||
&[
|
||||
r"^webpack\b",
|
||||
r"^npx\s+webpack\b",
|
||||
r"^npm\s+run\s+build:webpack\b",
|
||||
],
|
||||
&[
|
||||
(r"webpack\s+\d", "i"),
|
||||
(r"compiled (?:successfully|with \d+ errors?)", "i"),
|
||||
(r"asset .+\.js", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"build-vite",
|
||||
"build",
|
||||
&[
|
||||
r"^vite\s+build\b",
|
||||
r"^npm\s+run\s+build\b",
|
||||
r"^pnpm\s+build\b",
|
||||
],
|
||||
&[
|
||||
(r"vite v[\d.]+", "i"),
|
||||
(r"✓ built in", "i"),
|
||||
(r"transforming \(\d+\)", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"biome",
|
||||
"build",
|
||||
&[r"^biome\b", r"^npx\s+biome\b"],
|
||||
&[
|
||||
(r"lint/[A-Za-z0-9/.-]+", ""),
|
||||
(r"Checked \d+ files? in", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"prettier",
|
||||
"build",
|
||||
&[r"^prettier\b", r"^npx\s+prettier\b"],
|
||||
&[
|
||||
(r"^Checking formatting\.\.\.", "m"),
|
||||
(r"Code style issues found", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"turbo",
|
||||
"build",
|
||||
&[r"^turbo\b", r"^npx\s+turbo\b"],
|
||||
&[
|
||||
(r"^• Packages in scope:", "m"),
|
||||
(r"^Tasks:\s+\d+\s+successful", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"nx",
|
||||
"build",
|
||||
&[r"^nx\b", r"^npx\s+nx\b"],
|
||||
&[(r"^NX\s+", "m"), (r"^> nx run ", "m")],
|
||||
),
|
||||
(
|
||||
"playwright",
|
||||
"test",
|
||||
&[r"^playwright\s+test\b", r"^npx\s+playwright\s+test\b"],
|
||||
&[
|
||||
(r"Running \d+ tests? using \d+ workers?", "i"),
|
||||
(r"^\s+\d+ failed", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"npm-install",
|
||||
"package",
|
||||
&[r"^(?:npm|pnpm|yarn)\s+(?:install|add|update)\b"],
|
||||
&[
|
||||
(r"added \d+ packages", "i"),
|
||||
(r"packages are looking for funding", "i"),
|
||||
(r"audited \d+ packages", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"npm-audit",
|
||||
"package",
|
||||
&[r"^(?:npm|pnpm|yarn)\s+audit\b"],
|
||||
&[
|
||||
(r"found \d+ vulnerabilities", "i"),
|
||||
(r"\b(?:low|moderate|high|critical)\b", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"ruff",
|
||||
"build",
|
||||
&[r"^ruff\b", r"^uv\s+run\s+ruff\b"],
|
||||
&[
|
||||
(r"^[\w./-]+\.py:\d+:\d+:\s+[A-Z]\d+", "m"),
|
||||
(r"Found \d+ errors?\.", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"mypy",
|
||||
"build",
|
||||
&[r"^mypy\b", r"^python\s+-m\s+mypy\b"],
|
||||
&[
|
||||
(r"^[\w./-]+\.py:\d+:\s+error:", "m"),
|
||||
(r"Found \d+ errors? in \d+ files?", "i"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"pip",
|
||||
"package",
|
||||
&[
|
||||
r"^pip\s+(?:install|download|uninstall)\b",
|
||||
r"^python\s+-m\s+pip\b",
|
||||
],
|
||||
&[(r"^Collecting ", "m"), (r"^Successfully installed ", "m")],
|
||||
),
|
||||
(
|
||||
"uv-sync",
|
||||
"package",
|
||||
&[r"^uv\s+sync\b", r"^uv\s+pip\s+install\b"],
|
||||
&[
|
||||
(r"^Resolved \d+ packages?", "m"),
|
||||
(r"^Installed \d+ packages?", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"poetry-install",
|
||||
"package",
|
||||
&[r"^poetry\s+install\b"],
|
||||
&[
|
||||
(r"^Installing dependencies from lock file", "m"),
|
||||
(r"^Package operations:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"golangci-lint",
|
||||
"build",
|
||||
&[r"^golangci-lint\b"],
|
||||
&[(r"^[\w./-]+\.go:\d+:\d+:", "m"), (r"^\d+ issues?:", "m")],
|
||||
),
|
||||
(
|
||||
"bundle-install",
|
||||
"package",
|
||||
&[r"^bundle\s+install\b"],
|
||||
&[
|
||||
(r"^Fetching gem metadata from ", "m"),
|
||||
(r"^Bundle complete!", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"rubocop",
|
||||
"build",
|
||||
&[r"^rubocop\b", r"^bundle\s+exec\s+rubocop\b"],
|
||||
&[
|
||||
(r"^Inspecting \d+ files", "m"),
|
||||
(r"^[\w./-]+\.rb:\d+:\d+:\s+[A-Z]:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"docker-ps",
|
||||
"docker",
|
||||
&[r"^docker\s+ps\b"],
|
||||
&[(r"^CONTAINER ID\s+IMAGE\s+COMMAND", "m")],
|
||||
),
|
||||
(
|
||||
"docker-logs",
|
||||
"docker",
|
||||
&[r"^docker\s+logs\b", r"^docker\s+compose\s+logs\b"],
|
||||
&[
|
||||
(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}", "m"),
|
||||
(r"\b(?:ERROR|WARN|INFO)\b", ""),
|
||||
(r"^Attaching to ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"aws",
|
||||
"cloud",
|
||||
&[r"^aws\b"],
|
||||
&[
|
||||
(r"An error occurred \([A-Za-z0-9]+\) when calling", ""),
|
||||
(r"^(?:upload|download): ", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"gcloud",
|
||||
"cloud",
|
||||
&[r"^gcloud\b"],
|
||||
&[(r"^ERROR: \(gcloud\.", "m"), (r"^Updated property \[", "m")],
|
||||
),
|
||||
(
|
||||
"ssh",
|
||||
"cloud",
|
||||
&[r"^ssh\b"],
|
||||
&[
|
||||
(r"Permission denied \(", ""),
|
||||
(r"Host key verification failed", ""),
|
||||
(r"Connection timed out", ""),
|
||||
],
|
||||
),
|
||||
(
|
||||
"rsync",
|
||||
"cloud",
|
||||
&[r"^rsync\b"],
|
||||
&[
|
||||
(r"^sending incremental file list", "m"),
|
||||
(r"^rsync error:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"curl",
|
||||
"cloud",
|
||||
&[r"^curl\b"],
|
||||
&[(r"curl: \(\d+\)", ""), (r"^HTTP/\d(?:\.\d)? \d{3}", "m")],
|
||||
),
|
||||
(
|
||||
"wget",
|
||||
"cloud",
|
||||
&[r"^wget\b"],
|
||||
&[(r"^--\d{4}-\d{2}-\d{2}", "m"), (r"^ERROR \d{3}:", "m")],
|
||||
),
|
||||
(
|
||||
"json-output",
|
||||
"generic",
|
||||
&[r"^jq\b", r"^cat\s+.*\.json\b"],
|
||||
&[(r"^\s*[\[{][\s\S]*[\]}]\s*$", "")],
|
||||
),
|
||||
(
|
||||
"shell-ls",
|
||||
"shell",
|
||||
&[r"^ls(?:\s+-[A-Za-z]+)?\b"],
|
||||
&[
|
||||
(r"^total \d+", "m"),
|
||||
(r"^\S+\s+\S+\s+\d+\s+\w+\s+\d{1,2}\s+", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"shell-find",
|
||||
"shell",
|
||||
&[r"^find\b"],
|
||||
&[(r"^(?:\.{1,2}|/|[\w.-]+/).+", "m")],
|
||||
),
|
||||
(
|
||||
"shell-grep",
|
||||
"shell",
|
||||
&[r"^(?:grep|rg|ag)\b"],
|
||||
&[
|
||||
(
|
||||
r"^[\w./-]+\.(?:ts|tsx|js|jsx|py|go|rs|java|rb|md|json|ya?ml|txt):\d*:",
|
||||
"m",
|
||||
),
|
||||
(r"^[\w./-]+/[\w./-]+:\d*:", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"shell-ps",
|
||||
"shell",
|
||||
&[r"^ps\b"],
|
||||
&[(r"^(?:USER\s+PID|\s*PID\s+)", "m")],
|
||||
),
|
||||
(
|
||||
"shell-df",
|
||||
"shell",
|
||||
&[r"^df\b"],
|
||||
&[(r"^Filesystem\s+.*Use%", "m")],
|
||||
),
|
||||
(
|
||||
"shell-du",
|
||||
"shell",
|
||||
&[r"^du\b"],
|
||||
&[(r"^\d+(?:\.\d+)?[KMGTP]?\s+\S+", "m")],
|
||||
),
|
||||
(
|
||||
"error-stacktrace",
|
||||
"generic",
|
||||
&[],
|
||||
&[
|
||||
(r"Traceback \(most recent call last\):", ""),
|
||||
(r"^\s+at\s+\S+\s+\(.+:\d+:\d+\)", "m"),
|
||||
(r"^panic: ", "m"),
|
||||
(r"^thread '[^']+' panicked at", "m"),
|
||||
],
|
||||
),
|
||||
(
|
||||
"generic-error",
|
||||
"generic",
|
||||
&[],
|
||||
&[
|
||||
(r"Error:", ""),
|
||||
(r"Exception:", ""),
|
||||
(r"Traceback \(most recent call last\):", ""),
|
||||
],
|
||||
),
|
||||
];
|
||||
|
||||
table
|
||||
.iter()
|
||||
.map(|(t, cat, cmds, contents)| Detector {
|
||||
r#type: t,
|
||||
category: cat,
|
||||
command_patterns: cmds.iter().map(|p| rx(p, "i")).collect(),
|
||||
content_patterns: contents.iter().map(|(p, f)| rx(p, f)).collect(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -257,13 +257,11 @@ impl RtkFilter {
|
||||
.build()
|
||||
.ok()?;
|
||||
let unless = match r.unless {
|
||||
Some(u) => {
|
||||
RegexBuilder::new(&u)
|
||||
.case_insensitive(true)
|
||||
.multi_line(true)
|
||||
.build()
|
||||
.ok()
|
||||
}
|
||||
Some(u) => RegexBuilder::new(&u)
|
||||
.case_insensitive(true)
|
||||
.multi_line(true)
|
||||
.build()
|
||||
.ok(),
|
||||
None => None,
|
||||
};
|
||||
Some(MatchOutputRule {
|
||||
|
||||
@@ -31,11 +31,14 @@ fn build_lookup(messages: &[Value]) -> ToolLookup {
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let command = part.get("input").and_then(command_from_input);
|
||||
lookup.insert(id.to_string(), ToolMeta {
|
||||
name,
|
||||
command,
|
||||
arguments_str: None,
|
||||
});
|
||||
lookup.insert(
|
||||
id.to_string(),
|
||||
ToolMeta {
|
||||
name,
|
||||
command,
|
||||
arguments_str: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
lookup
|
||||
|
||||
@@ -40,11 +40,14 @@ fn build_lookup(messages: &[Value]) -> ToolLookup {
|
||||
// strings). For OpenAI the parse is deferred — build_lookup only
|
||||
// stores the raw string.
|
||||
let command = None;
|
||||
lookup.insert(id.to_string(), ToolMeta {
|
||||
name,
|
||||
command,
|
||||
arguments_str,
|
||||
});
|
||||
lookup.insert(
|
||||
id.to_string(),
|
||||
ToolMeta {
|
||||
name,
|
||||
command,
|
||||
arguments_str,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
lookup
|
||||
|
||||
@@ -1,682 +1,20 @@
|
||||
// Anthropic Messages API <-> Gemini generateContent API message mapping.
|
||||
//
|
||||
// Pure translation functions, no IO. Converts Anthropic request types into
|
||||
// Gemini request types and Gemini response types back into Anthropic responses.
|
||||
// Pure translation functions, no IO. Split by direction:
|
||||
// - `request` Anthropic -> Gemini request
|
||||
// - `response` Gemini -> Anthropic response
|
||||
// - `reverse` Gemini -> Anthropic request & Anthropic -> Gemini response
|
||||
// - `helpers` shared utilities (role merging, tool-id mapping)
|
||||
|
||||
use std::collections::HashMap;
|
||||
mod helpers;
|
||||
mod request;
|
||||
mod response;
|
||||
mod reverse;
|
||||
|
||||
use crate::anthropic::messages as anthropic;
|
||||
use crate::gemini::request as gemini;
|
||||
use crate::gemini::response as gemini_resp;
|
||||
use crate::mapping::tools_map::sanitize_schema_for_gemini;
|
||||
use crate::util::ids::{generate_message_id, generate_tool_use_id};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Request direction: Anthropic -> Gemini
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Compute degradation warnings for a Gemini-bound request.
|
||||
///
|
||||
/// Call before translating to surface features that are silently dropped during
|
||||
/// Anthropic → Gemini translation. Emit the result as an `x-anyllm-degradation`
|
||||
/// response header so clients can detect lossy translations.
|
||||
pub fn compute_gemini_request_warnings(
|
||||
req: &anthropic::MessageCreateRequest,
|
||||
) -> crate::mapping::warnings::TranslationWarnings {
|
||||
use crate::mapping::warnings::TranslationWarnings;
|
||||
let mut w = TranslationWarnings::default();
|
||||
|
||||
// Single pass: collect all per-block warning flags at once.
|
||||
let mut has_thinking = false;
|
||||
let mut has_document = false;
|
||||
let mut has_url_image = false;
|
||||
for msg in &req.messages {
|
||||
if let anthropic::Content::Blocks(blocks) = &msg.content {
|
||||
for b in blocks {
|
||||
match b {
|
||||
// Thinking/RedactedThinking: no Gemini Content equivalent.
|
||||
anthropic::ContentBlock::Thinking { .. }
|
||||
| anthropic::ContentBlock::RedactedThinking { .. } => has_thinking = true,
|
||||
// Document blocks have no Gemini equivalent.
|
||||
anthropic::ContentBlock::Document { .. } => has_document = true,
|
||||
// URL-type images: Gemini only accepts inline base64 data.
|
||||
anthropic::ContentBlock::Image { source } if source.source_type != "base64" => {
|
||||
has_url_image = true
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if has_thinking {
|
||||
w.add("thinking_blocks");
|
||||
}
|
||||
if has_document {
|
||||
w.add("document_blocks");
|
||||
}
|
||||
if has_url_image {
|
||||
w.add("url_images");
|
||||
}
|
||||
|
||||
// cache_control on system blocks is dropped; Gemini has no prompt-caching API.
|
||||
if let Some(anthropic::System::Blocks(blocks)) = &req.system {
|
||||
if blocks.iter().any(|b| b.cache_control.is_some()) {
|
||||
w.add("cache_control");
|
||||
}
|
||||
}
|
||||
|
||||
w
|
||||
}
|
||||
|
||||
/// Convert an Anthropic `MessageCreateRequest` into a Gemini `GenerateContentRequest`.
|
||||
///
|
||||
/// Maps `thinking_config` to `generationConfig.thinkingConfig` for Gemini 2.5
|
||||
/// thinking models. Drops unsupported features (thinking content blocks in prior
|
||||
/// messages, document blocks, cache_control) and merges consecutive same-role
|
||||
/// messages to satisfy Gemini's strict alternation requirement.
|
||||
pub fn anthropic_to_gemini_request(
|
||||
req: &anthropic::MessageCreateRequest,
|
||||
) -> gemini::GenerateContentRequest {
|
||||
let tool_id_map = build_tool_id_map(&req.messages);
|
||||
|
||||
// System instruction
|
||||
let system_instruction = req.system.as_ref().map(|sys| {
|
||||
let text = match sys {
|
||||
anthropic::System::Text(s) => s.clone(),
|
||||
anthropic::System::Blocks(blocks) => blocks
|
||||
.iter()
|
||||
.map(|b| b.text.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
};
|
||||
gemini::Content {
|
||||
role: None,
|
||||
parts: vec![gemini::Part::text(text)],
|
||||
}
|
||||
});
|
||||
|
||||
// Convert messages
|
||||
let mut contents: Vec<gemini::Content> = Vec::new();
|
||||
for msg in &req.messages {
|
||||
let role = match msg.role {
|
||||
anthropic::Role::User => "user",
|
||||
anthropic::Role::Assistant => "model",
|
||||
};
|
||||
let parts = content_blocks_to_parts(&msg.content, &tool_id_map);
|
||||
if !parts.is_empty() {
|
||||
contents.push(gemini::Content {
|
||||
role: Some(role.to_string()),
|
||||
parts,
|
||||
});
|
||||
}
|
||||
}
|
||||
contents = merge_consecutive_roles(contents);
|
||||
|
||||
// Tools
|
||||
let tools = req.tools.as_ref().map(|tools| {
|
||||
vec![gemini::Tool {
|
||||
function_declarations: tools
|
||||
.iter()
|
||||
.map(|t| gemini::FunctionDeclaration {
|
||||
name: t.name.clone(),
|
||||
description: t.description.clone(),
|
||||
parameters: Some(sanitize_schema_for_gemini(t.input_schema.clone())),
|
||||
})
|
||||
.collect(),
|
||||
}]
|
||||
});
|
||||
|
||||
// Tool config
|
||||
let tool_config = req.tool_choice.as_ref().map(|tc| {
|
||||
if matches!(
|
||||
tc,
|
||||
anthropic::ToolChoice::Auto {
|
||||
disable_parallel_tool_use: Some(true)
|
||||
} | anthropic::ToolChoice::Any {
|
||||
disable_parallel_tool_use: Some(true)
|
||||
}
|
||||
) {
|
||||
tracing::warn!(
|
||||
"disable_parallel_tool_use=true is not supported by Gemini; \
|
||||
parallel tool calls may still occur"
|
||||
);
|
||||
}
|
||||
let (mode, allowed) = match tc {
|
||||
anthropic::ToolChoice::Auto { .. } => ("AUTO", None),
|
||||
anthropic::ToolChoice::Any { .. } => ("ANY", None),
|
||||
anthropic::ToolChoice::None => ("NONE", None),
|
||||
// Gemini ANY + allowedFunctionNames restricts to a specific tool.
|
||||
anthropic::ToolChoice::Tool { name, .. } => ("ANY", Some(vec![name.clone()])),
|
||||
};
|
||||
gemini::ToolConfig {
|
||||
function_calling_config: gemini::FunctionCallingConfig {
|
||||
mode: mode.to_string(),
|
||||
allowed_function_names: allowed,
|
||||
},
|
||||
}
|
||||
});
|
||||
|
||||
// Generation config
|
||||
let generation_config = {
|
||||
let thinking_config =
|
||||
if let Some(anthropic::ThinkingConfig::Enabled { budget_tokens }) = &req.thinking {
|
||||
Some(gemini::ThinkingConfig {
|
||||
thinking_budget: *budget_tokens,
|
||||
include_thoughts: Some(true),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let gc = gemini::GenerationConfig {
|
||||
max_output_tokens: Some(req.max_tokens),
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
top_k: req.top_k,
|
||||
stop_sequences: req.stop_sequences.clone(),
|
||||
thinking_config,
|
||||
..Default::default()
|
||||
};
|
||||
Some(gc)
|
||||
};
|
||||
|
||||
gemini::GenerateContentRequest {
|
||||
contents,
|
||||
system_instruction,
|
||||
generation_config,
|
||||
tools,
|
||||
tool_config,
|
||||
safety_settings: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a map from Anthropic tool_use IDs to tool names.
|
||||
///
|
||||
/// Scans all messages for `ToolUse` blocks so that `ToolResult` translation can
|
||||
/// look up the function name Gemini expects.
|
||||
pub fn build_tool_id_map(messages: &[anthropic::InputMessage]) -> HashMap<String, String> {
|
||||
let mut map = HashMap::new();
|
||||
for msg in messages {
|
||||
let blocks = match &msg.content {
|
||||
anthropic::Content::Text(_) => continue,
|
||||
anthropic::Content::Blocks(b) => b,
|
||||
};
|
||||
for block in blocks {
|
||||
if let anthropic::ContentBlock::ToolUse { id, name, .. } = block {
|
||||
map.insert(id.clone(), name.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
map
|
||||
}
|
||||
|
||||
/// Merge consecutive same-role `Content` entries by concatenating their parts.
|
||||
///
|
||||
/// Gemini requires strict user/model role alternation. When the Anthropic
|
||||
/// conversation has two consecutive user (or model) turns, this merges them
|
||||
/// into a single turn.
|
||||
pub fn merge_consecutive_roles(contents: Vec<gemini::Content>) -> Vec<gemini::Content> {
|
||||
let mut merged: Vec<gemini::Content> = Vec::with_capacity(contents.len());
|
||||
for c in contents {
|
||||
if let Some(last) = merged.last_mut() {
|
||||
if last.role == c.role {
|
||||
last.parts.extend(c.parts);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
merged.push(c);
|
||||
}
|
||||
|
||||
// Gemini requires the first content turn to have role "user". An Anthropic
|
||||
// client may legally send an assistant-first conversation (for few-shot
|
||||
// prompting). Prepend a dummy user turn so Gemini does not return a 400.
|
||||
if merged.first().and_then(|c| c.role.as_deref()) == Some("model") {
|
||||
merged.insert(
|
||||
0,
|
||||
gemini::Content {
|
||||
role: Some("user".to_string()),
|
||||
parts: vec![gemini::Part::text(String::new())],
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
merged
|
||||
}
|
||||
|
||||
/// Convert Anthropic message content into a vec of Gemini Parts.
|
||||
fn content_blocks_to_parts(
|
||||
content: &anthropic::Content,
|
||||
tool_id_map: &HashMap<String, String>,
|
||||
) -> Vec<gemini::Part> {
|
||||
match content {
|
||||
anthropic::Content::Text(s) => vec![gemini::Part::text(s.clone())],
|
||||
anthropic::Content::Blocks(blocks) => blocks
|
||||
.iter()
|
||||
.filter_map(|block| content_block_to_part(block, tool_id_map))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single Anthropic ContentBlock to a Gemini Part, or None if dropped.
|
||||
fn content_block_to_part(
|
||||
block: &anthropic::ContentBlock,
|
||||
tool_id_map: &HashMap<String, String>,
|
||||
) -> Option<gemini::Part> {
|
||||
match block {
|
||||
anthropic::ContentBlock::Text { text } => Some(gemini::Part::text(text.clone())),
|
||||
|
||||
anthropic::ContentBlock::Image { source } => {
|
||||
// Gemini only supports inline base64 data, not URLs.
|
||||
if source.source_type == "base64" {
|
||||
let mime = source
|
||||
.media_type
|
||||
.clone()
|
||||
.unwrap_or_else(|| "image/png".into());
|
||||
let data = source.data.clone().unwrap_or_default();
|
||||
Some(gemini::Part::inline_data(mime, data))
|
||||
} else {
|
||||
// URL-type images cannot be sent as inline_data; drop.
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
anthropic::ContentBlock::ToolUse { name, input, .. } => {
|
||||
// Strip the Anthropic tool_use id; Gemini uses name-based correlation.
|
||||
Some(gemini::Part::function_call(name.clone(), input.clone()))
|
||||
}
|
||||
|
||||
anthropic::ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let Some(name) = tool_id_map.get(tool_use_id).cloned() else {
|
||||
// Gemini requires the function name to match a declared FunctionDeclaration.
|
||||
// Emitting an unknown name causes a 400; drop the result instead.
|
||||
tracing::warn!(
|
||||
tool_use_id,
|
||||
"dropping ToolResult: tool_use_id not found in tool_id_map"
|
||||
);
|
||||
return None;
|
||||
};
|
||||
|
||||
let response_value = tool_result_to_json(content, *is_error);
|
||||
Some(gemini::Part::function_response(name, response_value))
|
||||
}
|
||||
|
||||
// Thinking, RedactedThinking, Document: not supported by Gemini, drop.
|
||||
anthropic::ContentBlock::Thinking { .. }
|
||||
| anthropic::ContentBlock::RedactedThinking { .. }
|
||||
| anthropic::ContentBlock::Document { .. } => None,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Anthropic ToolResult content into a JSON value for Gemini FunctionResponse.
|
||||
fn tool_result_to_json(
|
||||
content: &Option<anthropic::ToolResultContent>,
|
||||
is_error: Option<bool>,
|
||||
) -> serde_json::Value {
|
||||
let text = match content {
|
||||
Some(anthropic::ToolResultContent::Text(s)) => s.clone(),
|
||||
Some(anthropic::ToolResultContent::Blocks(blocks)) => {
|
||||
// Concatenate text blocks; other block types (e.g., images) cannot be
|
||||
// represented in Gemini FunctionResponse and are replaced with a placeholder.
|
||||
blocks
|
||||
.iter()
|
||||
.map(|b| match b {
|
||||
anthropic::ContentBlock::Text { text } => text.clone(),
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
"tool_result contains non-text block; \
|
||||
replacing with \"[non-text]\" placeholder for Gemini"
|
||||
);
|
||||
"[non-text]".into()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
if is_error == Some(true) {
|
||||
serde_json::json!({ "error": text })
|
||||
} else {
|
||||
serde_json::json!({ "result": text })
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Response direction: Gemini -> Anthropic
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Convert a Gemini `GenerateContentResponse` into an Anthropic `MessageResponse`.
|
||||
///
|
||||
/// Uses only the first candidate. Synthesizes Anthropic-format tool IDs for any
|
||||
/// function calls.
|
||||
pub fn gemini_to_anthropic_response(
|
||||
resp: &gemini_resp::GenerateContentResponse,
|
||||
model: &str,
|
||||
) -> anthropic::MessageResponse {
|
||||
let candidate = resp.candidates.first();
|
||||
|
||||
let content = candidate
|
||||
.map(|c| {
|
||||
c.content
|
||||
.parts
|
||||
.iter()
|
||||
.filter_map(gemini_part_to_content_block)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let has_function_call = content
|
||||
.iter()
|
||||
.any(|b| matches!(b, anthropic::ContentBlock::ToolUse { .. }));
|
||||
|
||||
let stop_reason = candidate
|
||||
.and_then(|c| c.finish_reason.as_ref())
|
||||
.map(|fr| match fr {
|
||||
gemini_resp::FinishReason::STOP if has_function_call => anthropic::StopReason::ToolUse,
|
||||
gemini_resp::FinishReason::STOP => anthropic::StopReason::EndTurn,
|
||||
gemini_resp::FinishReason::MAX_TOKENS => anthropic::StopReason::MaxTokens,
|
||||
// SAFETY, RECITATION, LANGUAGE, OTHER, Unknown all map to EndTurn.
|
||||
_ => anthropic::StopReason::EndTurn,
|
||||
})
|
||||
// No finish_reason at all (e.g. empty candidates) -> EndTurn.
|
||||
.or(if candidate.is_some() {
|
||||
Some(anthropic::StopReason::EndTurn)
|
||||
} else {
|
||||
None
|
||||
});
|
||||
|
||||
let usage = resp
|
||||
.usage_metadata
|
||||
.as_ref()
|
||||
.map(|u| anthropic::Usage {
|
||||
input_tokens: u.prompt_token_count,
|
||||
output_tokens: u.candidates_token_count,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
anthropic::MessageResponse {
|
||||
id: generate_message_id(),
|
||||
response_type: "message".into(),
|
||||
role: anthropic::Role::Assistant,
|
||||
content,
|
||||
model: model.to_string(),
|
||||
stop_reason,
|
||||
stop_sequence: None,
|
||||
usage,
|
||||
created: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single Gemini Part to an Anthropic ContentBlock, or None if not mappable.
|
||||
fn gemini_part_to_content_block(part: &gemini::Part) -> Option<anthropic::ContentBlock> {
|
||||
// Thought parts from thinking models map to Anthropic thinking blocks.
|
||||
if part.thought == Some(true) {
|
||||
return part
|
||||
.text
|
||||
.as_ref()
|
||||
.map(|text| anthropic::ContentBlock::Thinking {
|
||||
thinking: text.clone(),
|
||||
signature: None,
|
||||
});
|
||||
}
|
||||
if let Some(text) = &part.text {
|
||||
return Some(anthropic::ContentBlock::Text { text: text.clone() });
|
||||
}
|
||||
if let Some(fc) = &part.function_call {
|
||||
return Some(anthropic::ContentBlock::ToolUse {
|
||||
id: generate_tool_use_id(),
|
||||
name: fc.name.clone(),
|
||||
input: fc.args.clone(),
|
||||
});
|
||||
}
|
||||
// inline_data, file_data, function_response: not expected in model output,
|
||||
// or have no Anthropic equivalent. Drop.
|
||||
None
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Input direction: Gemini CLI -> Anthropic (for accepting Gemini-format input)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Convert a Gemini CLI `GenerateContentRequest` into an Anthropic `MessageCreateRequest`.
|
||||
///
|
||||
/// `model` is the model name extracted from the URL path (e.g. `gemini-2.5-pro` from
|
||||
/// `POST /v1beta/models/gemini-2.5-pro:generateContent`). All generation config fields
|
||||
/// map directly; unsupported Gemini features (safety settings, response_schema) are dropped.
|
||||
pub fn gemini_to_anthropic_request(
|
||||
req: &gemini::GenerateContentRequest,
|
||||
model: &str,
|
||||
) -> anthropic::MessageCreateRequest {
|
||||
// Build name->id map so that function_response parts reference the same
|
||||
// synthetic tool_use_id as the corresponding function_call part.
|
||||
let mut name_to_id: HashMap<String, String> = HashMap::new();
|
||||
for content in &req.contents {
|
||||
for part in &content.parts {
|
||||
if let Some(ref fc) = part.function_call {
|
||||
name_to_id
|
||||
.entry(fc.name.clone())
|
||||
.or_insert_with(generate_tool_use_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// System instruction -> system
|
||||
let system = req.system_instruction.as_ref().map(|si| {
|
||||
let text = si
|
||||
.parts
|
||||
.iter()
|
||||
.filter_map(|p| p.text.as_deref())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
anthropic::System::Text(text)
|
||||
});
|
||||
|
||||
// Contents -> messages
|
||||
let messages: Vec<anthropic::InputMessage> = req
|
||||
.contents
|
||||
.iter()
|
||||
.filter_map(|c| gemini_content_to_input_message(c, &name_to_id))
|
||||
.collect();
|
||||
|
||||
// Generation config
|
||||
let gc = req.generation_config.as_ref();
|
||||
let max_tokens = gc.and_then(|g| g.max_output_tokens).unwrap_or(8192);
|
||||
let temperature = gc.and_then(|g| g.temperature);
|
||||
let top_p = gc.and_then(|g| g.top_p);
|
||||
let top_k = gc.and_then(|g| g.top_k);
|
||||
let stop_sequences = gc
|
||||
.and_then(|g| g.stop_sequences.clone())
|
||||
.filter(|v| !v.is_empty());
|
||||
|
||||
// Tools
|
||||
let tools = req.tools.as_ref().map(|ts| {
|
||||
ts.iter()
|
||||
.flat_map(|t| t.function_declarations.iter())
|
||||
.map(|fd| anthropic::Tool {
|
||||
name: fd.name.clone(),
|
||||
description: fd.description.clone(),
|
||||
input_schema: fd
|
||||
.parameters
|
||||
.clone()
|
||||
.unwrap_or_else(|| serde_json::json!({"type": "object"})),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
|
||||
// Tool choice
|
||||
let tool_choice = req.tool_config.as_ref().map(|tc| {
|
||||
match tc.function_calling_config.mode.as_str() {
|
||||
"NONE" => anthropic::ToolChoice::None,
|
||||
"ANY" => match tc.function_calling_config.allowed_function_names.as_deref() {
|
||||
Some([name]) => anthropic::ToolChoice::Tool { name: name.clone() },
|
||||
_ => anthropic::ToolChoice::Any {
|
||||
disable_parallel_tool_use: None,
|
||||
},
|
||||
},
|
||||
// AUTO or anything else
|
||||
_ => anthropic::ToolChoice::Auto {
|
||||
disable_parallel_tool_use: None,
|
||||
},
|
||||
}
|
||||
});
|
||||
|
||||
anthropic::MessageCreateRequest {
|
||||
model: model.to_string(),
|
||||
max_tokens,
|
||||
messages,
|
||||
system,
|
||||
temperature,
|
||||
top_p,
|
||||
top_k,
|
||||
stop_sequences,
|
||||
tools,
|
||||
tool_choice,
|
||||
metadata: None,
|
||||
thinking: None,
|
||||
stream: None,
|
||||
extra: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single Gemini `Content` turn into an Anthropic `InputMessage`.
|
||||
/// Returns `None` if the turn produces no content blocks (e.g. all parts were dropped).
|
||||
fn gemini_content_to_input_message(
|
||||
content: &gemini::Content,
|
||||
name_to_id: &HashMap<String, String>,
|
||||
) -> Option<anthropic::InputMessage> {
|
||||
let role = match content.role.as_deref() {
|
||||
Some("model") => anthropic::Role::Assistant,
|
||||
// "user", None, or anything unrecognised -> user.
|
||||
_ => anthropic::Role::User,
|
||||
};
|
||||
let blocks: Vec<anthropic::ContentBlock> = content
|
||||
.parts
|
||||
.iter()
|
||||
.filter_map(|p| gemini_input_part_to_block(p, name_to_id))
|
||||
.collect();
|
||||
if blocks.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(anthropic::InputMessage {
|
||||
role,
|
||||
content: anthropic::Content::Blocks(blocks),
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert a single Gemini `Part` from a user/model message into an Anthropic `ContentBlock`.
|
||||
fn gemini_input_part_to_block(
|
||||
part: &gemini::Part,
|
||||
name_to_id: &HashMap<String, String>,
|
||||
) -> Option<anthropic::ContentBlock> {
|
||||
if let Some(ref text) = part.text {
|
||||
return Some(anthropic::ContentBlock::Text { text: text.clone() });
|
||||
}
|
||||
if let Some(ref fc) = part.function_call {
|
||||
let id = name_to_id
|
||||
.get(&fc.name)
|
||||
.cloned()
|
||||
.unwrap_or_else(generate_tool_use_id);
|
||||
return Some(anthropic::ContentBlock::ToolUse {
|
||||
id,
|
||||
name: fc.name.clone(),
|
||||
input: fc.args.clone(),
|
||||
});
|
||||
}
|
||||
if let Some(ref fr) = part.function_response {
|
||||
let tool_use_id = name_to_id
|
||||
.get(&fr.name)
|
||||
.cloned()
|
||||
.unwrap_or_else(generate_tool_use_id);
|
||||
return Some(anthropic::ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content: Some(anthropic::ToolResultContent::Text(
|
||||
serde_json::to_string(&fr.response).unwrap_or_default(),
|
||||
)),
|
||||
is_error: None,
|
||||
});
|
||||
}
|
||||
if let Some(ref data) = part.inline_data {
|
||||
if data.mime_type.starts_with("image/") {
|
||||
return Some(anthropic::ContentBlock::Image {
|
||||
source: anthropic::ImageSource {
|
||||
source_type: "base64".to_string(),
|
||||
media_type: Some(data.mime_type.clone()),
|
||||
data: Some(data.data.clone()),
|
||||
url: None,
|
||||
},
|
||||
});
|
||||
}
|
||||
// Non-image inline data (audio, video, etc.) has no Anthropic equivalent — drop.
|
||||
tracing::warn!(
|
||||
mime_type = %data.mime_type,
|
||||
"dropping inline_data part with non-image mime type (no Anthropic equivalent)"
|
||||
);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Convert an Anthropic `MessageResponse` into a Gemini `GenerateContentResponse`.
|
||||
///
|
||||
/// Used when the proxy accepts Gemini CLI input and must return Gemini-format output.
|
||||
/// This is the inverse of `gemini_to_anthropic_response`.
|
||||
pub fn anthropic_to_gemini_response(
|
||||
resp: &anthropic::MessageResponse,
|
||||
) -> gemini_resp::GenerateContentResponse {
|
||||
let parts: Vec<gemini::Part> = resp
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
anthropic::ContentBlock::Text { text } => Some(gemini::Part::text(text.clone())),
|
||||
anthropic::ContentBlock::ToolUse { name, input, .. } => {
|
||||
Some(gemini::Part::function_call(name.clone(), input.clone()))
|
||||
}
|
||||
anthropic::ContentBlock::Thinking { thinking, .. } => Some(gemini::Part {
|
||||
thought: Some(true),
|
||||
text: Some(thinking.clone()),
|
||||
..Default::default()
|
||||
}),
|
||||
// ToolResult, Image, Document, RedactedThinking: not expected in model output.
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let finish_reason = resp.stop_reason.as_ref().map(|sr| match sr {
|
||||
anthropic::StopReason::EndTurn | anthropic::StopReason::ToolUse => {
|
||||
gemini_resp::FinishReason::STOP
|
||||
}
|
||||
anthropic::StopReason::MaxTokens => gemini_resp::FinishReason::MAX_TOKENS,
|
||||
anthropic::StopReason::StopSequence => gemini_resp::FinishReason::STOP,
|
||||
_ => gemini_resp::FinishReason::STOP,
|
||||
});
|
||||
|
||||
let candidate = gemini_resp::Candidate {
|
||||
content: gemini::Content {
|
||||
role: Some("model".to_string()),
|
||||
parts,
|
||||
},
|
||||
finish_reason,
|
||||
safety_ratings: None,
|
||||
};
|
||||
|
||||
gemini_resp::GenerateContentResponse {
|
||||
candidates: vec![candidate],
|
||||
usage_metadata: Some(gemini_resp::UsageMetadata {
|
||||
prompt_token_count: resp.usage.input_tokens,
|
||||
candidates_token_count: resp.usage.output_tokens,
|
||||
total_token_count: resp.usage.input_tokens + resp.usage.output_tokens,
|
||||
cached_content_token_count: 0,
|
||||
}),
|
||||
model_version: Some(resp.model.clone()),
|
||||
}
|
||||
}
|
||||
pub use helpers::{build_tool_id_map, merge_consecutive_roles};
|
||||
pub use request::{anthropic_to_gemini_request, compute_gemini_request_warnings};
|
||||
pub use response::gemini_to_anthropic_response;
|
||||
pub use reverse::{anthropic_to_gemini_response, gemini_to_anthropic_request};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
// Common translation utilities shared across the Gemini message-map directions:
|
||||
// tool-id lookup and Gemini role-alternation merging.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::anthropic::messages as anthropic;
|
||||
use crate::gemini::request as gemini;
|
||||
|
||||
/// Build a map from Anthropic tool_use IDs to tool names.
|
||||
///
|
||||
/// Scans all messages for `ToolUse` blocks so that `ToolResult` translation can
|
||||
/// look up the function name Gemini expects.
|
||||
pub fn build_tool_id_map(messages: &[anthropic::InputMessage]) -> HashMap<String, String> {
|
||||
let mut map = HashMap::new();
|
||||
for msg in messages {
|
||||
let blocks = match &msg.content {
|
||||
anthropic::Content::Text(_) => continue,
|
||||
anthropic::Content::Blocks(b) => b,
|
||||
};
|
||||
for block in blocks {
|
||||
if let anthropic::ContentBlock::ToolUse { id, name, .. } = block {
|
||||
map.insert(id.clone(), name.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
map
|
||||
}
|
||||
|
||||
/// Merge consecutive same-role `Content` entries by concatenating their parts.
|
||||
///
|
||||
/// Gemini requires strict user/model role alternation. When the Anthropic
|
||||
/// conversation has two consecutive user (or model) turns, this merges them
|
||||
/// into a single turn.
|
||||
pub fn merge_consecutive_roles(contents: Vec<gemini::Content>) -> Vec<gemini::Content> {
|
||||
let mut merged: Vec<gemini::Content> = Vec::with_capacity(contents.len());
|
||||
for c in contents {
|
||||
if let Some(last) = merged.last_mut() {
|
||||
if last.role == c.role {
|
||||
last.parts.extend(c.parts);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
merged.push(c);
|
||||
}
|
||||
|
||||
// Gemini requires the first content turn to have role "user". An Anthropic
|
||||
// client may legally send an assistant-first conversation (for few-shot
|
||||
// prompting). Prepend a dummy user turn so Gemini does not return a 400.
|
||||
if merged.first().and_then(|c| c.role.as_deref()) == Some("model") {
|
||||
merged.insert(
|
||||
0,
|
||||
gemini::Content {
|
||||
role: Some("user".to_string()),
|
||||
parts: vec![gemini::Part::text(String::new())],
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
merged
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
// Request direction: Anthropic -> Gemini.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::helpers::{build_tool_id_map, merge_consecutive_roles};
|
||||
use crate::anthropic::messages as anthropic;
|
||||
use crate::gemini::request as gemini;
|
||||
use crate::mapping::tools_map::sanitize_schema_for_gemini;
|
||||
|
||||
/// Compute degradation warnings for a Gemini-bound request.
|
||||
///
|
||||
/// Call before translating to surface features that are silently dropped during
|
||||
/// Anthropic → Gemini translation. Emit the result as an `x-anyllm-degradation`
|
||||
/// response header so clients can detect lossy translations.
|
||||
pub fn compute_gemini_request_warnings(
|
||||
req: &anthropic::MessageCreateRequest,
|
||||
) -> crate::mapping::warnings::TranslationWarnings {
|
||||
use crate::mapping::warnings::TranslationWarnings;
|
||||
let mut w = TranslationWarnings::default();
|
||||
|
||||
// Single pass: collect all per-block warning flags at once.
|
||||
let mut has_thinking = false;
|
||||
let mut has_document = false;
|
||||
let mut has_url_image = false;
|
||||
for msg in &req.messages {
|
||||
if let anthropic::Content::Blocks(blocks) = &msg.content {
|
||||
for b in blocks {
|
||||
match b {
|
||||
// Thinking/RedactedThinking: no Gemini Content equivalent.
|
||||
anthropic::ContentBlock::Thinking { .. }
|
||||
| anthropic::ContentBlock::RedactedThinking { .. } => has_thinking = true,
|
||||
// Document blocks have no Gemini equivalent.
|
||||
anthropic::ContentBlock::Document { .. } => has_document = true,
|
||||
// URL-type images: Gemini only accepts inline base64 data.
|
||||
anthropic::ContentBlock::Image { source } if source.source_type != "base64" => {
|
||||
has_url_image = true
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if has_thinking {
|
||||
w.add("thinking_blocks");
|
||||
}
|
||||
if has_document {
|
||||
w.add("document_blocks");
|
||||
}
|
||||
if has_url_image {
|
||||
w.add("url_images");
|
||||
}
|
||||
|
||||
// cache_control on system blocks is dropped; Gemini has no prompt-caching API.
|
||||
if let Some(anthropic::System::Blocks(blocks)) = &req.system {
|
||||
if blocks.iter().any(|b| b.cache_control.is_some()) {
|
||||
w.add("cache_control");
|
||||
}
|
||||
}
|
||||
|
||||
w
|
||||
}
|
||||
|
||||
/// Convert an Anthropic `MessageCreateRequest` into a Gemini `GenerateContentRequest`.
|
||||
///
|
||||
/// Maps `thinking_config` to `generationConfig.thinkingConfig` for Gemini 2.5
|
||||
/// thinking models. Drops unsupported features (thinking content blocks in prior
|
||||
/// messages, document blocks, cache_control) and merges consecutive same-role
|
||||
/// messages to satisfy Gemini's strict alternation requirement.
|
||||
pub fn anthropic_to_gemini_request(
|
||||
req: &anthropic::MessageCreateRequest,
|
||||
) -> gemini::GenerateContentRequest {
|
||||
let tool_id_map = build_tool_id_map(&req.messages);
|
||||
|
||||
// System instruction
|
||||
let system_instruction = req.system.as_ref().map(|sys| {
|
||||
let text = match sys {
|
||||
anthropic::System::Text(s) => s.clone(),
|
||||
anthropic::System::Blocks(blocks) => blocks
|
||||
.iter()
|
||||
.map(|b| b.text.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
};
|
||||
gemini::Content {
|
||||
role: None,
|
||||
parts: vec![gemini::Part::text(text)],
|
||||
}
|
||||
});
|
||||
|
||||
// Convert messages
|
||||
let mut contents: Vec<gemini::Content> = Vec::new();
|
||||
for msg in &req.messages {
|
||||
let role = match msg.role {
|
||||
anthropic::Role::User => "user",
|
||||
anthropic::Role::Assistant => "model",
|
||||
};
|
||||
let parts = content_blocks_to_parts(&msg.content, &tool_id_map);
|
||||
if !parts.is_empty() {
|
||||
contents.push(gemini::Content {
|
||||
role: Some(role.to_string()),
|
||||
parts,
|
||||
});
|
||||
}
|
||||
}
|
||||
contents = merge_consecutive_roles(contents);
|
||||
|
||||
// Tools
|
||||
let tools = req.tools.as_ref().map(|tools| {
|
||||
vec![gemini::Tool {
|
||||
function_declarations: tools
|
||||
.iter()
|
||||
.map(|t| gemini::FunctionDeclaration {
|
||||
name: t.name.clone(),
|
||||
description: t.description.clone(),
|
||||
parameters: Some(sanitize_schema_for_gemini(t.input_schema.clone())),
|
||||
})
|
||||
.collect(),
|
||||
}]
|
||||
});
|
||||
|
||||
// Tool config
|
||||
let tool_config = req.tool_choice.as_ref().map(|tc| {
|
||||
if matches!(
|
||||
tc,
|
||||
anthropic::ToolChoice::Auto {
|
||||
disable_parallel_tool_use: Some(true)
|
||||
} | anthropic::ToolChoice::Any {
|
||||
disable_parallel_tool_use: Some(true)
|
||||
}
|
||||
) {
|
||||
tracing::warn!(
|
||||
"disable_parallel_tool_use=true is not supported by Gemini; \
|
||||
parallel tool calls may still occur"
|
||||
);
|
||||
}
|
||||
let (mode, allowed) = match tc {
|
||||
anthropic::ToolChoice::Auto { .. } => ("AUTO", None),
|
||||
anthropic::ToolChoice::Any { .. } => ("ANY", None),
|
||||
anthropic::ToolChoice::None => ("NONE", None),
|
||||
// Gemini ANY + allowedFunctionNames restricts to a specific tool.
|
||||
anthropic::ToolChoice::Tool { name, .. } => ("ANY", Some(vec![name.clone()])),
|
||||
};
|
||||
gemini::ToolConfig {
|
||||
function_calling_config: gemini::FunctionCallingConfig {
|
||||
mode: mode.to_string(),
|
||||
allowed_function_names: allowed,
|
||||
},
|
||||
}
|
||||
});
|
||||
|
||||
// Generation config
|
||||
let generation_config = {
|
||||
let thinking_config =
|
||||
if let Some(anthropic::ThinkingConfig::Enabled { budget_tokens }) = &req.thinking {
|
||||
Some(gemini::ThinkingConfig {
|
||||
thinking_budget: *budget_tokens,
|
||||
include_thoughts: Some(true),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let gc = gemini::GenerationConfig {
|
||||
max_output_tokens: Some(req.max_tokens),
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
top_k: req.top_k,
|
||||
stop_sequences: req.stop_sequences.clone(),
|
||||
thinking_config,
|
||||
..Default::default()
|
||||
};
|
||||
Some(gc)
|
||||
};
|
||||
|
||||
gemini::GenerateContentRequest {
|
||||
contents,
|
||||
system_instruction,
|
||||
generation_config,
|
||||
tools,
|
||||
tool_config,
|
||||
safety_settings: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Anthropic message content into a vec of Gemini Parts.
|
||||
fn content_blocks_to_parts(
|
||||
content: &anthropic::Content,
|
||||
tool_id_map: &HashMap<String, String>,
|
||||
) -> Vec<gemini::Part> {
|
||||
match content {
|
||||
anthropic::Content::Text(s) => vec![gemini::Part::text(s.clone())],
|
||||
anthropic::Content::Blocks(blocks) => blocks
|
||||
.iter()
|
||||
.filter_map(|block| content_block_to_part(block, tool_id_map))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single Anthropic ContentBlock to a Gemini Part, or None if dropped.
|
||||
fn content_block_to_part(
|
||||
block: &anthropic::ContentBlock,
|
||||
tool_id_map: &HashMap<String, String>,
|
||||
) -> Option<gemini::Part> {
|
||||
match block {
|
||||
anthropic::ContentBlock::Text { text } => Some(gemini::Part::text(text.clone())),
|
||||
|
||||
anthropic::ContentBlock::Image { source } => {
|
||||
// Gemini only supports inline base64 data, not URLs.
|
||||
if source.source_type == "base64" {
|
||||
let mime = source
|
||||
.media_type
|
||||
.clone()
|
||||
.unwrap_or_else(|| "image/png".into());
|
||||
let data = source.data.clone().unwrap_or_default();
|
||||
Some(gemini::Part::inline_data(mime, data))
|
||||
} else {
|
||||
// URL-type images cannot be sent as inline_data; drop.
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
anthropic::ContentBlock::ToolUse { name, input, .. } => {
|
||||
// Strip the Anthropic tool_use id; Gemini uses name-based correlation.
|
||||
Some(gemini::Part::function_call(name.clone(), input.clone()))
|
||||
}
|
||||
|
||||
anthropic::ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let Some(name) = tool_id_map.get(tool_use_id).cloned() else {
|
||||
// Gemini requires the function name to match a declared FunctionDeclaration.
|
||||
// Emitting an unknown name causes a 400; drop the result instead.
|
||||
tracing::warn!(
|
||||
tool_use_id,
|
||||
"dropping ToolResult: tool_use_id not found in tool_id_map"
|
||||
);
|
||||
return None;
|
||||
};
|
||||
|
||||
let response_value = tool_result_to_json(content, *is_error);
|
||||
Some(gemini::Part::function_response(name, response_value))
|
||||
}
|
||||
|
||||
// Thinking, RedactedThinking, Document: not supported by Gemini, drop.
|
||||
anthropic::ContentBlock::Thinking { .. }
|
||||
| anthropic::ContentBlock::RedactedThinking { .. }
|
||||
| anthropic::ContentBlock::Document { .. } => None,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Anthropic ToolResult content into a JSON value for Gemini FunctionResponse.
|
||||
fn tool_result_to_json(
|
||||
content: &Option<anthropic::ToolResultContent>,
|
||||
is_error: Option<bool>,
|
||||
) -> serde_json::Value {
|
||||
let text = match content {
|
||||
Some(anthropic::ToolResultContent::Text(s)) => s.clone(),
|
||||
Some(anthropic::ToolResultContent::Blocks(blocks)) => {
|
||||
// Concatenate text blocks; other block types (e.g., images) cannot be
|
||||
// represented in Gemini FunctionResponse and are replaced with a placeholder.
|
||||
blocks
|
||||
.iter()
|
||||
.map(|b| match b {
|
||||
anthropic::ContentBlock::Text { text } => text.clone(),
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
"tool_result contains non-text block; \
|
||||
replacing with \"[non-text]\" placeholder for Gemini"
|
||||
);
|
||||
"[non-text]".into()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
if is_error == Some(true) {
|
||||
serde_json::json!({ "error": text })
|
||||
} else {
|
||||
serde_json::json!({ "result": text })
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// Response direction: Gemini -> Anthropic.
|
||||
|
||||
use crate::anthropic::messages as anthropic;
|
||||
use crate::gemini::request as gemini;
|
||||
use crate::gemini::response as gemini_resp;
|
||||
use crate::util::ids::{generate_message_id, generate_tool_use_id};
|
||||
|
||||
/// Convert a Gemini `GenerateContentResponse` into an Anthropic `MessageResponse`.
|
||||
///
|
||||
/// Uses only the first candidate. Synthesizes Anthropic-format tool IDs for any
|
||||
/// function calls.
|
||||
pub fn gemini_to_anthropic_response(
|
||||
resp: &gemini_resp::GenerateContentResponse,
|
||||
model: &str,
|
||||
) -> anthropic::MessageResponse {
|
||||
let candidate = resp.candidates.first();
|
||||
|
||||
let content = candidate
|
||||
.map(|c| {
|
||||
c.content
|
||||
.parts
|
||||
.iter()
|
||||
.filter_map(gemini_part_to_content_block)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let has_function_call = content
|
||||
.iter()
|
||||
.any(|b| matches!(b, anthropic::ContentBlock::ToolUse { .. }));
|
||||
|
||||
let stop_reason = candidate
|
||||
.and_then(|c| c.finish_reason.as_ref())
|
||||
.map(|fr| match fr {
|
||||
gemini_resp::FinishReason::STOP if has_function_call => anthropic::StopReason::ToolUse,
|
||||
gemini_resp::FinishReason::STOP => anthropic::StopReason::EndTurn,
|
||||
gemini_resp::FinishReason::MAX_TOKENS => anthropic::StopReason::MaxTokens,
|
||||
// SAFETY, RECITATION, LANGUAGE, OTHER, Unknown all map to EndTurn.
|
||||
_ => anthropic::StopReason::EndTurn,
|
||||
})
|
||||
// No finish_reason at all (e.g. empty candidates) -> EndTurn.
|
||||
.or(if candidate.is_some() {
|
||||
Some(anthropic::StopReason::EndTurn)
|
||||
} else {
|
||||
None
|
||||
});
|
||||
|
||||
let usage = resp
|
||||
.usage_metadata
|
||||
.as_ref()
|
||||
.map(|u| anthropic::Usage {
|
||||
input_tokens: u.prompt_token_count,
|
||||
output_tokens: u.candidates_token_count,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
anthropic::MessageResponse {
|
||||
id: generate_message_id(),
|
||||
response_type: "message".into(),
|
||||
role: anthropic::Role::Assistant,
|
||||
content,
|
||||
model: model.to_string(),
|
||||
stop_reason,
|
||||
stop_sequence: None,
|
||||
usage,
|
||||
created: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single Gemini Part to an Anthropic ContentBlock, or None if not mappable.
|
||||
fn gemini_part_to_content_block(part: &gemini::Part) -> Option<anthropic::ContentBlock> {
|
||||
// Thought parts from thinking models map to Anthropic thinking blocks.
|
||||
if part.thought == Some(true) {
|
||||
return part
|
||||
.text
|
||||
.as_ref()
|
||||
.map(|text| anthropic::ContentBlock::Thinking {
|
||||
thinking: text.clone(),
|
||||
signature: None,
|
||||
});
|
||||
}
|
||||
if let Some(text) = &part.text {
|
||||
return Some(anthropic::ContentBlock::Text { text: text.clone() });
|
||||
}
|
||||
if let Some(fc) = &part.function_call {
|
||||
return Some(anthropic::ContentBlock::ToolUse {
|
||||
id: generate_tool_use_id(),
|
||||
name: fc.name.clone(),
|
||||
input: fc.args.clone(),
|
||||
});
|
||||
}
|
||||
// inline_data, file_data, function_response: not expected in model output,
|
||||
// or have no Anthropic equivalent. Drop.
|
||||
None
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user