|
| 1 | +//! Example: host a perfect-replica of Anthropic's `POST /v1/messages`. |
| 2 | +//! |
| 3 | +//! Exercises both branches of the typed response enum: |
| 4 | +//! - `body.stream == Some(true)` → `OkStream(Sse<...>)` |
| 5 | +//! - otherwise → `Ok(Message)` (single JSON body) |
| 6 | +//! |
| 7 | +//! NOTE: Anthropic's published spec declares only `application/json` |
| 8 | +//! on the 200 response. This example pulls in a small overlay |
| 9 | +//! (`sse-overlay.json`) via the generator's `schema_extensions` |
| 10 | +//! mechanism, which declares `text/event-stream` on the 200 so the |
| 11 | +//! generated trait gets the `OkStream` variant. |
| 12 | +//! |
| 13 | +//! Run: |
| 14 | +//! 1. `cargo run -p openapi-to-rust -- generate \ |
| 15 | +//! --config examples/server-anthropic-messages/openapi-to-rust.toml` |
| 16 | +//! 2. `cargo run --manifest-path examples/server-anthropic-messages/Cargo.toml` |
| 17 | +//! 3. Unary: |
| 18 | +//! `curl -s http://127.0.0.1:3001/v1/messages \ |
| 19 | +//! -H 'content-type: application/json' \ |
| 20 | +//! -d '{"model":"claude-x","max_tokens":50, |
| 21 | +//! "messages":[{"role":"user","content":"hi"}]}'` |
| 22 | +//! 4. SSE: |
| 23 | +//! `curl -N -s http://127.0.0.1:3001/v1/messages \ |
| 24 | +//! -H 'content-type: application/json' \ |
| 25 | +//! -d '{"model":"claude-x","max_tokens":50,"stream":true, |
| 26 | +//! "messages":[{"role":"user","content":"hi"}]}'` |
| 27 | +
|
| 28 | +pub mod gen; |
| 29 | + |
| 30 | +use axum::response::sse::{Event, KeepAlive, Sse}; |
| 31 | +use futures_util::stream; |
| 32 | +use gen::CreateMessageParams; |
| 33 | +use gen::server::{MessagesPostResponse, ServerApi, ServerEventStream, server_api_router}; |
| 34 | +use std::convert::Infallible; |
| 35 | +use std::time::Duration; |
| 36 | + |
| 37 | +#[derive(Clone)] |
| 38 | +struct AppState; |
| 39 | + |
| 40 | +#[axum::async_trait] |
| 41 | +impl ServerApi for AppState { |
| 42 | + async fn messages_post(&self, body: CreateMessageParams) -> MessagesPostResponse { |
| 43 | + if body.stream == Some(true) { |
| 44 | + messages_streaming() |
| 45 | + } else { |
| 46 | + messages_unary() |
| 47 | + } |
| 48 | + } |
| 49 | +} |
| 50 | + |
| 51 | +fn messages_unary() -> MessagesPostResponse { |
| 52 | + let msg = gen::Message { |
| 53 | + container: None, |
| 54 | + content: vec![gen::ContentBlock::TextBlock(gen::ResponseTextBlock { |
| 55 | + citations: None, |
| 56 | + text: "hello (unary)".into(), |
| 57 | + })], |
| 58 | + id: "msg_demo".into(), |
| 59 | + model: gen::Model::Custom("claude-demo".into()), |
| 60 | + role: gen::MessageRole::Assistant, |
| 61 | + stop_details: None, |
| 62 | + stop_reason: None, |
| 63 | + stop_sequence: None, |
| 64 | + r#type: gen::MessageType::Message, |
| 65 | + usage: gen::Usage { |
| 66 | + cache_creation: None, |
| 67 | + cache_creation_input_tokens: None, |
| 68 | + cache_read_input_tokens: None, |
| 69 | + inference_geo: None, |
| 70 | + input_tokens: 0, |
| 71 | + output_tokens: 0, |
| 72 | + server_tool_use: None, |
| 73 | + service_tier: None, |
| 74 | + }, |
| 75 | + }; |
| 76 | + MessagesPostResponse::Ok(msg) |
| 77 | +} |
| 78 | + |
| 79 | +fn messages_streaming() -> MessagesPostResponse { |
| 80 | + // The real Anthropic stream emits message_start → |
| 81 | + // content_block_start → content_block_delta* → content_block_stop |
| 82 | + // → message_delta → message_stop. The example fires a short |
| 83 | + // subset; production code mirrors the full sequence from the |
| 84 | + // upstream model. |
| 85 | + let events = stream::iter(vec![ |
| 86 | + sse_event("message_start", r#"{"type":"message_start","message":{"id":"msg_demo"}}"#), |
| 87 | + sse_event( |
| 88 | + "content_block_start", |
| 89 | + r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, |
| 90 | + ), |
| 91 | + sse_event( |
| 92 | + "content_block_delta", |
| 93 | + r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello "}}"#, |
| 94 | + ), |
| 95 | + sse_event( |
| 96 | + "content_block_delta", |
| 97 | + r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"world"}}"#, |
| 98 | + ), |
| 99 | + sse_event("content_block_stop", r#"{"type":"content_block_stop","index":0}"#), |
| 100 | + sse_event("message_stop", r#"{"type":"message_stop"}"#), |
| 101 | + ]); |
| 102 | + let pinned: ServerEventStream = Box::pin(events); |
| 103 | + MessagesPostResponse::OkStream( |
| 104 | + Sse::new(pinned).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))), |
| 105 | + ) |
| 106 | +} |
| 107 | + |
| 108 | +fn sse_event(name: &str, data: &str) -> Result<Event, Infallible> { |
| 109 | + Ok(Event::default().event(name).data(data)) |
| 110 | +} |
| 111 | + |
| 112 | +#[tokio::main] |
| 113 | +async fn main() { |
| 114 | + let app = server_api_router(AppState); |
| 115 | + let listener = tokio::net::TcpListener::bind("127.0.0.1:3001").await.unwrap(); |
| 116 | + println!("listening on http://{}", listener.local_addr().unwrap()); |
| 117 | + axum::serve(listener, app).await.unwrap(); |
| 118 | +} |
| 119 | + |
| 120 | +#[cfg(test)] |
| 121 | +mod tests { |
| 122 | + use super::*; |
| 123 | + |
| 124 | + fn make_body(stream: Option<bool>) -> CreateMessageParams { |
| 125 | + let mut json = serde_json::json!({ |
| 126 | + "model": "claude-x", |
| 127 | + "max_tokens": 50, |
| 128 | + "messages": [{"role": "user", "content": "hi"}], |
| 129 | + }); |
| 130 | + if let Some(s) = stream { |
| 131 | + json["stream"] = serde_json::Value::Bool(s); |
| 132 | + } |
| 133 | + serde_json::from_value(json).expect("minimal CreateMessageParams must deserialize") |
| 134 | + } |
| 135 | + |
| 136 | + #[tokio::test] |
| 137 | + async fn unary_path_returns_ok_variant() { |
| 138 | + let r = AppState.messages_post(make_body(None)).await; |
| 139 | + assert!(matches!(r, MessagesPostResponse::Ok(_))); |
| 140 | + } |
| 141 | + |
| 142 | + #[tokio::test] |
| 143 | + async fn stream_path_returns_ok_stream_variant() { |
| 144 | + let r = AppState.messages_post(make_body(Some(true))).await; |
| 145 | + assert!(matches!(r, MessagesPostResponse::OkStream(_))); |
| 146 | + } |
| 147 | +} |
0 commit comments