Skip to content

Commit bbdbfef

Browse files
author
Linuxdazhao
committed
feat: add RateLimitLayer middleware (#320)
Uses governor for RPM limiting, placed below the retry layer so retries are also throttled. Reads x-ratelimit-remaining-requests and x-ratelimit-reset-requests headers to apply backpressure when the server quota is exhausted. Gated behind the rate-limit feature.
1 parent 884aff9 commit bbdbfef

5 files changed

Lines changed: 753 additions & 0 deletions

File tree

async-openai/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ assistant = ["assistant-types", "_api", ]
5252
administration = ["administration-types", "_api"]
5353
completions = ["completion-types", "_api"]
5454
middleware = ["dep:tower", "_api"]
55+
rate-limit = ["dep:governor", "middleware"]
5556

5657
# Type feature flags - these enable only the types
5758
response-types = ["dep:derive_builder"]
@@ -132,6 +133,7 @@ full = [
132133
"types",
133134
"byot",
134135
"middleware",
136+
"rate-limit",
135137
]
136138

137139
# Internal feature to enable API dependencies
@@ -180,6 +182,7 @@ secrecy = { version = "0.10", features = ["serde"], optional = true }
180182
serde_urlencoded = { version = "0.7", optional = true }
181183
url = { version = "2.5", optional = true }
182184
tower = { version = "0.5", features = ["limit", "retry", "timeout", "util"], optional = true }
185+
governor = { version = "0.8", optional = true, default-features = false, features = ["std"] }
183186
eventsource-stream = { version = "0.2", optional = true }
184187
## For Webhook signature verification
185188
hmac = { version = "0.12", optional = true, default-features = false}
@@ -197,6 +200,7 @@ tokio-tungstenite = { version = "0.28", optional = true, default-features = fals
197200
## API WASM dependencies
198201
[target.wasm32-unknown-unknown.dependencies]
199202
getrandom = { version = "0.3", features = ["wasm_js"] }
203+
governor = { version = "0.8", optional = true, default-features = false, features = ["std"] }
200204

201205
[dev-dependencies]
202206
http = "1"

async-openai/MIDDLEWARE.md

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,22 @@ Custom tower retry policies can call `middleware::retry::should_retry` to reuse
9191

9292
On native targets retries wait using `tokio::time::sleep`. On WASM retries are immediate.
9393

94+
## Rate limit layer
95+
96+
Enable the `rate-limit` feature to use `RateLimitLayer`.
97+
98+
```rust
99+
use async_openai::middleware::{retry::OpenAIRetryLayer, ReqwestService};
100+
use async_openai::middleware::rate_limit::RateLimitLayer;
101+
102+
let service = tower::ServiceBuilder::new()
103+
.layer(OpenAIRetryLayer::default())
104+
.layer(RateLimitLayer::per_minute(60))
105+
.service(ReqwestService::new(reqwest::Client::new()));
106+
```
107+
108+
`RateLimitLayer` should sit below `OpenAIRetryLayer` so retries are also throttled. On native targets, the layer provides local request-per-minute limiting and server-feedback backpressure from `x-ratelimit-remaining-requests` plus `x-ratelimit-reset-requests`. On WASM targets, it records local governor state but does not delay requests.
109+
94110
## Error Handling
95111

96112
`OpenAIError::Boxed` is available only when the `middleware` feature is enabled.

async-openai/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
//! ## Making requests
3030
//!
3131
//!```
32+
//!# #[cfg(feature = "responses")]
3233
//!# tokio_test::block_on(async {
3334
//! use async_openai::{Client, types::responses::{CreateResponseArgs}};
3435
//!
@@ -122,6 +123,7 @@
122123
//!
123124
//! For demonstration:
124125
//! ```
126+
//! # #[cfg(feature = "chat-completion")]
125127
//! # tokio_test::block_on(async {
126128
//! # use async_openai::Client;
127129
//! # use async_openai::traits::RequestOptionsBuilder;

async-openai/src/middleware/mod.rs

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,35 @@
103103
//!
104104
//! On native targets retries wait using `tokio::time::sleep`. On WASM retries are immediate.
105105
//!
106+
//! ## Rate limit layer
107+
//!
108+
//! Enable the `rate-limit` feature to use [`rate_limit::RateLimitLayer`].
109+
//! Place it below [`retry::OpenAIRetryLayer`] so retry attempts also pass
110+
//! through the limiter:
111+
//!
112+
//! ```no_run
113+
//! # #[cfg(feature = "rate-limit")]
114+
//! # fn main() {
115+
//! # use async_openai::{Client, config::OpenAIConfig};
116+
//! # use async_openai::middleware::{ReqwestService, retry::OpenAIRetryLayer};
117+
//! # use async_openai::middleware::rate_limit::RateLimitLayer;
118+
//! let service = tower::ServiceBuilder::new()
119+
//! .layer(OpenAIRetryLayer::default())
120+
//! .layer(RateLimitLayer::per_minute(60))
121+
//! .service(ReqwestService::new(reqwest::Client::new()));
122+
//!
123+
//! let client = Client::with_config(OpenAIConfig::default())
124+
//! .with_http_service(service);
125+
//! # }
126+
//! # #[cfg(not(feature = "rate-limit"))]
127+
//! # fn main() {}
128+
//! ```
129+
//!
130+
//! On native targets, the layer limits request RPM locally and applies
131+
//! server-feedback backpressure when `x-ratelimit-remaining-requests`
132+
//! reaches zero and `x-ratelimit-reset-requests` is parseable. On WASM
133+
//! targets, it records local governor state but does not delay requests.
134+
//!
106135
//! ## Error Handling
107136
//!
108137
//! `OpenAIError::Boxed` is available only when the `middleware` feature is enabled.
@@ -118,6 +147,10 @@
118147
//! trait bounds with or without middleware. JSON request bodies are serialized
119148
//! before they enter the replayable middleware request factory.
120149
150+
/// Request rate limiting middleware.
151+
#[cfg(feature = "rate-limit")]
152+
pub mod rate_limit;
153+
121154
/// Retry layers and policies for middleware.
122155
pub mod retry {
123156
#[doc(inline)]

0 commit comments

Comments
 (0)