Compare commits
10
Commits
20a41f404d
...
7ab1fab51b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7ab1fab51b | ||
|
|
1d0cf515ed | ||
|
|
80f5bcd5ce | ||
|
|
6449bb5f4a | ||
|
|
732ca01cc7 | ||
|
|
68e6c8c46e | ||
|
|
e19fcd5417 | ||
|
|
3916ef3158 | ||
|
|
5ec2011b4e | ||
|
|
776cdbab69 |
Generated
+1
@@ -1314,6 +1314,7 @@ dependencies = [
|
|||||||
"sqlx",
|
"sqlx",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-test",
|
"tokio-test",
|
||||||
|
"tower",
|
||||||
"tower-http",
|
"tower-http",
|
||||||
"tracing",
|
"tracing",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ tracing = "0.1"
|
|||||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||||
anyhow = "1"
|
anyhow = "1"
|
||||||
tower-http = { version = "0.6", features = ["cors", "trace"] }
|
tower-http = { version = "0.6", features = ["cors", "trace"] }
|
||||||
|
tower = { version = "0.5", features = ["util"] }
|
||||||
futures = "0.3"
|
futures = "0.3"
|
||||||
bytes = "1"
|
bytes = "1"
|
||||||
dotenv = "0.15"
|
dotenv = "0.15"
|
||||||
|
|||||||
@@ -0,0 +1,33 @@
|
|||||||
|
# Build stage
|
||||||
|
FROM rust:1.82-slim AS builder
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y \
|
||||||
|
pkg-config \
|
||||||
|
libssl-dev \
|
||||||
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
COPY Cargo.toml Cargo.lock ./
|
||||||
|
RUN mkdir src && echo 'fn main(){}' > src/main.rs
|
||||||
|
RUN cargo build --release
|
||||||
|
RUN rm src/main.rs
|
||||||
|
|
||||||
|
COPY src ./src
|
||||||
|
RUN touch src/main.rs && cargo build --release
|
||||||
|
|
||||||
|
# Runtime stage
|
||||||
|
FROM debian:bookworm-slim AS runtime
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y \
|
||||||
|
ca-certificates \
|
||||||
|
libssl3 \
|
||||||
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
COPY --from=builder /app/target/release/openai-proxy /usr/local/bin/openai-proxy
|
||||||
|
|
||||||
|
EXPOSE 8080
|
||||||
|
|
||||||
|
ENTRYPOINT ["openai-proxy"]
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
version: "3.9"
|
||||||
|
|
||||||
|
services:
|
||||||
|
postgres:
|
||||||
|
image: postgres:16-alpine
|
||||||
|
restart: unless-stopped
|
||||||
|
environment:
|
||||||
|
POSTGRES_USER: proxy
|
||||||
|
POSTGRES_PASSWORD: proxy
|
||||||
|
POSTGRES_DB: proxydb
|
||||||
|
volumes:
|
||||||
|
- postgres_data:/var/lib/postgresql/data
|
||||||
|
- ./init.sql:/docker-entrypoint-initdb.d/init.sql
|
||||||
|
ports:
|
||||||
|
- "5432:5432"
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD-SHELL", "pg_isready -U proxy -d proxydb"]
|
||||||
|
interval: 5s
|
||||||
|
timeout: 5s
|
||||||
|
retries: 5
|
||||||
|
|
||||||
|
proxy:
|
||||||
|
build:
|
||||||
|
context: .
|
||||||
|
dockerfile: Dockerfile
|
||||||
|
restart: unless-stopped
|
||||||
|
depends_on:
|
||||||
|
postgres:
|
||||||
|
condition: service_healthy
|
||||||
|
environment:
|
||||||
|
DATABASE_URL: postgres://proxy:proxy@postgres:5432/proxydb
|
||||||
|
UPSTREAM_URL: ${UPSTREAM_URL:-https://api.openai.com}
|
||||||
|
UPSTREAM_API_KEY: ${UPSTREAM_API_KEY}
|
||||||
|
BIND_ADDR: 0.0.0.0
|
||||||
|
PORT: 8080
|
||||||
|
RUST_LOG: ${RUST_LOG:-info}
|
||||||
|
ports:
|
||||||
|
- "${PORT:-8080}:8080"
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
postgres_data:
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS message_cache (
|
||||||
|
id BIGSERIAL PRIMARY KEY,
|
||||||
|
hash TEXT NOT NULL UNIQUE,
|
||||||
|
role TEXT NOT NULL,
|
||||||
|
content_json TEXT NOT NULL,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
hit_count BIGINT NOT NULL DEFAULT 0
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_message_cache_hash ON message_cache (hash);
|
||||||
@@ -0,0 +1,407 @@
|
|||||||
|
//! Stress bench: in-process mock upstream + full proxy stack.
|
||||||
|
//!
|
||||||
|
//! Three bench modes selectable via env:
|
||||||
|
//! MODE=cold (default) fresh random payloads — tests miss path
|
||||||
|
//! MODE=hot same 100 payloads repeated — tests L1 DRAM hit path
|
||||||
|
//!
|
||||||
|
//! Usage (no postgres needed for quick bench):
|
||||||
|
//! BENCH_SKIP_DB=1 cargo run --bin stress --release
|
||||||
|
//!
|
||||||
|
//! With postgres + redis:
|
||||||
|
//! DATABASE_URL=... REDIS_URL=redis://localhost:6379 cargo run --bin stress --release
|
||||||
|
|
||||||
|
use std::{
|
||||||
|
sync::{
|
||||||
|
atomic::{AtomicU64, Ordering},
|
||||||
|
Arc,
|
||||||
|
},
|
||||||
|
time::Instant,
|
||||||
|
};
|
||||||
|
|
||||||
|
use axum::{
|
||||||
|
extract::State,
|
||||||
|
http::StatusCode,
|
||||||
|
routing::{get, post},
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use hdrhistogram::Histogram;
|
||||||
|
use openai_proxy::{
|
||||||
|
cache::CacheStore,
|
||||||
|
proxy::ProxyClient,
|
||||||
|
routes::{self, AppState},
|
||||||
|
};
|
||||||
|
use rand::{Rng, SeedableRng};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use sqlx::postgres::PgPoolOptions;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tower_http::cors::{Any, CorsLayer};
|
||||||
|
|
||||||
|
// ── Mock upstream ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Default, Clone)]
|
||||||
|
struct MockCounters {
|
||||||
|
hits: Arc<AtomicU64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_completions(
|
||||||
|
State(c): State<MockCounters>,
|
||||||
|
Json(_body): Json<Value>,
|
||||||
|
) -> (StatusCode, Json<Value>) {
|
||||||
|
c.hits.fetch_add(1, Ordering::Relaxed);
|
||||||
|
// Reply size controlled by MOCK_REPLY_SIZE env (default 256 bytes for 10k test)
|
||||||
|
let sz: usize = std::env::var("MOCK_REPLY_SIZE")
|
||||||
|
.ok().and_then(|s| s.parse().ok()).unwrap_or(256);
|
||||||
|
(
|
||||||
|
StatusCode::OK,
|
||||||
|
Json(json!({
|
||||||
|
"id": "cmpl-mock",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"created": 1699000000u64,
|
||||||
|
"model": "mock",
|
||||||
|
"choices": [{
|
||||||
|
"index": 0,
|
||||||
|
"message": {"role": "assistant", "content": "x".repeat(sz)},
|
||||||
|
"finish_reason": "stop"
|
||||||
|
}],
|
||||||
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_health() -> &'static str { "ok" }
|
||||||
|
|
||||||
|
// ── Payload generator ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn msg_size() -> usize {
|
||||||
|
std::env::var("MSG_SIZE").ok().and_then(|s| s.parse().ok()).unwrap_or(512)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn gen_msg(rng: &mut impl Rng, role: &str) -> Value {
|
||||||
|
let base = msg_size();
|
||||||
|
// Add up to 50% jitter so hashes are unique even in hot pool
|
||||||
|
let len = base + (rng.gen::<usize>() % (base / 2 + 1));
|
||||||
|
let content: String = (0..len).map(|_| rng.gen_range('a'..='z')).collect();
|
||||||
|
json!({"role": role, "content": content})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn gen_payload(rng: &mut impl Rng) -> Value {
|
||||||
|
let turns = 1 + (rng.gen::<usize>() % 8);
|
||||||
|
let mut msgs = vec![gen_msg(rng, "system")];
|
||||||
|
for i in 0..turns {
|
||||||
|
msgs.push(gen_msg(rng, if i % 2 == 0 { "user" } else { "assistant" }));
|
||||||
|
}
|
||||||
|
json!({
|
||||||
|
"model": "openai/gpt-4o",
|
||||||
|
"messages": msgs,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"max_tokens": 50,
|
||||||
|
"stream": false
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pre-generate a fixed pool of payloads for hot-path testing
|
||||||
|
fn gen_pool(n: usize) -> Vec<Value> {
|
||||||
|
let mut rng = rand::rngs::SmallRng::seed_from_u64(0xcafe);
|
||||||
|
(0..n).map(|_| gen_payload(&mut rng)).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Sweep ────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Returns (elapsed_s, MiB/s, req/s, p50_ms, p95_ms, p99_ms, p999_ms, errors)
|
||||||
|
///
|
||||||
|
/// Each worker maintains a local Histogram — zero shared-lock contention.
|
||||||
|
/// Histograms are merged after all workers finish.
|
||||||
|
async fn run_sweep(
|
||||||
|
proxy_url: &str,
|
||||||
|
concurrency: usize,
|
||||||
|
iters_per_worker: usize,
|
||||||
|
hot_pool: Option<Arc<Vec<Value>>>,
|
||||||
|
) -> (f64, f64, f64, f64, f64, f64, f64, u64) {
|
||||||
|
let req_b = Arc::new(AtomicU64::new(0));
|
||||||
|
let res_b = Arc::new(AtomicU64::new(0));
|
||||||
|
let errs = Arc::new(AtomicU64::new(0));
|
||||||
|
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.pool_max_idle_per_host(concurrency * 2)
|
||||||
|
.tcp_nodelay(true)
|
||||||
|
.build().unwrap();
|
||||||
|
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let mut handles = Vec::with_capacity(concurrency);
|
||||||
|
|
||||||
|
for w in 0..concurrency {
|
||||||
|
let client = client.clone();
|
||||||
|
let url = proxy_url.to_string();
|
||||||
|
let req_b = req_b.clone();
|
||||||
|
let res_b = res_b.clone();
|
||||||
|
let errs = errs.clone();
|
||||||
|
let pool_ref = hot_pool.clone();
|
||||||
|
|
||||||
|
handles.push(tokio::spawn(async move {
|
||||||
|
// Per-worker local histogram — NO shared lock
|
||||||
|
let mut local_hist = Histogram::<u64>::new_with_bounds(1, 60_000_000, 3).unwrap();
|
||||||
|
let mut rng = rand::rngs::SmallRng::seed_from_u64(w as u64 * 0xdeadbeef + 1);
|
||||||
|
|
||||||
|
for i in 0..iters_per_worker {
|
||||||
|
let payload = match &pool_ref {
|
||||||
|
Some(p) => p[i % p.len()].clone(),
|
||||||
|
None => gen_payload(&mut rng),
|
||||||
|
};
|
||||||
|
let bytes = serde_json::to_vec(&payload).unwrap();
|
||||||
|
req_b.fetch_add(bytes.len() as u64, Ordering::Relaxed);
|
||||||
|
|
||||||
|
let ts = Instant::now();
|
||||||
|
match client.post(&url)
|
||||||
|
.header("Authorization", "Bearer sk-bench")
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(bytes)
|
||||||
|
.send().await
|
||||||
|
{
|
||||||
|
Ok(r) => {
|
||||||
|
let us = ts.elapsed().as_micros() as u64;
|
||||||
|
if let Ok(b) = r.bytes().await {
|
||||||
|
res_b.fetch_add(b.len() as u64, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
local_hist.record(us).ok();
|
||||||
|
}
|
||||||
|
Err(_) => { errs.fetch_add(1, Ordering::Relaxed); }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
local_hist
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Merge all per-worker histograms
|
||||||
|
let mut merged = Histogram::<u64>::new_with_bounds(1, 60_000_000, 3).unwrap();
|
||||||
|
for h in handles {
|
||||||
|
if let Ok(local) = h.await {
|
||||||
|
merged.add(local).ok();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let elapsed = t0.elapsed().as_secs_f64();
|
||||||
|
let total_b = req_b.load(Ordering::Relaxed) + res_b.load(Ordering::Relaxed);
|
||||||
|
let total_r = (concurrency * iters_per_worker) as u64 - errs.load(Ordering::Relaxed);
|
||||||
|
|
||||||
|
(
|
||||||
|
elapsed,
|
||||||
|
total_b as f64 / elapsed / 1_048_576.0,
|
||||||
|
total_r as f64 / elapsed,
|
||||||
|
merged.value_at_percentile(50.0) as f64 / 1000.0,
|
||||||
|
merged.value_at_percentile(95.0) as f64 / 1000.0,
|
||||||
|
merged.value_at_percentile(99.0) as f64 / 1000.0,
|
||||||
|
merged.value_at_percentile(99.9) as f64 / 1000.0,
|
||||||
|
errs.load(Ordering::Relaxed),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn print_header() {
|
||||||
|
println!(
|
||||||
|
"{:<6} {:<5} {:>10} {:>10} {:>8} {:>8} {:>8} {:>9} {:>6}",
|
||||||
|
"mode", "conc", "req/s", "MiB/s", "p50ms", "p95ms", "p99ms", "p99.9ms", "errs"
|
||||||
|
);
|
||||||
|
println!("{}", "-".repeat(82));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn print_row(mode: &str, conc: usize, el: f64, mib: f64, rps: f64, p50: f64, p95: f64, p99: f64, p999: f64, errs: u64) {
|
||||||
|
println!(
|
||||||
|
"{:<6} {:>5} {:>10.1} {:>10.1} {:>8.2} {:>8.2} {:>8.2} {:>9.2} {:>6}",
|
||||||
|
mode, conc, rps, mib, p50, p95, p99, p999, errs
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Main ────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
let ncpus = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4);
|
||||||
|
let skip_db = std::env::var("BENCH_SKIP_DB").is_ok();
|
||||||
|
|
||||||
|
eprintln!("\n╔══════════════════════════════════════════════════════════╗");
|
||||||
|
eprintln!(" openai-proxy stress bench | {} CPUs | DB={}",
|
||||||
|
ncpus, if skip_db { "SKIPPED" } else { "postgres" });
|
||||||
|
eprintln!("╚══════════════════════════════════════════════════════════╝");
|
||||||
|
|
||||||
|
// ── 1. Mock upstream ─────────────────────────────────────────────────────
|
||||||
|
let mock_port: u16 = 19099;
|
||||||
|
let mock_base = format!("http://127.0.0.1:{}/v1", mock_port);
|
||||||
|
let counters = MockCounters::default();
|
||||||
|
let mock_app = Router::new()
|
||||||
|
.route("/v1/chat/completions", post(mock_completions))
|
||||||
|
.route("/health", get(mock_health))
|
||||||
|
.with_state(counters.clone());
|
||||||
|
let mock_listener = TcpListener::bind(format!("127.0.0.1:{}", mock_port)).await?;
|
||||||
|
tokio::spawn(async move { axum::serve(mock_listener, mock_app).await.unwrap() });
|
||||||
|
unsafe { std::env::set_var("BENCH_UPSTREAM_URL", &mock_base) };
|
||||||
|
eprintln!("Mock upstream → 127.0.0.1:{}", mock_port);
|
||||||
|
|
||||||
|
// ── 2. Cache layer ───────────────────────────────────────────────────────
|
||||||
|
let db_url = std::env::var("DATABASE_URL")
|
||||||
|
.unwrap_or_else(|_| "postgres://proxy:proxy@localhost:5432/proxydb".to_string());
|
||||||
|
|
||||||
|
let pool = if skip_db {
|
||||||
|
// Lazy pool that never actually connects unless we call it
|
||||||
|
PgPoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect_lazy(&db_url)?
|
||||||
|
} else {
|
||||||
|
eprintln!("DB → {}", db_url);
|
||||||
|
PgPoolOptions::new()
|
||||||
|
.max_connections((ncpus * 4) as u32)
|
||||||
|
.connect(&db_url).await?
|
||||||
|
};
|
||||||
|
|
||||||
|
let redis_pool = std::env::var("REDIS_URL").ok().map(|url| {
|
||||||
|
eprintln!("Redis → {}", url);
|
||||||
|
deadpool_redis::Config::from_url(&url)
|
||||||
|
.create_pool(Some(deadpool_redis::Runtime::Tokio1))
|
||||||
|
.expect("Redis pool")
|
||||||
|
});
|
||||||
|
if redis_pool.is_none() {
|
||||||
|
eprintln!("Redis → disabled");
|
||||||
|
}
|
||||||
|
|
||||||
|
let cache = CacheStore::new(pool.clone(), redis_pool);
|
||||||
|
if !skip_db {
|
||||||
|
cache.init_schema().await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 3. Proxy server ──────────────────────────────────────────────────────
|
||||||
|
let proxy_port: u16 = 18080;
|
||||||
|
let proxy_url = format!("http://127.0.0.1:{}/v1/chat/completions", proxy_port);
|
||||||
|
let proxy = ProxyClient::new(cache, openai_proxy::middleware::Pipeline::new())?;
|
||||||
|
let state = AppState { proxy };
|
||||||
|
let proxy_app = Router::new()
|
||||||
|
.route("/v1/chat/completions", post(routes::chat_completions))
|
||||||
|
.route("/health", get(routes::health))
|
||||||
|
.layer(CorsLayer::new().allow_origin(Any).allow_methods(Any).allow_headers(Any))
|
||||||
|
.with_state(state);
|
||||||
|
let proxy_listener = TcpListener::bind(format!("127.0.0.1:{}", proxy_port)).await?;
|
||||||
|
tokio::spawn(async move { axum::serve(proxy_listener, proxy_app).await.unwrap() });
|
||||||
|
eprintln!("Proxy → 127.0.0.1:{}", proxy_port);
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(150)).await;
|
||||||
|
|
||||||
|
// ── 4. Pre-generate hot payload pool ────────────────────────────────────────
|
||||||
|
let hot_pool = Arc::new(gen_pool(100));
|
||||||
|
|
||||||
|
// ── 5. Concurrency sweeps ─────────────────────────────────────────────────
|
||||||
|
// Iterations scale with concurrency so total work is constant per level
|
||||||
|
let base_iters = 200; // per worker
|
||||||
|
let conc_levels: &[usize] = &[1, 2, 4, 8, 16, 32, 64, 128];
|
||||||
|
|
||||||
|
// Warmup
|
||||||
|
eprintln!("\nWarmup...");
|
||||||
|
run_sweep(&proxy_url, 8, 50, Some(hot_pool.clone())).await;
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||||
|
|
||||||
|
eprintln!();
|
||||||
|
print_header();
|
||||||
|
|
||||||
|
let mut results: Vec<(String, usize, f64, f64, f64, f64, f64, f64, f64, u64)> = Vec::new();
|
||||||
|
|
||||||
|
// Cold mode: fresh random payloads (measures full proxy+upstream+cache-write path)
|
||||||
|
for &conc in conc_levels {
|
||||||
|
let (el, mib, rps, p50, p95, p99, p999, errs) =
|
||||||
|
run_sweep(&proxy_url, conc, base_iters, None).await;
|
||||||
|
print_row("cold", conc, el, mib, rps, p50, p95, p99, p999, errs);
|
||||||
|
results.push(("cold".into(), conc, el, mib, rps, p50, p95, p99, p999, errs));
|
||||||
|
}
|
||||||
|
|
||||||
|
println!();
|
||||||
|
|
||||||
|
// Hot mode: repeated payloads from fixed pool (L1 DashMap hits after warmup)
|
||||||
|
// First run primes the message-content L1; second run is measured.
|
||||||
|
run_sweep(&proxy_url, 32, base_iters, Some(hot_pool.clone())).await; // prime L1
|
||||||
|
for &conc in conc_levels {
|
||||||
|
let (el, mib, rps, p50, p95, p99, p999, errs) =
|
||||||
|
run_sweep(&proxy_url, conc, base_iters, Some(hot_pool.clone())).await;
|
||||||
|
print_row("hot", conc, el, mib, rps, p50, p95, p99, p999, errs);
|
||||||
|
results.push(("hot".into(), conc, el, mib, rps, p50, p95, p99, p999, errs));
|
||||||
|
}
|
||||||
|
|
||||||
|
println!();
|
||||||
|
|
||||||
|
// Conv-cache mode: same pool, but now conv_l1 is warm (full response cached).
|
||||||
|
// Every request should short-circuit at get_conv() — zero upstream calls.
|
||||||
|
let mock_hits_before = counters.hits.load(Ordering::Relaxed);
|
||||||
|
run_sweep(&proxy_url, 32, base_iters, Some(hot_pool.clone())).await; // prime conv_l1
|
||||||
|
let conv_hits_baseline = counters.hits.load(Ordering::Relaxed) - mock_hits_before;
|
||||||
|
for &conc in conc_levels {
|
||||||
|
let mock_before = counters.hits.load(Ordering::Relaxed);
|
||||||
|
let (el, mib, rps, p50, p95, p99, p999, errs) =
|
||||||
|
run_sweep(&proxy_url, conc, base_iters, Some(hot_pool.clone())).await;
|
||||||
|
let mock_after = counters.hits.load(Ordering::Relaxed);
|
||||||
|
let upstream_calls = mock_after - mock_before;
|
||||||
|
print_row("conv", conc, el, mib, rps, p50, p95, p99, p999, errs);
|
||||||
|
// Print upstream call count inline as a sanity check
|
||||||
|
if upstream_calls > 0 {
|
||||||
|
eprintln!(" [conv conc={conc}] upstream calls: {upstream_calls} (expected ~0 for conv cache hits)");
|
||||||
|
}
|
||||||
|
results.push(("conv".into(), conc, el, mib, rps, p50, p95, p99, p999, errs));
|
||||||
|
}
|
||||||
|
let _ = conv_hits_baseline; // used only to prime
|
||||||
|
|
||||||
|
// ── 6. Summary ────────────────────────────────────────────────────────────
|
||||||
|
let peak_cold = results.iter().filter(|r| r.0 == "cold")
|
||||||
|
.max_by(|a, b| a.4.partial_cmp(&b.4).unwrap()).unwrap();
|
||||||
|
let peak_hot = results.iter().filter(|r| r.0 == "hot")
|
||||||
|
.max_by(|a, b| a.4.partial_cmp(&b.4).unwrap()).unwrap();
|
||||||
|
let peak_conv = results.iter().filter(|r| r.0 == "conv")
|
||||||
|
.max_by(|a, b| a.4.partial_cmp(&b.4).unwrap()).unwrap();
|
||||||
|
|
||||||
|
let cold_base = results.iter().find(|r| r.0 == "cold" && r.1 == 1).map(|r| r.4).unwrap_or(1.0);
|
||||||
|
let hot_base = results.iter().find(|r| r.0 == "hot" && r.1 == 1).map(|r| r.4).unwrap_or(1.0);
|
||||||
|
let conv_base = results.iter().find(|r| r.0 == "conv" && r.1 == 1).map(|r| r.4).unwrap_or(1.0);
|
||||||
|
|
||||||
|
eprintln!("\n╔══════════════════════════════════════════════════════════╗");
|
||||||
|
eprintln!(" Peak COLD (unique payloads, full proxy path)");
|
||||||
|
eprintln!(" {:>9.1} req/s {:>7.1} MiB/s conc={}", peak_cold.4, peak_cold.3, peak_cold.1);
|
||||||
|
eprintln!(" p50={:.2}ms p95={:.2}ms p99={:.2}ms", peak_cold.5, peak_cold.6, peak_cold.7);
|
||||||
|
eprintln!();
|
||||||
|
eprintln!(" Peak HOT (msg-content L1 DRAM hits)");
|
||||||
|
eprintln!(" {:>9.1} req/s {:>7.1} MiB/s conc={}", peak_hot.4, peak_hot.3, peak_hot.1);
|
||||||
|
eprintln!(" p50={:.2}ms p95={:.2}ms p99={:.2}ms", peak_hot.5, peak_hot.6, peak_hot.7);
|
||||||
|
eprintln!();
|
||||||
|
eprintln!(" Peak CONV (conv-id L1 full-response cache hits)");
|
||||||
|
eprintln!(" {:>9.1} req/s {:>7.1} MiB/s conc={}", peak_conv.4, peak_conv.3, peak_conv.1);
|
||||||
|
eprintln!(" p50={:.2}ms p95={:.2}ms p99={:.2}ms", peak_conv.5, peak_conv.6, peak_conv.7);
|
||||||
|
eprintln!();
|
||||||
|
eprintln!(" Scaling (cold, relative to conc=1 = {:.1} req/s)", cold_base);
|
||||||
|
for r in results.iter().filter(|r| r.0 == "cold") {
|
||||||
|
let eff = (r.4 / cold_base) / r.1 as f64 * 100.0;
|
||||||
|
eprintln!(" conc={:>3} {:>9.1} req/s {:.2}x eff={:.1}%", r.1, r.4, r.4/cold_base, eff);
|
||||||
|
}
|
||||||
|
eprintln!();
|
||||||
|
eprintln!(" Scaling (hot, relative to conc=1 = {:.1} req/s)", hot_base);
|
||||||
|
for r in results.iter().filter(|r| r.0 == "hot") {
|
||||||
|
let eff = (r.4 / hot_base) / r.1 as f64 * 100.0;
|
||||||
|
eprintln!(" conc={:>3} {:>9.1} req/s {:.2}x eff={:.1}%", r.1, r.4, r.4/hot_base, eff);
|
||||||
|
}
|
||||||
|
eprintln!();
|
||||||
|
eprintln!(" Scaling (conv, relative to conc=1 = {:.1} req/s)", conv_base);
|
||||||
|
for r in results.iter().filter(|r| r.0 == "conv") {
|
||||||
|
let eff = (r.4 / conv_base) / r.1 as f64 * 100.0;
|
||||||
|
eprintln!(" conc={:>3} {:>9.1} req/s {:.2}x eff={:.1}%", r.1, r.4, r.4/conv_base, eff);
|
||||||
|
}
|
||||||
|
eprintln!("╚══════════════════════════════════════════════════════════╝");
|
||||||
|
|
||||||
|
// Check 10k target
|
||||||
|
let target = 10_000.0f64;
|
||||||
|
if peak_hot.4 >= target {
|
||||||
|
eprintln!("\n✅ TARGET MET: {:.0} req/s >= {:.0} req/s (hot path)", peak_hot.4, target);
|
||||||
|
} else {
|
||||||
|
eprintln!("\n⚠ HOT peak: {:.0} req/s (target: {:.0})", peak_hot.4, target);
|
||||||
|
eprintln!(" COLD peak: {:.0} req/s", peak_cold.4);
|
||||||
|
}
|
||||||
|
|
||||||
|
if !skip_db {
|
||||||
|
// Count rows written by the async background task
|
||||||
|
if let Ok(row) = sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM message_cache")
|
||||||
|
.fetch_one(&pool).await
|
||||||
|
{
|
||||||
|
eprintln!(" Postgres rows (unique hashes): {}", row.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -0,0 +1,541 @@
|
|||||||
|
use std::collections::VecDeque;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use sqlx::PgPool;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
use xxhash_rust::xxh3::xxh3_128;
|
||||||
|
|
||||||
|
use crate::types::{ChatCompletionRequest, Message, MessageContent};
|
||||||
|
|
||||||
|
// ── Constants ─────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Max entries in per-node DRAM cache. At avg 32 KiB each ≈ 320 MiB worst-case.
|
||||||
|
const L1_CAP: usize = 10_000;
|
||||||
|
/// Max conversation responses in per-node DRAM cache. At avg 4 KiB ≈ 40 MiB worst-case.
|
||||||
|
const CONV_L1_CAP: usize = 10_000;
|
||||||
|
/// Redis TTL in seconds (24 h).
|
||||||
|
const REDIS_TTL: usize = 86_400;
|
||||||
|
/// Redis key prefix for individual message cache.
|
||||||
|
const REDIS_PFX: &str = "mc:";
|
||||||
|
/// Redis key prefix for conversation response cache.
|
||||||
|
const CONV_REDIS_PFX: &str = "cv:";
|
||||||
|
|
||||||
|
// ── Public types ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct CachedEntry {
|
||||||
|
pub role: Arc<str>,
|
||||||
|
pub content_json: Arc<str>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Internal flush messages ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
struct Flush {
|
||||||
|
hash_hex: String,
|
||||||
|
role: String,
|
||||||
|
content_json: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ConvFlush {
|
||||||
|
conv_id_hex: String,
|
||||||
|
response_json: Arc<[u8]>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── FIFO eviction ring ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Single-writer FIFO ring that tracks insertion order for L1 eviction.
|
||||||
|
/// Lives entirely inside the background flush task — zero contention on reads.
|
||||||
|
struct FifoRing {
|
||||||
|
buf: VecDeque<u128>,
|
||||||
|
cap: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FifoRing {
|
||||||
|
fn new(cap: usize) -> Self {
|
||||||
|
Self { buf: VecDeque::with_capacity(cap + 1), cap }
|
||||||
|
}
|
||||||
|
/// Returns the hash to evict if over capacity, else None.
|
||||||
|
fn push(&mut self, h: u128) -> Option<u128> {
|
||||||
|
self.buf.push_back(h);
|
||||||
|
if self.buf.len() > self.cap { self.buf.pop_front() } else { None }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── CacheStore ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Three-tier write-through cache:
|
||||||
|
/// L1 DashMap (DRAM, this node, FIFO eviction) ~50 ns
|
||||||
|
/// L2 Redis (shared all nodes, TTL 24 h, optional) ~200 µs
|
||||||
|
/// L3 Postgres (durable, async background, analytics) ~2 ms — NEVER on hot path
|
||||||
|
///
|
||||||
|
/// Clone is O(1) — all state is behind Arc.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct CacheStore {
|
||||||
|
// ── message-content cache (existing) ──
|
||||||
|
l1: Arc<DashMap<u128, CachedEntry>>,
|
||||||
|
flush_tx: mpsc::UnboundedSender<Flush>,
|
||||||
|
// ── conversation-response cache (new) ──
|
||||||
|
/// conv_id → full serialised response bytes
|
||||||
|
conv_l1: Arc<DashMap<u128, Arc<[u8]>>>,
|
||||||
|
conv_flush_tx: mpsc::UnboundedSender<ConvFlush>,
|
||||||
|
redis: Option<deadpool_redis::Pool>,
|
||||||
|
/// Kept for init_schema only; L3 writes go through flush channels.
|
||||||
|
pg: PgPool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CacheStore {
|
||||||
|
pub fn new(pg: PgPool, redis: Option<deadpool_redis::Pool>) -> Self {
|
||||||
|
let l1: Arc<DashMap<u128, CachedEntry>> = Arc::new(DashMap::with_capacity(L1_CAP));
|
||||||
|
let (tx, rx) = mpsc::unbounded_channel();
|
||||||
|
tokio::spawn(flush_task(rx, l1.clone(), pg.clone()));
|
||||||
|
|
||||||
|
let conv_l1: Arc<DashMap<u128, Arc<[u8]>>> =
|
||||||
|
Arc::new(DashMap::with_capacity(CONV_L1_CAP));
|
||||||
|
let (conv_tx, conv_rx) = mpsc::unbounded_channel();
|
||||||
|
tokio::spawn(conv_flush_task(conv_rx, conv_l1.clone(), pg.clone()));
|
||||||
|
|
||||||
|
Self { l1, flush_tx: tx, conv_l1, conv_flush_tx: conv_tx, redis, pg }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn init_schema(&self) -> anyhow::Result<()> {
|
||||||
|
sqlx::query(
|
||||||
|
r#"CREATE TABLE IF NOT EXISTS message_cache (
|
||||||
|
id BIGSERIAL PRIMARY KEY,
|
||||||
|
hash TEXT NOT NULL UNIQUE,
|
||||||
|
role TEXT NOT NULL,
|
||||||
|
content_json TEXT NOT NULL,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
hit_count BIGINT NOT NULL DEFAULT 0
|
||||||
|
)"#,
|
||||||
|
).execute(&self.pg).await?;
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_message_cache_hash ON message_cache (hash)",
|
||||||
|
).execute(&self.pg).await?;
|
||||||
|
// conversation_cache stores full response bytes keyed by conv_id
|
||||||
|
sqlx::query(
|
||||||
|
r#"CREATE TABLE IF NOT EXISTS conversation_cache (
|
||||||
|
conv_id TEXT NOT NULL PRIMARY KEY,
|
||||||
|
response BYTEA NOT NULL,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
hit_count BIGINT NOT NULL DEFAULT 0
|
||||||
|
)"#,
|
||||||
|
).execute(&self.pg).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Look up a hash. L1 hit → 50 ns. L2 hit → 200 µs, backfills L1.
|
||||||
|
/// Returns None only on complete miss (write path still caches on put).
|
||||||
|
pub async fn get(&self, hash: u128) -> Option<CachedEntry> {
|
||||||
|
if let Some(e) = self.l1.get(&hash) {
|
||||||
|
return Some(e.clone());
|
||||||
|
}
|
||||||
|
if let Some(e) = self.redis_get(hash).await {
|
||||||
|
self.l1.insert(hash, e.clone());
|
||||||
|
return Some(e);
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write through all three tiers. L1+L2 are synchronous (fast).
|
||||||
|
/// L3 (Postgres) is fire-and-forget via unbounded channel.
|
||||||
|
pub async fn put(&self, hash: u128, role: &str, content_json: &str) {
|
||||||
|
let entry = CachedEntry {
|
||||||
|
role: Arc::from(role),
|
||||||
|
content_json: Arc::from(content_json),
|
||||||
|
};
|
||||||
|
self.l1.insert(hash, entry);
|
||||||
|
self.redis_put(hash, role, content_json).await;
|
||||||
|
// Non-blocking: L3 write queued to background task
|
||||||
|
let _ = self.flush_tx.send(Flush {
|
||||||
|
hash_hex: hash_to_hex(hash),
|
||||||
|
role: role.to_string(),
|
||||||
|
content_json: content_json.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cache all messages from a request. Infallible — errors are absorbed.
|
||||||
|
pub async fn cache_messages(&self, req: &ChatCompletionRequest) {
|
||||||
|
for msg in &req.messages {
|
||||||
|
if let Some((role, cj)) = extract_cacheable(msg) {
|
||||||
|
self.put(xxh3_128(cj.as_bytes()), role, &cj).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cache an assistant reply. Returns hex hash.
|
||||||
|
pub async fn cache_assistant_reply(&self, content: &str) -> String {
|
||||||
|
let cj = serde_json::to_string(content).unwrap_or_else(|_| content.to_string());
|
||||||
|
let h = xxh3_128(cj.as_bytes());
|
||||||
|
self.put(h, "assistant", &cj).await;
|
||||||
|
hash_to_hex(h)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Redis helpers ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async fn redis_get(&self, hash: u128) -> Option<CachedEntry> {
|
||||||
|
let pool = self.redis.as_ref()?;
|
||||||
|
let mut conn = pool.get().await.ok()?;
|
||||||
|
let key = format!("{}{}", REDIS_PFX, hash_to_hex(hash));
|
||||||
|
let raw: Option<String> = redis::cmd("GET")
|
||||||
|
.arg(&key)
|
||||||
|
.query_async(&mut *conn)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten();
|
||||||
|
let raw = raw?;
|
||||||
|
let (role, cj) = raw.split_once('\n')?;
|
||||||
|
Some(CachedEntry { role: Arc::from(role), content_json: Arc::from(cj) })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn redis_put(&self, hash: u128, role: &str, content_json: &str) {
|
||||||
|
let pool = match &self.redis { Some(p) => p, None => return };
|
||||||
|
let mut conn = match pool.get().await { Ok(c) => c, Err(_) => return };
|
||||||
|
let key = format!("{}{}", REDIS_PFX, hash_to_hex(hash));
|
||||||
|
let val = format!("{}\n{}", role, content_json);
|
||||||
|
let _: Result<(), _> = redis::cmd("SETEX")
|
||||||
|
.arg(&key).arg(REDIS_TTL).arg(&val)
|
||||||
|
.query_async(&mut *conn).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Conversation-response cache ───────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Look up a full response by conversation ID.
|
||||||
|
/// L1 hit → 50 ns. L2 (Redis) hit → 200 µs, backfills L1.
|
||||||
|
pub async fn get_conv(&self, conv_id: u128) -> Option<Arc<[u8]>> {
|
||||||
|
if let Some(e) = self.conv_l1.get(&conv_id) {
|
||||||
|
return Some(e.clone());
|
||||||
|
}
|
||||||
|
if let Some(bytes) = self.conv_redis_get(conv_id).await {
|
||||||
|
self.conv_l1.insert(conv_id, bytes.clone());
|
||||||
|
return Some(bytes);
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Store a full response keyed by conversation ID.
|
||||||
|
/// L1+L2 synchronous; L3 (Postgres) fire-and-forget.
|
||||||
|
pub async fn put_conv(&self, conv_id: u128, response: &[u8]) {
|
||||||
|
let arc: Arc<[u8]> = Arc::from(response);
|
||||||
|
self.conv_l1.insert(conv_id, arc.clone());
|
||||||
|
self.conv_redis_put(conv_id, response).await;
|
||||||
|
let _ = self.conv_flush_tx.send(ConvFlush {
|
||||||
|
conv_id_hex: hash_to_hex(conv_id),
|
||||||
|
response_json: arc,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn conv_redis_get(&self, conv_id: u128) -> Option<Arc<[u8]>> {
|
||||||
|
let pool = self.redis.as_ref()?;
|
||||||
|
let mut conn = pool.get().await.ok()?;
|
||||||
|
let key = format!("{}{}", CONV_REDIS_PFX, hash_to_hex(conv_id));
|
||||||
|
let raw: Option<Vec<u8>> = redis::cmd("GET")
|
||||||
|
.arg(&key)
|
||||||
|
.query_async(&mut *conn)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten();
|
||||||
|
raw.map(|v| Arc::from(v.as_slice()))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn conv_redis_put(&self, conv_id: u128, response: &[u8]) {
|
||||||
|
let pool = match &self.redis { Some(p) => p, None => return };
|
||||||
|
let mut conn = match pool.get().await { Ok(c) => c, Err(_) => return };
|
||||||
|
let key = format!("{}{}", CONV_REDIS_PFX, hash_to_hex(conv_id));
|
||||||
|
let _: Result<(), _> = redis::cmd("SETEX")
|
||||||
|
.arg(&key).arg(REDIS_TTL).arg(response)
|
||||||
|
.query_async(&mut *conn).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Background flush task ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Drains the flush channel, evicts L1 FIFO, writes Postgres.
|
||||||
|
/// Single task — the FifoRing has zero contention.
|
||||||
|
async fn flush_task(
|
||||||
|
mut rx: mpsc::UnboundedReceiver<Flush>,
|
||||||
|
l1: Arc<DashMap<u128, CachedEntry>>,
|
||||||
|
pg: PgPool,
|
||||||
|
) {
|
||||||
|
let mut ring = FifoRing::new(L1_CAP);
|
||||||
|
while let Some(msg) = rx.recv().await {
|
||||||
|
let h = u128::from_str_radix(&msg.hash_hex, 16).unwrap_or(0);
|
||||||
|
if let Some(evict) = ring.push(h) {
|
||||||
|
l1.remove(&evict);
|
||||||
|
}
|
||||||
|
// Best-effort Postgres upsert; never blocks the hot path
|
||||||
|
let _ = sqlx::query(
|
||||||
|
r#"INSERT INTO message_cache (hash, role, content_json)
|
||||||
|
VALUES ($1, $2, $3)
|
||||||
|
ON CONFLICT (hash) DO UPDATE
|
||||||
|
SET hit_count = message_cache.hit_count + 1"#,
|
||||||
|
)
|
||||||
|
.bind(&msg.hash_hex)
|
||||||
|
.bind(&msg.role)
|
||||||
|
.bind(&msg.content_json)
|
||||||
|
.execute(&pg)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drains the conv flush channel, evicts conv L1 FIFO, writes Postgres.
|
||||||
|
async fn conv_flush_task(
|
||||||
|
mut rx: mpsc::UnboundedReceiver<ConvFlush>,
|
||||||
|
conv_l1: Arc<DashMap<u128, Arc<[u8]>>>,
|
||||||
|
pg: PgPool,
|
||||||
|
) {
|
||||||
|
let mut ring = FifoRing::new(CONV_L1_CAP);
|
||||||
|
while let Some(msg) = rx.recv().await {
|
||||||
|
let h = u128::from_str_radix(&msg.conv_id_hex, 16).unwrap_or(0);
|
||||||
|
if let Some(evict) = ring.push(h) {
|
||||||
|
conv_l1.remove(&evict);
|
||||||
|
}
|
||||||
|
let _ = sqlx::query(
|
||||||
|
r#"INSERT INTO conversation_cache (conv_id, response)
|
||||||
|
VALUES ($1, $2)
|
||||||
|
ON CONFLICT (conv_id) DO UPDATE
|
||||||
|
SET hit_count = conversation_cache.hit_count + 1"#,
|
||||||
|
)
|
||||||
|
.bind(&msg.conv_id_hex)
|
||||||
|
.bind(msg.response_json.as_ref())
|
||||||
|
.execute(&pg)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Public helpers ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
pub fn compute_hash(content: &str) -> u128 {
|
||||||
|
xxh3_128(content.as_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn hash_to_hex(h: u128) -> String {
|
||||||
|
format!("{:032x}", h)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compute a Merkle-chain conversation ID from an ordered message slice.
|
||||||
|
///
|
||||||
|
/// The ID is derived by chaining: `running = xxh3(running_bytes ‖ msg_hash_bytes)`
|
||||||
|
/// for each cacheable message in sequence. System and Developer messages are
|
||||||
|
/// excluded — they are configuration, not conversation turns, so two requests
|
||||||
|
/// that differ only in their system prompt share the same conversation chain.
|
||||||
|
///
|
||||||
|
/// Properties:
|
||||||
|
/// - Two requests with identical message history produce identical conv_ids.
|
||||||
|
/// - A single added message produces a new ID dependent on the full prior chain.
|
||||||
|
/// - O(N) in message count, single allocation-free pass.
|
||||||
|
/// - Returns 0 for an empty / all-system message list (no cacheable turns).
|
||||||
|
pub fn compute_conv_id(messages: &[crate::types::Message]) -> u128 {
|
||||||
|
let mut running = 0u128;
|
||||||
|
for msg in messages {
|
||||||
|
if let Some((_, cj)) = extract_cacheable(msg) {
|
||||||
|
let msg_hash = xxh3_128(cj.as_bytes());
|
||||||
|
let mut buf = [0u8; 32];
|
||||||
|
buf[..16].copy_from_slice(&running.to_le_bytes());
|
||||||
|
buf[16..].copy_from_slice(&msg_hash.to_le_bytes());
|
||||||
|
running = xxh3_128(&buf);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
running
|
||||||
|
}
|
||||||
|
|
||||||
|
fn extract_cacheable(msg: &Message) -> Option<(&'static str, String)> {
|
||||||
|
match msg {
|
||||||
|
Message::User { content, .. } =>
|
||||||
|
Some(("user", serde_json::to_string(content).unwrap_or_default())),
|
||||||
|
Message::Assistant { .. } =>
|
||||||
|
Some(("assistant", serde_json::to_string(msg).unwrap_or_default())),
|
||||||
|
Message::Tool { .. } =>
|
||||||
|
Some(("tool", serde_json::to_string(msg).unwrap_or_default())),
|
||||||
|
Message::Function { .. } =>
|
||||||
|
Some(("function", serde_json::to_string(msg).unwrap_or_default())),
|
||||||
|
Message::System { .. } | Message::Developer { .. } => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn content_to_json(content: &MessageContent) -> String {
|
||||||
|
serde_json::to_string(content).unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compute_hash_deterministic() {
|
||||||
|
assert_eq!(compute_hash("hello world"), compute_hash("hello world"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_compute_hash_different_inputs() {
|
||||||
|
assert_ne!(compute_hash("hello"), compute_hash("world"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_hash_to_hex_32_chars() {
|
||||||
|
let h = hash_to_hex(compute_hash("test input"));
|
||||||
|
assert_eq!(h.len(), 32);
|
||||||
|
assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_hash_to_hex_empty() {
|
||||||
|
assert_eq!(hash_to_hex(compute_hash("")), format!("{:032x}", xxh3_128(b"")));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_fifo_ring_no_eviction_below_cap() {
|
||||||
|
let mut r = FifoRing::new(3);
|
||||||
|
assert!(r.push(1).is_none());
|
||||||
|
assert!(r.push(2).is_none());
|
||||||
|
assert!(r.push(3).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_fifo_ring_evicts_oldest() {
|
||||||
|
let mut r = FifoRing::new(3);
|
||||||
|
r.push(1); r.push(2); r.push(3);
|
||||||
|
assert_eq!(r.push(4), Some(1));
|
||||||
|
assert_eq!(r.push(5), Some(2));
|
||||||
|
assert_eq!(r.push(6), Some(3));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_cacheable_user() {
|
||||||
|
let msg = Message::User {
|
||||||
|
content: MessageContent::Text("hi".into()),
|
||||||
|
name: None,
|
||||||
|
};
|
||||||
|
let (role, _) = extract_cacheable(&msg).unwrap();
|
||||||
|
assert_eq!(role, "user");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_cacheable_system_none() {
|
||||||
|
let msg = Message::System {
|
||||||
|
content: MessageContent::Text("be helpful".into()),
|
||||||
|
name: None,
|
||||||
|
};
|
||||||
|
assert!(extract_cacheable(&msg).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_cacheable_tool() {
|
||||||
|
let msg = Message::Tool {
|
||||||
|
content: MessageContent::Text("result".into()),
|
||||||
|
tool_call_id: "call_1".into(),
|
||||||
|
};
|
||||||
|
let (role, _) = extract_cacheable(&msg).unwrap();
|
||||||
|
assert_eq!(role, "tool");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_same_content_same_hash() {
|
||||||
|
let msg = Message::User {
|
||||||
|
content: MessageContent::Text("same".into()),
|
||||||
|
name: None,
|
||||||
|
};
|
||||||
|
let (_, c1) = extract_cacheable(&msg).unwrap();
|
||||||
|
let (_, c2) = extract_cacheable(&msg).unwrap();
|
||||||
|
assert_eq!(compute_hash(&c1), compute_hash(&c2));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_redis_value_roundtrip() {
|
||||||
|
let role = "user";
|
||||||
|
let cj = r#""hello\nworld""#;
|
||||||
|
let val = format!("{}\n{}", role, cj);
|
||||||
|
let (r, c) = val.split_once('\n').unwrap();
|
||||||
|
assert_eq!(r, role);
|
||||||
|
assert_eq!(c, cj);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── compute_conv_id tests ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn user(text: &str) -> Message {
|
||||||
|
Message::User { content: MessageContent::Text(text.into()), name: None }
|
||||||
|
}
|
||||||
|
fn assistant(text: &str) -> Message {
|
||||||
|
Message::Assistant {
|
||||||
|
content: Some(MessageContent::Text(text.into())),
|
||||||
|
name: None, tool_calls: None, function_call: None,
|
||||||
|
refusal: None, audio: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
fn system(text: &str) -> Message {
|
||||||
|
Message::System { content: MessageContent::Text(text.into()), name: None }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_empty_is_zero() {
|
||||||
|
assert_eq!(compute_conv_id(&[]), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_system_only_is_zero() {
|
||||||
|
// System messages don't advance the chain
|
||||||
|
assert_eq!(compute_conv_id(&[system("be helpful")]), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_deterministic() {
|
||||||
|
let msgs = vec![user("hello"), assistant("hi there"), user("what's 2+2?")];
|
||||||
|
assert_eq!(compute_conv_id(&msgs), compute_conv_id(&msgs));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_different_content_differs() {
|
||||||
|
let a = vec![user("hello")];
|
||||||
|
let b = vec![user("goodbye")];
|
||||||
|
assert_ne!(compute_conv_id(&a), compute_conv_id(&b));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_order_matters() {
|
||||||
|
let a = vec![user("hello"), user("world")];
|
||||||
|
let b = vec![user("world"), user("hello")];
|
||||||
|
assert_ne!(compute_conv_id(&a), compute_conv_id(&b));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_prefix_shared() {
|
||||||
|
// First two turns are shared; only the third differs
|
||||||
|
let prefix = vec![user("hello"), assistant("hi")];
|
||||||
|
let id_at_prefix = compute_conv_id(&prefix);
|
||||||
|
|
||||||
|
let mut branch_a = prefix.clone();
|
||||||
|
branch_a.push(user("what's 2+2?"));
|
||||||
|
let mut branch_b = prefix.clone();
|
||||||
|
branch_b.push(user("what's the time?"));
|
||||||
|
|
||||||
|
// Both branches produce IDs different from the shared prefix ID
|
||||||
|
assert_ne!(compute_conv_id(&branch_a), id_at_prefix);
|
||||||
|
assert_ne!(compute_conv_id(&branch_b), id_at_prefix);
|
||||||
|
// And different from each other
|
||||||
|
assert_ne!(compute_conv_id(&branch_a), compute_conv_id(&branch_b));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_system_ignored_in_chain() {
|
||||||
|
// Changing only the system prompt must not change the conv_id
|
||||||
|
let a = vec![system("be helpful"), user("hello"), assistant("hi")];
|
||||||
|
let b = vec![system("be terse"), user("hello"), assistant("hi")];
|
||||||
|
assert_eq!(compute_conv_id(&a), compute_conv_id(&b));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_conv_id_extends_parent() {
|
||||||
|
// ID after turn N must differ from ID after turn N-1
|
||||||
|
let m0 = user("turn 0");
|
||||||
|
let m1 = user("turn 1");
|
||||||
|
|
||||||
|
let id0 = compute_conv_id(&[m0.clone()]);
|
||||||
|
let id1 = compute_conv_id(&[m0.clone(), m1.clone()]);
|
||||||
|
assert_ne!(id0, id1);
|
||||||
|
assert_ne!(id0, 0);
|
||||||
|
assert_ne!(id1, 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
pub mod cache;
|
||||||
|
pub mod config;
|
||||||
|
pub mod middleware;
|
||||||
|
pub mod proxy;
|
||||||
|
pub mod router;
|
||||||
|
pub mod routes;
|
||||||
|
pub mod types;
|
||||||
+27
-12
@@ -1,14 +1,9 @@
|
|||||||
mod cache;
|
|
||||||
mod config;
|
|
||||||
mod proxy;
|
|
||||||
mod routes;
|
|
||||||
mod types;
|
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use axum::{
|
use axum::{
|
||||||
routing::{get, post},
|
routing::{get, post},
|
||||||
Router,
|
Router,
|
||||||
};
|
};
|
||||||
|
use deadpool_redis::{Config as RedisConfig, Runtime as RedisRuntime};
|
||||||
use dotenv::dotenv;
|
use dotenv::dotenv;
|
||||||
use sqlx::postgres::PgPoolOptions;
|
use sqlx::postgres::PgPoolOptions;
|
||||||
use tower_http::{
|
use tower_http::{
|
||||||
@@ -18,10 +13,11 @@ use tower_http::{
|
|||||||
use tracing::info;
|
use tracing::info;
|
||||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
||||||
|
|
||||||
use crate::cache::CacheStore;
|
use openai_proxy::cache::CacheStore;
|
||||||
use crate::config::Config;
|
use openai_proxy::config::Config;
|
||||||
use crate::proxy::ProxyClient;
|
use openai_proxy::middleware::{Pipeline, StructuredOutputMiddleware};
|
||||||
use crate::routes::AppState;
|
use openai_proxy::proxy::ProxyClient;
|
||||||
|
use openai_proxy::routes::{self, AppState};
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> Result<()> {
|
async fn main() -> Result<()> {
|
||||||
@@ -40,11 +36,30 @@ async fn main() -> Result<()> {
|
|||||||
.connect(&config.database_url)
|
.connect(&config.database_url)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let cache = CacheStore::new(pool);
|
let redis_pool = config.redis_url.as_deref().map(|url| {
|
||||||
|
RedisConfig::from_url(url)
|
||||||
|
.create_pool(Some(RedisRuntime::Tokio1))
|
||||||
|
.expect("failed to create Redis pool")
|
||||||
|
});
|
||||||
|
|
||||||
|
if redis_pool.is_some() {
|
||||||
|
info!("Redis cache layer enabled (L2)");
|
||||||
|
}
|
||||||
|
|
||||||
|
let cache = CacheStore::new(pool, redis_pool);
|
||||||
cache.init_schema().await?;
|
cache.init_schema().await?;
|
||||||
info!("Database schema ready");
|
info!("Database schema ready");
|
||||||
|
|
||||||
let proxy = ProxyClient::new(config.clone(), cache)?;
|
let pipeline = Pipeline::new()
|
||||||
|
.add(StructuredOutputMiddleware::new(vec![
|
||||||
|
"gpt-4o".to_string(),
|
||||||
|
"gpt-4o-mini".to_string(),
|
||||||
|
"gpt-4.1".to_string(),
|
||||||
|
"o3-mini".to_string(),
|
||||||
|
"o1".to_string(),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let proxy = ProxyClient::new(cache, pipeline)?;
|
||||||
let state = AppState { proxy };
|
let state = AppState { proxy };
|
||||||
|
|
||||||
let cors = CorsLayer::new()
|
let cors = CorsLayer::new()
|
||||||
|
|||||||
@@ -0,0 +1,238 @@
|
|||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
use crate::types::{ChatCompletionRequest, Message, MessageContent, ResponseFormat};
|
||||||
|
|
||||||
|
pub trait Middleware: Send + Sync {
|
||||||
|
fn pre_send(&self, _req: &mut ChatCompletionRequest, _auth: &mut String) {}
|
||||||
|
fn post_receive(&self, _req: &ChatCompletionRequest, _resp: &mut Value) {}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
pub struct Pipeline {
|
||||||
|
middlewares: Vec<Box<dyn Middleware>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Pipeline {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self { middlewares: Vec::new() }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn add<M: Middleware + 'static>(mut self, m: M) -> Self {
|
||||||
|
self.middlewares.push(Box::new(m));
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn run_pre(&self, req: &mut ChatCompletionRequest, auth: &mut String) {
|
||||||
|
for mw in &self.middlewares {
|
||||||
|
mw.pre_send(req, auth);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn run_post(&self, req: &ChatCompletionRequest, resp: &mut Value) {
|
||||||
|
for mw in &self.middlewares {
|
||||||
|
mw.post_receive(req, resp);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Injects a JSON schema formatting system prompt for models that do not
|
||||||
|
/// natively support `response_format: { type: "json_schema", ... }`.
|
||||||
|
/// Models in `native_models` are left untouched.
|
||||||
|
pub struct StructuredOutputMiddleware {
|
||||||
|
native_models: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StructuredOutputMiddleware {
|
||||||
|
pub fn new(native_models: Vec<String>) -> Self {
|
||||||
|
Self { native_models }
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_native(&self, model: &str) -> bool {
|
||||||
|
self.native_models.iter().any(|n| model.contains(n.as_str()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Middleware for StructuredOutputMiddleware {
|
||||||
|
fn pre_send(&self, req: &mut ChatCompletionRequest, _auth: &mut String) {
|
||||||
|
let json_schema = match &req.response_format {
|
||||||
|
Some(ResponseFormat { format_type, json_schema: Some(schema) })
|
||||||
|
if format_type == "json_schema" || format_type == "json_object" =>
|
||||||
|
schema.clone(),
|
||||||
|
_ => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
if self.is_native(&req.model) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema_str = serde_json::to_string_pretty(&json_schema).unwrap_or_default();
|
||||||
|
let instruction = format!(
|
||||||
|
"You MUST respond with valid JSON only, conforming exactly to this JSON schema:\n\n```json\n{}\n```\n\nDo NOT include any text before or after the JSON. Output ONLY the JSON object.",
|
||||||
|
schema_str
|
||||||
|
);
|
||||||
|
|
||||||
|
let inject_pos = req.messages.iter()
|
||||||
|
.rposition(|m| matches!(m, Message::User { .. }))
|
||||||
|
.unwrap_or(req.messages.len());
|
||||||
|
|
||||||
|
req.messages.insert(inject_pos, Message::System {
|
||||||
|
content: MessageContent::Text(instruction),
|
||||||
|
name: None,
|
||||||
|
});
|
||||||
|
|
||||||
|
req.response_format = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::types::*;
|
||||||
|
|
||||||
|
fn make_req(model: &str, messages: Vec<Message>, rf: Option<ResponseFormat>) -> ChatCompletionRequest {
|
||||||
|
ChatCompletionRequest {
|
||||||
|
model: model.to_string(),
|
||||||
|
messages,
|
||||||
|
stream: Some(false),
|
||||||
|
temperature: None, top_p: None, n: None,
|
||||||
|
max_tokens: None, max_completion_tokens: None, stop: None,
|
||||||
|
presence_penalty: None, frequency_penalty: None, logit_bias: None,
|
||||||
|
user: None, seed: None, logprobs: None, top_logprobs: None,
|
||||||
|
response_format: rf,
|
||||||
|
tools: None, tool_choice: None, parallel_tool_calls: None,
|
||||||
|
functions: None, function_call: None, service_tier: None,
|
||||||
|
modalities: None, audio: None, reasoning_effort: None,
|
||||||
|
metadata: None, store: None, prediction: None,
|
||||||
|
stream_options: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct AppendTextMiddleware(String);
|
||||||
|
|
||||||
|
impl Middleware for AppendTextMiddleware {
|
||||||
|
fn pre_send(&self, req: &mut ChatCompletionRequest, _auth: &mut String) {
|
||||||
|
if let Some(Message::System { content, .. }) = req.messages.first_mut() {
|
||||||
|
if let MessageContent::Text(t) = content {
|
||||||
|
t.push_str(&self.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_pipeline_applies_middleware_in_order() {
|
||||||
|
let mut req = make_req(
|
||||||
|
"gpt-4o",
|
||||||
|
vec![Message::System {
|
||||||
|
content: MessageContent::Text("hello".into()),
|
||||||
|
name: None,
|
||||||
|
}],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
let pipeline = Pipeline::new()
|
||||||
|
.add(AppendTextMiddleware(" world".into()))
|
||||||
|
.add(AppendTextMiddleware("!!!".into()));
|
||||||
|
|
||||||
|
pipeline.run_pre(&mut req, &mut auth);
|
||||||
|
|
||||||
|
let content = match &req.messages[0] {
|
||||||
|
Message::System { content: MessageContent::Text(t), .. } => t.as_str(),
|
||||||
|
_ => panic!("expected system text"),
|
||||||
|
};
|
||||||
|
assert_eq!(content, "hello world!!!");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_empty_pipeline_is_noop() {
|
||||||
|
let mut req = make_req(
|
||||||
|
"gpt-4o",
|
||||||
|
vec![Message::User { content: MessageContent::Text("hi".into()), name: None }],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
let pipeline = Pipeline::new();
|
||||||
|
pipeline.run_pre(&mut req, &mut auth);
|
||||||
|
assert_eq!(req.messages.len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_structured_output_injects_prompt_for_non_native_model() {
|
||||||
|
let schema = serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}});
|
||||||
|
let m = StructuredOutputMiddleware::new(vec!["gpt-4o".to_string()]);
|
||||||
|
|
||||||
|
let mut req = make_req(
|
||||||
|
"dummy-model",
|
||||||
|
vec![
|
||||||
|
Message::System { content: MessageContent::Text("be helpful".into()), name: None },
|
||||||
|
Message::User { content: MessageContent::Text("hi".into()), name: None },
|
||||||
|
],
|
||||||
|
Some(ResponseFormat {
|
||||||
|
format_type: "json_schema".into(),
|
||||||
|
json_schema: Some(schema),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
m.pre_send(&mut req, &mut auth);
|
||||||
|
|
||||||
|
assert_eq!(req.messages.len(), 3, "should have 3 messages");
|
||||||
|
// Injected before the last User message (index 1), so injected is at index 1
|
||||||
|
let injected = match &req.messages[1] {
|
||||||
|
Message::System { content: MessageContent::Text(t), .. } => t.clone(),
|
||||||
|
_ => panic!("expected injected system message at index 1, got: {:?}", &req.messages[1]),
|
||||||
|
};
|
||||||
|
assert!(injected.contains("JSON schema"), "injected prompt missing 'JSON schema': {}", injected);
|
||||||
|
assert!(injected.contains("name"), "injected prompt should reference schema fields: {}", injected);
|
||||||
|
assert!(req.response_format.is_none(), "response_format should be stripped");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_structured_output_skips_native_model() {
|
||||||
|
let schema = serde_json::json!({"type": "object"});
|
||||||
|
let m = StructuredOutputMiddleware::new(vec!["gpt-4o".to_string()]);
|
||||||
|
|
||||||
|
let mut req = make_req(
|
||||||
|
"gpt-4o",
|
||||||
|
vec![Message::User { content: MessageContent::Text("hi".into()), name: None }],
|
||||||
|
Some(ResponseFormat {
|
||||||
|
format_type: "json_schema".into(),
|
||||||
|
json_schema: Some(schema),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
m.pre_send(&mut req, &mut auth);
|
||||||
|
|
||||||
|
assert_eq!(req.messages.len(), 1, "should not inject for native models");
|
||||||
|
assert!(req.response_format.is_some(), "should keep response_format for native models");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_structured_output_noops_when_no_response_format() {
|
||||||
|
let m = StructuredOutputMiddleware::new(vec![]);
|
||||||
|
|
||||||
|
let mut req = make_req(
|
||||||
|
"dummy-model",
|
||||||
|
vec![Message::User { content: MessageContent::Text("hi".into()), name: None }],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
let orig_count = req.messages.len();
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
m.pre_send(&mut req, &mut auth);
|
||||||
|
assert_eq!(req.messages.len(), orig_count, "no injection when no response_format");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_post_receive_default_is_noop() {
|
||||||
|
let m = StructuredOutputMiddleware::new(vec![]);
|
||||||
|
let req = make_req("dummy", vec![], None);
|
||||||
|
let mut resp = serde_json::json!({"choices": [{"message": {"content": "test"}}]});
|
||||||
|
let original = resp.clone();
|
||||||
|
m.post_receive(&req, &mut resp);
|
||||||
|
assert_eq!(resp, original);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,252 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use bytes::Bytes;
|
||||||
|
use reqwest::Client;
|
||||||
|
use serde_json::Value;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::cache::{self, CacheStore};
|
||||||
|
use crate::middleware::Pipeline;
|
||||||
|
use crate::router::{self, Provider};
|
||||||
|
use crate::types::ChatCompletionRequest;
|
||||||
|
|
||||||
|
fn num_cpus() -> usize {
|
||||||
|
std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(4)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inner holds the two hot shared resources on their own 64-byte cache lines
|
||||||
|
/// so tasks running on different cores do not bounce the same cache line.
|
||||||
|
#[repr(align(64))]
|
||||||
|
struct Inner {
|
||||||
|
client: Client,
|
||||||
|
cache: CacheStore,
|
||||||
|
pipeline: Pipeline,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// ProxyClient is a newtype around Arc<Inner>.
|
||||||
|
/// Clone is one atomic increment — hand one to every spawned task for free.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ProxyClient(Arc<Inner>);
|
||||||
|
|
||||||
|
pub enum ProxyResponse {
|
||||||
|
Buffered(Bytes),
|
||||||
|
Streamed(reqwest::Response),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProxyClient {
|
||||||
|
pub fn new(cache: CacheStore, pipeline: Pipeline) -> Result<Self> {
|
||||||
|
let client = Client::builder()
|
||||||
|
.pool_max_idle_per_host(num_cpus() * 2)
|
||||||
|
.build()
|
||||||
|
.context("failed to build reqwest client")?;
|
||||||
|
Ok(Self(Arc::new(Inner { client, cache, pipeline })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Forward a ChatCompletionRequest upstream.
|
||||||
|
///
|
||||||
|
/// - Model prefix is stripped and used to select the upstream URL.
|
||||||
|
/// - Caller's `Authorization` header is forwarded verbatim.
|
||||||
|
/// - Inbound messages and the assistant reply are cached in Postgres.
|
||||||
|
pub async fn forward_request(
|
||||||
|
&self,
|
||||||
|
mut req: ChatCompletionRequest,
|
||||||
|
mut auth_header: String,
|
||||||
|
) -> Result<ProxyResponse> {
|
||||||
|
let target = router::resolve(&req.model);
|
||||||
|
debug!(provider = ?target.provider, model = %target.model, "routing");
|
||||||
|
req.model = target.model.clone();
|
||||||
|
|
||||||
|
// Compute conv_id from the original messages (before middleware mutation).
|
||||||
|
// Non-streaming only — streaming responses can't be cached.
|
||||||
|
let is_streaming = req.stream.unwrap_or(false);
|
||||||
|
let conv_id = if !is_streaming {
|
||||||
|
let id = cache::compute_conv_id(&req.messages);
|
||||||
|
if id != 0 {
|
||||||
|
if let Some(cached) = self.0.cache.get_conv(id).await {
|
||||||
|
debug!(conv_id = %cache::hash_to_hex(id), "conv cache hit");
|
||||||
|
return Ok(ProxyResponse::Buffered(Bytes::from(cached.to_vec())));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
id
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
|
||||||
|
// Pre-send: middleware may modify req and auth_header
|
||||||
|
self.0.pipeline.run_pre(&mut req, &mut auth_header);
|
||||||
|
|
||||||
|
self.0.cache.cache_messages(&req).await;
|
||||||
|
|
||||||
|
let url = format!("{}/chat/completions", target.base_url);
|
||||||
|
|
||||||
|
let mut builder = self
|
||||||
|
.0
|
||||||
|
.client
|
||||||
|
.post(&url)
|
||||||
|
.header("Authorization", &auth_header)
|
||||||
|
.json(&req);
|
||||||
|
|
||||||
|
if target.provider == Provider::Anthropic {
|
||||||
|
builder = builder.header("anthropic-version", "2023-06-01");
|
||||||
|
}
|
||||||
|
|
||||||
|
let http_resp = builder.send().await.context("upstream request failed")?;
|
||||||
|
|
||||||
|
if !http_resp.status().is_success() {
|
||||||
|
let status = http_resp.status();
|
||||||
|
let body = http_resp.bytes().await.unwrap_or_default();
|
||||||
|
return Err(anyhow::anyhow!(
|
||||||
|
"upstream returned {}: {}",
|
||||||
|
status,
|
||||||
|
String::from_utf8_lossy(&body)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
if is_streaming {
|
||||||
|
return Ok(ProxyResponse::Streamed(http_resp));
|
||||||
|
}
|
||||||
|
|
||||||
|
let body_bytes = http_resp.bytes().await.context("failed to read upstream body")?;
|
||||||
|
|
||||||
|
if let Ok(mut json) = serde_json::from_slice::<Value>(&body_bytes) {
|
||||||
|
// Post-receive: middleware may rewrite the response
|
||||||
|
self.0.pipeline.run_post(&req, &mut json);
|
||||||
|
|
||||||
|
if let Some(content) = json
|
||||||
|
.get("choices")
|
||||||
|
.and_then(|c| c.get(0))
|
||||||
|
.and_then(|c| c.get("message"))
|
||||||
|
.and_then(|m| m.get("content"))
|
||||||
|
.and_then(|c| c.as_str())
|
||||||
|
{
|
||||||
|
self.0.cache.cache_assistant_reply(content).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let modified = serde_json::to_vec(&json).context("failed to re-serialize response")?;
|
||||||
|
// Store full response in conv cache for future identical requests
|
||||||
|
if conv_id != 0 {
|
||||||
|
self.0.cache.put_conv(conv_id, &modified).await;
|
||||||
|
}
|
||||||
|
return Ok(ProxyResponse::Buffered(Bytes::from(modified)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(ProxyResponse::Buffered(body_bytes))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Forward GET /v1/models to OpenAI, forwarding the caller's auth.
|
||||||
|
pub async fn forward_get(&self, path: &str, auth_header: &str) -> Result<Bytes> {
|
||||||
|
let url = format!("https://api.openai.com/v1/{}", path.trim_start_matches('/'));
|
||||||
|
let resp = self
|
||||||
|
.0
|
||||||
|
.client
|
||||||
|
.get(&url)
|
||||||
|
.header("Authorization", auth_header)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.context("upstream GET failed")?;
|
||||||
|
let bytes = resp.bytes().await.context("failed to read upstream GET body")?;
|
||||||
|
Ok(bytes)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_url_construction_openai() {
|
||||||
|
let target = router::resolve("openai/gpt-4o");
|
||||||
|
let url = format!("{}/chat/completions", target.base_url);
|
||||||
|
assert_eq!(url, "https://api.openai.com/v1/chat/completions");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_url_construction_anthropic() {
|
||||||
|
let target = router::resolve("anthropic/claude-3-5-sonnet-20241022");
|
||||||
|
let url = format!("{}/chat/completions", target.base_url);
|
||||||
|
assert_eq!(url, "https://api.anthropic.com/v1/chat/completions");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_model_prefix_stripped() {
|
||||||
|
let target = router::resolve("anthropic/claude-3-5-haiku-20241022");
|
||||||
|
assert_eq!(target.model, "claude-3-5-haiku-20241022");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_assistant_content_extraction() {
|
||||||
|
let json = serde_json::json!({
|
||||||
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "42"}, "finish_reason": "stop"}]
|
||||||
|
});
|
||||||
|
let content = json
|
||||||
|
.get("choices").and_then(|c| c.get(0))
|
||||||
|
.and_then(|c| c.get("message"))
|
||||||
|
.and_then(|m| m.get("content"))
|
||||||
|
.and_then(|c| c.as_str());
|
||||||
|
assert_eq!(content, Some("42"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_pipeline_run_pre_called() {
|
||||||
|
use crate::middleware::{Middleware, Pipeline};
|
||||||
|
use crate::types::{Message, MessageContent};
|
||||||
|
use std::sync::{
|
||||||
|
atomic::{AtomicUsize, Ordering},
|
||||||
|
Arc,
|
||||||
|
};
|
||||||
|
|
||||||
|
struct CounterMiddleware(Arc<AtomicUsize>);
|
||||||
|
impl Middleware for CounterMiddleware {
|
||||||
|
fn pre_send(&self, req: &mut ChatCompletionRequest, _auth: &mut String) {
|
||||||
|
self.0.fetch_add(1, Ordering::SeqCst);
|
||||||
|
req.messages.push(Message::User {
|
||||||
|
content: MessageContent::Text("injected".into()),
|
||||||
|
name: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let count = Arc::new(AtomicUsize::new(0));
|
||||||
|
let pipeline = Pipeline::new().add(CounterMiddleware(count.clone()));
|
||||||
|
|
||||||
|
let mut req = ChatCompletionRequest {
|
||||||
|
model: "openai/gpt-4o".into(),
|
||||||
|
messages: vec![Message::User {
|
||||||
|
content: MessageContent::Text("hello".into()),
|
||||||
|
name: None,
|
||||||
|
}],
|
||||||
|
stream: Some(false),
|
||||||
|
temperature: None, top_p: None, n: None,
|
||||||
|
max_tokens: None, max_completion_tokens: None, stop: None,
|
||||||
|
presence_penalty: None, frequency_penalty: None, logit_bias: None,
|
||||||
|
user: None, seed: None, logprobs: None, top_logprobs: None,
|
||||||
|
response_format: None, tools: None, tool_choice: None,
|
||||||
|
parallel_tool_calls: None, functions: None, function_call: None,
|
||||||
|
service_tier: None, modalities: None, audio: None,
|
||||||
|
reasoning_effort: None, metadata: None, store: None, prediction: None,
|
||||||
|
stream_options: None,
|
||||||
|
};
|
||||||
|
let mut auth = "Bearer test".to_string();
|
||||||
|
|
||||||
|
pipeline.run_pre(&mut req, &mut auth);
|
||||||
|
|
||||||
|
assert_eq!(count.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(req.messages.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_null_content_not_cached() {
|
||||||
|
let json = serde_json::json!({
|
||||||
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": null}, "finish_reason": "tool_calls"}]
|
||||||
|
});
|
||||||
|
let content = json
|
||||||
|
.get("choices").and_then(|c| c.get(0))
|
||||||
|
.and_then(|c| c.get("message"))
|
||||||
|
.and_then(|m| m.get("content"))
|
||||||
|
.and_then(|c| c.as_str());
|
||||||
|
assert!(content.is_none());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,120 @@
|
|||||||
|
use axum::{
|
||||||
|
body::Body,
|
||||||
|
extract::State,
|
||||||
|
http::{HeaderMap, HeaderValue, StatusCode},
|
||||||
|
response::{IntoResponse, Response},
|
||||||
|
Json,
|
||||||
|
};
|
||||||
|
use futures::StreamExt;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use tracing::{error, info};
|
||||||
|
|
||||||
|
use crate::proxy::{ProxyClient, ProxyResponse};
|
||||||
|
use crate::types::ChatCompletionRequest;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AppState {
|
||||||
|
pub proxy: ProxyClient,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn health() -> impl IntoResponse {
|
||||||
|
Json(json!({"status": "ok"}))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn chat_completions(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(req): Json<ChatCompletionRequest>,
|
||||||
|
) -> Response {
|
||||||
|
let is_streaming = req.stream.unwrap_or(false);
|
||||||
|
info!(model = %req.model, streaming = is_streaming, "chat completion request");
|
||||||
|
|
||||||
|
match state.proxy.forward_request(req).await {
|
||||||
|
Ok(ProxyResponse::Buffered(bytes)) => {
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert(
|
||||||
|
axum::http::header::CONTENT_TYPE,
|
||||||
|
HeaderValue::from_static("application/json"),
|
||||||
|
);
|
||||||
|
(StatusCode::OK, headers, bytes).into_response()
|
||||||
|
}
|
||||||
|
Ok(ProxyResponse::Streamed(upstream_resp)) => {
|
||||||
|
let status = upstream_resp.status();
|
||||||
|
let upstream_headers = upstream_resp.headers().clone();
|
||||||
|
|
||||||
|
let stream = upstream_resp.bytes_stream().map(|result| {
|
||||||
|
result.map_err(|e| {
|
||||||
|
error!("stream error: {}", e);
|
||||||
|
std::io::Error::new(std::io::ErrorKind::Other, e)
|
||||||
|
})
|
||||||
|
});
|
||||||
|
|
||||||
|
let body = Body::from_stream(stream);
|
||||||
|
let mut resp = Response::new(body);
|
||||||
|
*resp.status_mut() = status;
|
||||||
|
|
||||||
|
if let Some(ct) = upstream_headers.get(axum::http::header::CONTENT_TYPE) {
|
||||||
|
resp.headers_mut().insert(axum::http::header::CONTENT_TYPE, ct.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
resp
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("proxy error: {}", e);
|
||||||
|
let body = Json(json!({
|
||||||
|
"error": {
|
||||||
|
"message": e.to_string(),
|
||||||
|
"type": "proxy_error",
|
||||||
|
"code": "upstream_error"
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
(StatusCode::BAD_GATEWAY, body).into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn models(State(state): State<AppState>) -> Response {
|
||||||
|
match state.proxy.forward_get("/v1/models").await {
|
||||||
|
Ok(bytes) => {
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert(
|
||||||
|
axum::http::header::CONTENT_TYPE,
|
||||||
|
HeaderValue::from_static("application/json"),
|
||||||
|
);
|
||||||
|
(StatusCode::OK, headers, bytes).into_response()
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("models proxy error: {}", e);
|
||||||
|
(
|
||||||
|
StatusCode::BAD_GATEWAY,
|
||||||
|
Json(json!({"error": {"message": e.to_string()}})),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use axum::{body::to_bytes, http::Request, routing::get};
|
||||||
|
use tower::util::ServiceExt;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_health_returns_200() {
|
||||||
|
let app = axum::Router::new().route("/health", get(health));
|
||||||
|
let response = app
|
||||||
|
.oneshot(
|
||||||
|
Request::builder()
|
||||||
|
.uri("/health")
|
||||||
|
.body(Body::empty())
|
||||||
|
.unwrap(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||||
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(json["status"], "ok");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,488 @@
|
|||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(tag = "type", rename_all = "snake_case")]
|
||||||
|
pub enum ContentPart {
|
||||||
|
Text {
|
||||||
|
text: String,
|
||||||
|
},
|
||||||
|
ImageUrl {
|
||||||
|
image_url: ImageUrl,
|
||||||
|
},
|
||||||
|
InputAudio {
|
||||||
|
input_audio: InputAudio,
|
||||||
|
},
|
||||||
|
File {
|
||||||
|
file: FileContent,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ImageUrl {
|
||||||
|
pub url: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub detail: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct InputAudio {
|
||||||
|
pub data: String,
|
||||||
|
pub format: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct FileContent {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub file_id: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub filename: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub file_data: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
pub enum MessageContent {
|
||||||
|
Text(String),
|
||||||
|
Parts(Vec<ContentPart>),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct Function {
|
||||||
|
pub name: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub description: Option<String>,
|
||||||
|
pub parameters: Value,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub strict: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct Tool {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub tool_type: String,
|
||||||
|
pub function: Function,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct FunctionCall {
|
||||||
|
pub name: String,
|
||||||
|
pub arguments: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ToolCall {
|
||||||
|
pub id: String,
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub call_type: String,
|
||||||
|
pub function: FunctionCall,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(tag = "role", rename_all = "lowercase")]
|
||||||
|
pub enum Message {
|
||||||
|
System {
|
||||||
|
content: MessageContent,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
name: Option<String>,
|
||||||
|
},
|
||||||
|
User {
|
||||||
|
content: MessageContent,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
name: Option<String>,
|
||||||
|
},
|
||||||
|
Assistant {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
content: Option<MessageContent>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
name: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
tool_calls: Option<Vec<ToolCall>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
function_call: Option<FunctionCall>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
refusal: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
audio: Option<AssistantAudio>,
|
||||||
|
},
|
||||||
|
Tool {
|
||||||
|
content: MessageContent,
|
||||||
|
tool_call_id: String,
|
||||||
|
},
|
||||||
|
Function {
|
||||||
|
content: MessageContent,
|
||||||
|
name: String,
|
||||||
|
},
|
||||||
|
Developer {
|
||||||
|
content: MessageContent,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
name: Option<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct AssistantAudio {
|
||||||
|
pub id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
pub enum ToolChoice {
|
||||||
|
String(String),
|
||||||
|
Object(ToolChoiceObject),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ToolChoiceObject {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub choice_type: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub function: Option<ToolChoiceFunction>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ToolChoiceFunction {
|
||||||
|
pub name: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ResponseFormat {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub format_type: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub json_schema: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct StreamOptions {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub include_usage: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct AudioOutput {
|
||||||
|
pub voice: String,
|
||||||
|
pub format: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct Prediction {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub prediction_type: String,
|
||||||
|
pub content: MessageContent,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ChatCompletionRequest {
|
||||||
|
pub model: String,
|
||||||
|
pub messages: Vec<Message>,
|
||||||
|
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub stream: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub stream_options: Option<StreamOptions>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub temperature: Option<f64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub top_p: Option<f64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub n: Option<u32>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub max_tokens: Option<u32>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub max_completion_tokens: Option<u32>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub stop: Option<Value>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub presence_penalty: Option<f64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub frequency_penalty: Option<f64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub logit_bias: Option<HashMap<String, f64>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub user: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub seed: Option<i64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub logprobs: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub top_logprobs: Option<u8>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub response_format: Option<ResponseFormat>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tools: Option<Vec<Tool>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tool_choice: Option<ToolChoice>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub parallel_tool_calls: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub functions: Option<Vec<Function>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub function_call: Option<Value>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub service_tier: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub modalities: Option<Vec<String>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub audio: Option<AudioOutput>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub reasoning_effort: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub metadata: Option<HashMap<String, Value>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub store: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub prediction: Option<Prediction>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Usage {
|
||||||
|
pub prompt_tokens: u32,
|
||||||
|
pub completion_tokens: u32,
|
||||||
|
pub total_tokens: u32,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub completion_tokens_details: Option<Value>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub prompt_tokens_details: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct LogprobContent {
|
||||||
|
pub token: String,
|
||||||
|
pub logprob: f64,
|
||||||
|
pub bytes: Option<Vec<u8>>,
|
||||||
|
pub top_logprobs: Vec<TopLogprob>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct TopLogprob {
|
||||||
|
pub token: String,
|
||||||
|
pub logprob: f64,
|
||||||
|
pub bytes: Option<Vec<u8>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Logprobs {
|
||||||
|
pub content: Option<Vec<LogprobContent>>,
|
||||||
|
pub refusal: Option<Vec<LogprobContent>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ResponseMessage {
|
||||||
|
pub role: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub content: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tool_calls: Option<Vec<ToolCall>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub function_call: Option<FunctionCall>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub refusal: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub audio: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Choice {
|
||||||
|
pub index: u32,
|
||||||
|
pub message: ResponseMessage,
|
||||||
|
pub finish_reason: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub logprobs: Option<Logprobs>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ChatCompletionResponse {
|
||||||
|
pub id: String,
|
||||||
|
pub object: String,
|
||||||
|
pub created: i64,
|
||||||
|
pub model: String,
|
||||||
|
pub choices: Vec<Choice>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub usage: Option<Usage>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub system_fingerprint: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub service_tier: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ChunkDelta {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub role: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub content: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tool_calls: Option<Vec<Value>>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub function_call: Option<Value>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub refusal: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct StreamChoice {
|
||||||
|
pub index: u32,
|
||||||
|
pub delta: ChunkDelta,
|
||||||
|
pub finish_reason: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub logprobs: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ChatCompletionChunk {
|
||||||
|
pub id: String,
|
||||||
|
pub object: String,
|
||||||
|
pub created: i64,
|
||||||
|
pub model: String,
|
||||||
|
pub choices: Vec<StreamChoice>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub usage: Option<Usage>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub system_fingerprint: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub service_tier: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_simple_user_message() {
|
||||||
|
let json = r#"{"role":"user","content":"Hello!"}"#;
|
||||||
|
let msg: Message = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(matches!(msg, Message::User { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_system_message() {
|
||||||
|
let json = r#"{"role":"system","content":"You are helpful."}"#;
|
||||||
|
let msg: Message = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(matches!(msg, Message::System { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_assistant_with_tool_calls() {
|
||||||
|
let json = r#"{
|
||||||
|
"role": "assistant",
|
||||||
|
"content": null,
|
||||||
|
"tool_calls": [{
|
||||||
|
"id": "call_abc",
|
||||||
|
"type": "function",
|
||||||
|
"function": {"name": "get_weather", "arguments": "{\"city\":\"NYC\"}"}
|
||||||
|
}]
|
||||||
|
}"#;
|
||||||
|
let msg: Message = serde_json::from_str(json).unwrap();
|
||||||
|
match msg {
|
||||||
|
Message::Assistant { tool_calls: Some(tc), .. } => {
|
||||||
|
assert_eq!(tc[0].id, "call_abc");
|
||||||
|
assert_eq!(tc[0].function.name, "get_weather");
|
||||||
|
}
|
||||||
|
_ => panic!("expected assistant with tool_calls"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_tool_message() {
|
||||||
|
let json = r#"{"role":"tool","content":"72 degrees","tool_call_id":"call_abc"}"#;
|
||||||
|
let msg: Message = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(matches!(msg, Message::Tool { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_content_parts() {
|
||||||
|
let json = r#"{"role":"user","content":[{"type":"text","text":"hello"},{"type":"image_url","image_url":{"url":"https://example.com/img.png"}}]}"#;
|
||||||
|
let msg: Message = serde_json::from_str(json).unwrap();
|
||||||
|
match msg {
|
||||||
|
Message::User { content: MessageContent::Parts(parts), .. } => {
|
||||||
|
assert_eq!(parts.len(), 2);
|
||||||
|
assert!(matches!(&parts[0], ContentPart::Text { text } if text == "hello"));
|
||||||
|
}
|
||||||
|
_ => panic!("expected user with parts"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_full_request() {
|
||||||
|
let json = r#"{
|
||||||
|
"model": "gpt-4o",
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "You are helpful."},
|
||||||
|
{"role": "user", "content": "What is 2+2?"}
|
||||||
|
],
|
||||||
|
"temperature": 0.7,
|
||||||
|
"max_tokens": 100,
|
||||||
|
"stream": false
|
||||||
|
}"#;
|
||||||
|
let req: ChatCompletionRequest = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(req.model, "gpt-4o");
|
||||||
|
assert_eq!(req.messages.len(), 2);
|
||||||
|
assert_eq!(req.temperature, Some(0.7));
|
||||||
|
assert_eq!(req.max_tokens, Some(100));
|
||||||
|
assert_eq!(req.stream, Some(false));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_streaming_chunk() {
|
||||||
|
let json = r#"{
|
||||||
|
"id": "chatcmpl-xyz",
|
||||||
|
"object": "chat.completion.chunk",
|
||||||
|
"created": 1699000000,
|
||||||
|
"model": "gpt-4o",
|
||||||
|
"choices": [{
|
||||||
|
"index": 0,
|
||||||
|
"delta": {"role": "assistant", "content": "Hello"},
|
||||||
|
"finish_reason": null
|
||||||
|
}]
|
||||||
|
}"#;
|
||||||
|
let chunk: ChatCompletionChunk = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(chunk.id, "chatcmpl-xyz");
|
||||||
|
assert_eq!(chunk.choices[0].delta.content.as_deref(), Some("Hello"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_roundtrip_request_serialization() {
|
||||||
|
let req = ChatCompletionRequest {
|
||||||
|
model: "gpt-4o".to_string(),
|
||||||
|
messages: vec![Message::User {
|
||||||
|
content: MessageContent::Text("hi".to_string()),
|
||||||
|
name: None,
|
||||||
|
}],
|
||||||
|
stream: Some(true),
|
||||||
|
temperature: Some(1.0),
|
||||||
|
tools: None,
|
||||||
|
tool_choice: None,
|
||||||
|
functions: None,
|
||||||
|
function_call: None,
|
||||||
|
top_p: None,
|
||||||
|
n: None,
|
||||||
|
max_tokens: None,
|
||||||
|
max_completion_tokens: None,
|
||||||
|
stop: None,
|
||||||
|
presence_penalty: None,
|
||||||
|
frequency_penalty: None,
|
||||||
|
logit_bias: None,
|
||||||
|
user: None,
|
||||||
|
seed: None,
|
||||||
|
logprobs: None,
|
||||||
|
top_logprobs: None,
|
||||||
|
response_format: None,
|
||||||
|
stream_options: None,
|
||||||
|
parallel_tool_calls: None,
|
||||||
|
service_tier: None,
|
||||||
|
modalities: None,
|
||||||
|
audio: None,
|
||||||
|
reasoning_effort: None,
|
||||||
|
metadata: None,
|
||||||
|
store: None,
|
||||||
|
prediction: None,
|
||||||
|
};
|
||||||
|
let json = serde_json::to_string(&req).unwrap();
|
||||||
|
let back: ChatCompletionRequest = serde_json::from_str(&json).unwrap();
|
||||||
|
assert_eq!(back.model, "gpt-4o");
|
||||||
|
assert_eq!(back.stream, Some(true));
|
||||||
|
}
|
||||||
|
}
|
||||||
Executable
+79
@@ -0,0 +1,79 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
BASE_URL="${PROXY_URL:-http://localhost:8080}"
|
||||||
|
API_KEY="${UPSTREAM_API_KEY:-test-key}"
|
||||||
|
|
||||||
|
echo "=== Smoke testing openai-proxy at $BASE_URL ==="
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "1. GET /health"
|
||||||
|
HEALTH=$(curl -sf "$BASE_URL/health")
|
||||||
|
echo " Response: $HEALTH"
|
||||||
|
echo "$HEALTH" | grep -q '"ok"' && echo " PASS" || { echo " FAIL"; exit 1; }
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "2. POST /v1/chat/completions (non-streaming)"
|
||||||
|
RESP=$(curl -sf "$BASE_URL/v1/chat/completions" \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-H "Authorization: Bearer $API_KEY" \
|
||||||
|
-d '{
|
||||||
|
"model": "gpt-4o-mini",
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "You are a concise assistant."},
|
||||||
|
{"role": "user", "content": "Reply with exactly: pong"}
|
||||||
|
],
|
||||||
|
"max_tokens": 10
|
||||||
|
}')
|
||||||
|
echo " Response: $RESP"
|
||||||
|
echo "$RESP" | grep -q '"choices"' && echo " PASS" || { echo " FAIL"; exit 1; }
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "3. POST /v1/chat/completions (streaming)"
|
||||||
|
STREAM_RESP=$(curl -sf "$BASE_URL/v1/chat/completions" \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-H "Authorization: Bearer $API_KEY" \
|
||||||
|
-d '{
|
||||||
|
"model": "gpt-4o-mini",
|
||||||
|
"messages": [{"role": "user", "content": "Say hi"}],
|
||||||
|
"stream": true,
|
||||||
|
"max_tokens": 10
|
||||||
|
}')
|
||||||
|
echo " First bytes: $(echo "$STREAM_RESP" | head -c 100)"
|
||||||
|
echo "$STREAM_RESP" | grep -q 'data:' && echo " PASS" || { echo " FAIL"; exit 1; }
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "4. GET /v1/models"
|
||||||
|
MODELS=$(curl -sf "$BASE_URL/v1/models" \
|
||||||
|
-H "Authorization: Bearer $API_KEY")
|
||||||
|
echo " Response (first 200 chars): $(echo "$MODELS" | head -c 200)"
|
||||||
|
echo "$MODELS" | grep -q '"data"' && echo " PASS" || { echo " FAIL"; exit 1; }
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "5. POST /v1/chat/completions (with tools)"
|
||||||
|
TOOL_RESP=$(curl -sf "$BASE_URL/v1/chat/completions" \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-H "Authorization: Bearer $API_KEY" \
|
||||||
|
-d '{
|
||||||
|
"model": "gpt-4o-mini",
|
||||||
|
"messages": [{"role": "user", "content": "What is the weather in NYC?"}],
|
||||||
|
"tools": [{
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "get_weather",
|
||||||
|
"description": "Get current weather",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {"city": {"type": "string"}},
|
||||||
|
"required": ["city"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}],
|
||||||
|
"tool_choice": "auto",
|
||||||
|
"max_tokens": 100
|
||||||
|
}')
|
||||||
|
echo " Response: $(echo "$TOOL_RESP" | head -c 300)"
|
||||||
|
echo "$TOOL_RESP" | grep -q '"choices"' && echo " PASS" || { echo " FAIL"; exit 1; }
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "=== All smoke tests PASSED ==="
|
||||||
Reference in New Issue
Block a user