feat: add Middleware trait, Pipeline, StructuredOutputMiddleware

2026-06-06 20:30:41 +07:00
parent 732ca01cc7
commit 6449bb5f4a
2 changed files with 245 additions and 0 deletions
@@ -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;
@@ -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);
+    }
+}