yapi_ai/api/
mod.rs

1//! Wire API implementations, one per API id, and dispatch by model.
2
3pub mod anthropic;
4pub mod bedrock;
5pub mod classify;
6mod codex_websocket;
7pub mod google;
8pub mod images;
9pub mod mistral;
10pub mod openai_completions;
11pub mod openai_responses;
12pub mod pi_messages;
13
14use std::sync::Arc;
15
16pub use codex_websocket::{
17    DebugStats as CodexWebSocketStats, close_sessions as close_codex_websockets,
18    debug_stats as codex_websocket_stats, reset_debug_stats as reset_codex_websocket_stats,
19};
20
21use futures_util::future::BoxFuture;
22use indexmap::IndexMap;
23use yapi_types::classify::{AssistantImages, ClassifierContext, ClassifierResult, ImagesContext};
24use yapi_types::message::{Content, ContentBlock, Message};
25use yapi_types::model::{ClassifierModel, ImageModel};
26
27use crate::http::Headers;
28use crate::stream::{EventStream, Provider, Request, new_output, now_ms, send_error};
29
30/// The built-in wire APIs. Each streams from a task that runs its module's
31/// `run`.
32#[derive(Clone, Copy, Debug)]
33enum Builtin {
34    Anthropic,
35    Completions,
36    Responses,
37    AzureResponses,
38    CodexResponses,
39    Gemini,
40    Vertex,
41    Mistral,
42    Bedrock,
43    PiMessages,
44}
45
46impl Builtin {
47    const ALL: [Builtin; 10] = [
48        Builtin::Anthropic,
49        Builtin::Completions,
50        Builtin::Responses,
51        Builtin::AzureResponses,
52        Builtin::CodexResponses,
53        Builtin::Gemini,
54        Builtin::Vertex,
55        Builtin::Mistral,
56        Builtin::Bedrock,
57        Builtin::PiMessages,
58    ];
59}
60
61impl Provider for Builtin {
62    fn api(&self) -> &str {
63        match self {
64            Builtin::Anthropic => "anthropic-messages",
65            Builtin::Completions => "openai-completions",
66            Builtin::Responses => "openai-responses",
67            Builtin::AzureResponses => "azure-openai-responses",
68            Builtin::CodexResponses => "openai-codex-responses",
69            Builtin::Gemini => "google-generative-ai",
70            Builtin::Vertex => "google-vertex",
71            Builtin::Mistral => "mistral-conversations",
72            Builtin::Bedrock => "bedrock-converse-stream",
73            Builtin::PiMessages => "pi-messages",
74        }
75    }
76
77    fn stream(&self, request: Request) -> EventStream {
78        use google::Flavor as Google;
79        use openai_responses::Flavor as Responses;
80        let (sender, stream) = EventStream::channel();
81        match self {
82            Builtin::Anthropic => tokio::spawn(anthropic::run(request, sender)),
83            Builtin::Completions => tokio::spawn(openai_completions::run(request, sender)),
84            Builtin::Responses => {
85                tokio::spawn(openai_responses::run(request, sender, Responses::OpenAi))
86            }
87            Builtin::AzureResponses => {
88                tokio::spawn(openai_responses::run(request, sender, Responses::Azure))
89            }
90            Builtin::CodexResponses => {
91                tokio::spawn(openai_responses::run(request, sender, Responses::Codex))
92            }
93            Builtin::Gemini => tokio::spawn(google::run(request, sender, Google::Gemini)),
94            Builtin::Vertex => tokio::spawn(google::run(request, sender, Google::Vertex)),
95            Builtin::Mistral => tokio::spawn(mistral::run(request, sender)),
96            Builtin::Bedrock => tokio::spawn(bedrock::run(request, sender)),
97            Builtin::PiMessages => tokio::spawn(pi_messages::run(request, sender)),
98        };
99        stream
100    }
101}
102
103/// The built-in implementation of a wire API id.
104pub fn builtin(api: &str) -> Option<Arc<dyn Provider>> {
105    let builtin = Builtin::ALL
106        .into_iter()
107        .find(|builtin| builtin.api() == api)?;
108    Some(Arc::new(builtin))
109}
110
111/// pi's `normalizeToolCallId` for Anthropic, Bedrock and Google:
112/// [`sanitize_id_part`], at most 64 characters.
113pub(crate) fn normalize_tool_call_id(id: &str) -> String {
114    sanitize_id_part(id).chars().take(64).collect()
115}
116
117/// Replaces every UTF-16 unit outside `[a-zA-Z0-9_-]` with `_`, as pi's id
118/// sanitizers do. A character outside the Basic Multilingual Plane (BMP)
119/// becomes two underscores.
120pub(crate) fn sanitize_id_part(text: &str) -> String {
121    text.chars()
122        .flat_map(|c| {
123            let keep = c.is_ascii_alphanumeric() || c == '_' || c == '-';
124            let (c, count) = if keep { (c, 1) } else { ('_', c.len_utf16()) };
125            std::iter::repeat_n(c, count)
126        })
127        .collect()
128}
129
130/// A stream that ends at once with `message` as its error.
131pub fn failed_stream(
132    model: &yapi_types::model::Model,
133    cancel: &tokio_util::sync::CancellationToken,
134    message: String,
135) -> EventStream {
136    let (sender, stream) = EventStream::channel();
137    send_error(&sender, new_output(model, now_ms()), cancel, message);
138    stream
139}
140
141/// The OpenAI SDK's request headers: the bearer `api_key`, JSON in and out, then
142/// the model's configured headers.
143pub(crate) fn openai_headers(
144    api_key: &str,
145    model_headers: Option<&IndexMap<String, String>>,
146) -> Headers {
147    let mut headers = Headers::default();
148    headers.set("authorization", Some(format!("Bearer {api_key}")));
149    headers.set("content-type", Some("application/json"));
150    headers.set("accept", Some("application/json"));
151    headers.extend_model(model_headers);
152    headers
153}
154
155/// GitHub Copilot's per-request headers: whether the user or the agent
156/// initiated the request, and the vision flag when images are sent. Port of
157/// `github-copilot-headers.ts`.
158pub(crate) fn copilot_headers(messages: &[Message]) -> Vec<(&'static str, String)> {
159    let initiator = match messages.last() {
160        Some(Message::User(_)) | None => "user",
161        Some(_) => "agent",
162    };
163    let has_image = |blocks: &[ContentBlock]| {
164        blocks
165            .iter()
166            .any(|block| matches!(block, ContentBlock::Image(_)))
167    };
168    let vision = messages.iter().any(|message| match message {
169        Message::User(user) => {
170            matches!(&user.content, Content::Blocks(blocks) if has_image(blocks))
171        }
172        Message::ToolResult(result) => has_image(&result.content),
173        _ => false,
174    });
175    let mut headers = vec![
176        ("X-Initiator", initiator.to_owned()),
177        ("Openai-Intent", "conversation-edits".to_owned()),
178    ];
179    if vision {
180        headers.push(("Copilot-Vision-Request", "true".to_owned()));
181    }
182    headers
183}
184
185/// The classifier and image APIs that a provider registered by an extension
186/// implements (pi-ai's `classify` and `generateImages` of the provider).
187pub trait ModelOperations: Send + Sync {
188    /// Classifies `context` with `model`. Failures are in the result.
189    fn classify<'a>(
190        &'a self,
191        model: &'a ClassifierModel,
192        context: &'a ClassifierContext,
193        options: &'a classify::ClassifyOptions,
194    ) -> BoxFuture<'a, ClassifierResult>;
195
196    /// Generates images with `model`. Failures are in the result.
197    fn generate_images<'a>(
198        &'a self,
199        model: &'a ImageModel,
200        context: &'a ImagesContext,
201        options: &'a images::ImagesOptions,
202    ) -> BoxFuture<'a, AssistantImages>;
203}
204
205/// Wire APIs by id: the built-ins plus any registered by extensions.
206#[derive(Clone, Default)]
207pub struct Apis {
208    registered: IndexMap<String, Arc<dyn Provider>>,
209    /// Streams of extension providers for their own models, by provider id
210    /// and API.
211    providers: IndexMap<(String, String), Arc<dyn Provider>>,
212}
213
214impl Apis {
215    /// Registers or replaces an implementation.
216    pub fn register(&mut self, provider: Arc<dyn Provider>) {
217        self.registered.insert(provider.api().to_owned(), provider);
218    }
219
220    /// Streams the models of `provider` whose API is `implementation`'s with
221    /// `implementation`, ahead of every other implementation of that API
222    /// (pi's `streamSimple` of an extension's provider, which may serve
223    /// several APIs).
224    pub fn register_for(&mut self, provider: &str, implementation: Arc<dyn Provider>) {
225        let key = (provider.to_owned(), implementation.api().to_owned());
226        self.providers.insert(key, implementation);
227    }
228
229    /// Drops the streams [`Apis::register_for`] added for `provider`.
230    pub fn unregister_for(&mut self, provider: &str) {
231        self.providers.retain(|(owner, _), _| owner != provider);
232    }
233
234    /// The implementation for an API id.
235    pub fn get(&self, api: &str) -> Option<Arc<dyn Provider>> {
236        self.registered.get(api).cloned().or_else(|| builtin(api))
237    }
238
239    /// Streams a request with the implementation of its model's API, adding the
240    /// headers some providers expect. An unknown API ends the stream with an error.
241    pub fn stream(&self, mut request: Request) -> EventStream {
242        if matches!(request.model.provider.as_str(), "opencode" | "opencode-go")
243            && let Some(session_id) = request.options.session_id.clone()
244            && !request
245                .options
246                .headers
247                .keys()
248                .any(|name| name.eq_ignore_ascii_case("x-opencode-session"))
249        {
250            request
251                .options
252                .headers
253                .insert("x-opencode-session".into(), Some(session_id));
254        }
255        if matches!(
256            request.model.provider.as_str(),
257            "cloudflare-ai-gateway" | "cloudflare-workers-ai"
258        ) {
259            request.model.base_url = crate::key_auth::cloudflare_base_url(
260                &request.model.base_url,
261                request.options.env.as_ref(),
262            );
263        }
264        let key = (request.model.provider.clone(), request.model.api.clone());
265        let own = self.providers.get(&key).cloned();
266        match own.or_else(|| self.get(&request.model.api)) {
267            Some(provider) => provider.stream(request),
268            None => {
269                let message = format!("No API provider registered for api: {}", request.model.api);
270                failed_stream(&request.model, &request.options.cancel, message)
271            }
272        }
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    /// A provider's own stream serves only its models of that API.
281    #[tokio::test]
282    async fn streams_serve_own_models() {
283        let mut apis = Apis::default();
284        apis.register_for(
285            "echo",
286            Arc::new(crate::faux::Faux::new([crate::faux::Response::text("hi")])),
287        );
288        let request = |provider: &str| Request {
289            model: serde_json::from_value(serde_json::json!({
290                "id": "m", "name": "M", "api": "faux", "provider": provider, "baseUrl": "",
291                "reasoning": false, "input": ["text"],
292                "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
293                "contextWindow": 1000, "maxTokens": 100,
294            }))
295            .unwrap(),
296            messages: Vec::new(),
297            options: crate::stream::StreamOptions::default(),
298        };
299        let own = apis.stream(request("echo")).result().await.unwrap();
300        assert_eq!(own.error_message, None);
301        let other = apis.stream(request("other")).result().await.unwrap();
302        assert_eq!(
303            other.error_message.as_deref(),
304            Some("No API provider registered for api: faux")
305        );
306    }
307
308    #[test]
309    fn tool_call_ids_utf16() {
310        // pi's regex replaces each half of a surrogate pair.
311        assert_eq!(normalize_tool_call_id("call|a😀b"), "call_a__b");
312        assert_eq!(normalize_tool_call_id(&"x".repeat(70)), "x".repeat(64));
313    }
314}
yapi_ai/api/mod.rs

yapi_ai/api/
mod.rs
