yapi_ai/api/
images.rs

1//! Image generation over OpenRouter's chat completions endpoint.
2//!
3//! Port of `packages/ai/src/api/openrouter-images.ts` in pi `v1.0.0`.
4
5use indexmap::IndexMap;
6use serde_json::{Value, json};
7use tokio_util::sync::CancellationToken;
8use yapi_types::classify::{AssistantImages, ImagesContent, ImagesContext, OutcomeReason};
9use yapi_types::message::{Cost, ImageContent, TextContent, Usage};
10use yapi_types::model::ImageModel;
11use yapi_types::models::InputKind;
12
13use crate::http::{self, Failure};
14use crate::stream::{RequestHooks, now_ms};
15
16/// Options of one image request.
17#[derive(Clone, Debug, Default)]
18pub struct ImagesOptions {
19    /// Credential for the provider.
20    pub api_key: Option<String>,
21    /// Extra headers. `None` removes one.
22    pub headers: IndexMap<String, Option<String>>,
23    /// Retries of retryable failures, or none when absent, as in the SDK call.
24    pub max_retries: Option<u32>,
25    /// Largest server-requested retry delay to honor.
26    pub max_retry_delay_ms: Option<u64>,
27    /// Per-attempt timeout.
28    pub timeout_ms: Option<u64>,
29    /// Cancels the request.
30    pub cancel: CancellationToken,
31    /// See the request body, which `payload` may replace, and the response.
32    pub hooks: RequestHooks,
33}
34
35/// Generates images with `model`. Reports failures in the result, as pi's
36/// `generateImages()` does.
37pub async fn generate_images(
38    model: &ImageModel,
39    context: &ImagesContext,
40    options: &ImagesOptions,
41) -> AssistantImages {
42    let mut output = AssistantImages {
43        api: model.api.clone(),
44        provider: model.provider.clone(),
45        model: model.id.clone(),
46        output: Vec::new(),
47        response_id: None,
48        usage: None,
49        stop_reason: OutcomeReason::Stop,
50        error_message: None,
51        timestamp: now_ms(),
52    };
53    let result = match model.api.as_str() {
54        "openrouter-images" => openrouter(model, context, options, &mut output).await,
55        api => Err(format!("No API provider registered for api: {api}")),
56    };
57    if let Err(message) = result {
58        output.stop_reason = super::classify::failure_reason(&options.cancel);
59        output.error_message = Some(message);
60    }
61    output
62}
63
64/// pi's `imageErrorResult`: `model`'s result for a request that failed with
65/// `message`, aborted when `cancel` fired.
66pub fn error_result(
67    model: &ImageModel,
68    message: String,
69    cancel: &CancellationToken,
70) -> AssistantImages {
71    AssistantImages {
72        api: model.api.clone(),
73        provider: model.provider.clone(),
74        model: model.id.clone(),
75        output: Vec::new(),
76        response_id: None,
77        usage: None,
78        stop_reason: super::classify::failure_reason(cancel),
79        error_message: Some(message),
80        timestamp: now_ms(),
81    }
82}
83
84fn params(model: &ImageModel, context: &ImagesContext) -> Value {
85    let content: Vec<Value> = context
86        .input
87        .iter()
88        .map(|item| match item {
89            // Rust strings hold no lone surrogates, so pi's sanitizing is a no-op.
90            ImagesContent::Text(text) => json!({"type": "text", "text": text.text}),
91            ImagesContent::Image(image) => json!({
92                "type": "image_url",
93                "image_url": {"url": format!("data:{};base64,{}", image.mime_type, image.data)},
94            }),
95        })
96        .collect();
97    let modalities = if model.output.contains(&InputKind::Text) {
98        json!(["image", "text"])
99    } else {
100        json!(["image"])
101    };
102    json!({
103        "model": model.id,
104        "messages": [{"role": "user", "content": content}],
105        "stream": false,
106        "modalities": modalities,
107    })
108}
109
110fn usage(raw: &Value, model: &ImageModel) -> Usage {
111    let count = |value: &Value| value.as_u64().unwrap_or(0);
112    let prompt = count(&raw["prompt_tokens"]);
113    let cached = count(&raw["prompt_tokens_details"]["cached_tokens"]);
114    let cache_write = count(&raw["prompt_tokens_details"]["cache_write_tokens"]);
115    let cache_read = if cache_write > 0 {
116        cached.saturating_sub(cache_write)
117    } else {
118        cached
119    };
120    let input = prompt.saturating_sub(cache_read + cache_write);
121    let output = count(&raw["completion_tokens"]);
122    let rate = |per_million: f64, tokens: u64| (per_million / 1_000_000.0) * tokens as f64;
123    let mut cost = Cost {
124        input: rate(model.cost.input, input),
125        output: rate(model.cost.output, output),
126        cache_read: rate(model.cost.cache_read, cache_read),
127        cache_write: rate(model.cost.cache_write, cache_write),
128        total: 0.0,
129    };
130    cost.total = cost.input + cost.output + cost.cache_read + cost.cache_write;
131    Usage {
132        input,
133        output,
134        cache_read,
135        cache_write,
136        reasoning: None,
137        total_tokens: Some(input + output + cache_read + cache_write),
138        cost,
139        cache_write_1h: None,
140    }
141}
142
143async fn openrouter(
144    model: &ImageModel,
145    context: &ImagesContext,
146    options: &ImagesOptions,
147    output: &mut AssistantImages,
148) -> Result<(), String> {
149    let Some(api_key) = options.api_key.as_deref().filter(|key| !key.is_empty()) else {
150        return Err(format!("No API key for provider: {}", model.provider));
151    };
152    let body = yapi_types::json::stringify(&options.hooks.payload(params(model, context)).await);
153    let url = format!("{}/chat/completions", model.base_url.trim_end_matches('/'));
154    let mut headers = super::openai_headers(api_key, model.headers.as_ref());
155    headers.extend(&options.headers);
156    let response = http::post(
157        &url,
158        &headers,
159        &body,
160        options.timeout_ms,
161        options.max_retries.unwrap_or(0),
162        options.max_retry_delay_ms,
163        &options.cancel,
164    )
165    .await;
166    let response = match response {
167        Ok(response) => response,
168        Err(Failure::Status { status, body }) => {
169            return Err(http::openai_status_message(status, &body, None));
170        }
171        Err(other) => {
172            return Err(other
173                .plain_message(http::openai_sdk_message)
174                .unwrap_or_default());
175        }
176    };
177    options.hooks.response(&response).await;
178    let Some(bytes) = options.cancel.run_until_cancelled(response.bytes()).await else {
179        return Err(http::ABORTED_DURING_STREAM.into());
180    };
181    let bytes = bytes.map_err(|_| "Connection error.".to_owned())?;
182    let response: Value = serde_json::from_slice(&bytes).map_err(|err| err.to_string())?;
183    output.response_id = response["id"].as_str().map(str::to_owned);
184    if response["usage"].is_object() {
185        output.usage = Some(usage(&response["usage"], model));
186    }
187    let message = &response["choices"][0]["message"];
188    if let Some(text) = message["content"].as_str().filter(|text| !text.is_empty()) {
189        output.output.push(ImagesContent::Text(TextContent {
190            text: text.to_owned(),
191            text_signature: None,
192        }));
193    }
194    for image in message["images"].as_array().into_iter().flatten() {
195        let url = image["image_url"]
196            .as_str()
197            .or_else(|| image["image_url"]["url"].as_str())
198            .unwrap_or_default();
199        let Some(data_url) = url.strip_prefix("data:") else {
200            continue;
201        };
202        let Some((mime_type, data)) = data_url.split_once(";base64,") else {
203            continue;
204        };
205        if mime_type.is_empty() || mime_type.contains(';') || data.is_empty() {
206            continue;
207        }
208        output.output.push(ImagesContent::Image(ImageContent {
209            data: data.to_owned(),
210            mime_type: mime_type.to_owned(),
211        }));
212    }
213    Ok(())
214}
yapi_ai/api/images.rs

yapi_ai/api/
images.rs
