yapi_ai/
stream.rs

1//! Streaming one assistant message from a provider.
2
3use futures_util::future::BoxFuture;
4use indexmap::IndexMap;
5use tokio::sync::mpsc;
6use tokio_util::sync::CancellationToken;
7use yapi_types::event::AssistantMessageEvent;
8use yapi_types::message::{
9    AssistantMessage, ContentBlock, Message, StopReason, ThinkingLevel, Usage,
10};
11use yapi_types::model::Model;
12pub use yapi_types::settings::Transport;
13
14use crate::credentials::{ProviderEnv, provider_env_value};
15
16/// A wire API: turns a model, a transcript and options into a streamed message.
17///
18/// Implementations never fail the call itself. Every outcome, including request
19/// errors and cancellation, arrives on the stream, which ends with exactly one
20/// [`StreamEvent::Done`] or [`StreamEvent::Error`].
21pub trait Provider: Send + Sync {
22    /// The wire API id this provider speaks, such as `anthropic-messages`.
23    fn api(&self) -> &str;
24
25    /// Starts streaming a response. Must be called inside a tokio runtime.
26    fn stream(&self, request: Request) -> EventStream;
27}
28
29/// Everything one provider request needs.
30#[derive(Clone, Debug)]
31pub struct Request {
32    /// The model to call.
33    pub model: Model,
34    /// The transcript. The first message may be the system prompt.
35    pub messages: Vec<Message>,
36    /// Request options.
37    pub options: StreamOptions,
38}
39
40/// How long the provider should keep the prompt cached.
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub enum CacheRetention {
43    /// No caching.
44    None,
45    /// The provider's default retention.
46    Short,
47    /// Extended retention where offered, such as Anthropic's one hour.
48    Long,
49}
50
51impl CacheRetention {
52    /// The name pi uses in its options.
53    pub fn as_str(self) -> &'static str {
54        match self {
55            CacheRetention::None => "none",
56            CacheRetention::Short => "short",
57            CacheRetention::Long => "long",
58        }
59    }
60
61    /// Parses a name from [`CacheRetention::as_str`].
62    pub fn parse(name: &str) -> Option<CacheRetention> {
63        [
64            CacheRetention::None,
65            CacheRetention::Short,
66            CacheRetention::Long,
67        ]
68        .into_iter()
69        .find(|retention| retention.as_str() == name)
70    }
71}
72
73/// Request options shared by every wire API.
74#[derive(Clone, Debug, Default)]
75pub struct StreamOptions {
76    /// Credential for the provider. Some providers take headers instead.
77    pub api_key: Option<String>,
78    /// Extra headers. `None` removes a header the provider would otherwise send.
79    pub headers: IndexMap<String, Option<String>>,
80    /// Thinking effort. `None` and `Off` disable thinking where the API allows.
81    pub reasoning: Option<ThinkingLevel>,
82    /// Output token limit, or the model's limit when absent.
83    pub max_tokens: Option<u64>,
84    /// Sampling temperature, where the model and thinking mode allow it.
85    pub temperature: Option<f64>,
86    /// Stable id for provider-side session affinity and caching.
87    pub session_id: Option<String>,
88    /// Preferred transport for providers that offer more than one. The
89    /// provider's default when absent.
90    pub transport: Option<Transport>,
91    /// How long opening a WebSocket may take, in milliseconds. 0 waits
92    /// forever. The provider's default when absent.
93    pub websocket_connect_timeout_ms: Option<u64>,
94    /// pi's `timeoutMs`: how long a request may wait, in milliseconds, as
95    /// each provider applies it. The provider's default when absent. A
96    /// number, as the settings hold it, so a provider that validates it can
97    /// report an invalid one as pi does.
98    pub timeout_ms: Option<f64>,
99    /// Prompt cache retention. `PI_CACHE_RETENTION=long` or short when absent.
100    pub cache_retention: Option<CacheRetention>,
101    /// Token budgets per level for budget-based thinking.
102    pub thinking_budgets: ThinkingBudgets,
103    /// Retries for retryable HTTP errors before the first byte. 0 disables them.
104    pub max_retries: u32,
105    /// Largest server-requested retry delay to honor. 0 means no limit.
106    pub max_retry_delay_ms: Option<u64>,
107    /// Cancels the request, and the stream then ends with an `aborted` error.
108    pub cancel: CancellationToken,
109    /// Provider settings from the credential, such as a Cloudflare account id
110    /// or an AWS profile, read ahead of the process environment.
111    pub env: Option<ProviderEnv>,
112    /// Observers of the request.
113    pub hooks: RequestHooks,
114}
115
116/// Request headers. `None` removes a header the provider would send.
117pub type RequestHeaders = IndexMap<String, Option<String>>;
118
119/// A session's callback into its requests.
120pub type Hook<A, R> = std::sync::Arc<dyn Fn(A) -> BoxFuture<'static, R> + Send + Sync>;
121
122/// A provider's answer, before its body is read.
123#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct ProviderResponse {
125    /// The HTTP status.
126    pub status: u16,
127    /// The response headers.
128    pub headers: IndexMap<String, String>,
129}
130
131/// What a session observes of its requests: pi-ai's `onPayload`,
132/// `onResponse` and `onProviderStreamEvent` options, and pi's
133/// `transformHeaders`.
134#[derive(Clone, Default)]
135pub struct RequestHooks {
136    /// Sees each request body, as the wire API would send it, and returns
137    /// the body to send.
138    pub payload: Option<Hook<serde_json::Value, serde_json::Value>>,
139    /// Sees the headers of each request, credentials' included, and returns
140    /// the headers to send. `None` removes one.
141    pub headers: Option<Hook<RequestHeaders, RequestHeaders>>,
142    /// Sees each successful response.
143    pub response: Option<Hook<ProviderResponse, ()>>,
144    /// Sees each streamed event, as parsed, before it is normalized.
145    pub stream_event: Option<Hook<serde_json::Value, ()>>,
146}
147
148impl RequestHooks {
149    /// The body to send in place of `payload`.
150    pub async fn payload(&self, payload: serde_json::Value) -> serde_json::Value {
151        match &self.payload {
152            Some(hook) => hook(payload).await,
153            None => payload,
154        }
155    }
156
157    /// Passes `request`'s headers through the headers hook. Call it once
158    /// credentials are applied.
159    pub async fn prepare(request: &mut Request) {
160        if let Some(hook) = request.options.hooks.headers.clone() {
161            request.options.headers = hook(std::mem::take(&mut request.options.headers)).await;
162        }
163    }
164
165    /// Reports a successful `response`.
166    pub async fn response(&self, response: &reqwest::Response) {
167        if let Some(hook) = &self.response {
168            let headers = response
169                .headers()
170                .iter()
171                .filter_map(|(name, value)| {
172                    Some((name.as_str().to_owned(), value.to_str().ok()?.to_owned()))
173                })
174                .collect();
175            hook(ProviderResponse {
176                status: response.status().as_u16(),
177                headers,
178            })
179            .await;
180        }
181    }
182
183    /// Reports a streamed event.
184    pub async fn stream_event(&self, data: &serde_json::Value) {
185        if let Some(hook) = &self.stream_event {
186            hook(data.clone()).await;
187        }
188    }
189}
190
191impl std::fmt::Debug for RequestHooks {
192    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
193        f.debug_struct("RequestHooks")
194            .field("payload", &self.payload.is_some())
195            .field("headers", &self.headers.is_some())
196            .field("response", &self.response.is_some())
197            .field("stream_event", &self.stream_event.is_some())
198            .finish()
199    }
200}
201
202impl StreamOptions {
203    /// The cache retention to use: [`Self::requested_cache_retention`], else
204    /// `short`.
205    pub fn resolved_cache_retention(&self) -> CacheRetention {
206        self.requested_cache_retention()
207            .unwrap_or(CacheRetention::Short)
208    }
209
210    /// The option, else `long` when `PI_CACHE_RETENTION=long` in `env` or the
211    /// process environment.
212    pub fn requested_cache_retention(&self) -> Option<CacheRetention> {
213        self.cache_retention.or_else(|| {
214            (provider_env_value("PI_CACHE_RETENTION", self.env.as_ref()).as_deref() == Some("long"))
215                .then_some(CacheRetention::Long)
216        })
217    }
218
219    /// Whether `headers` sets `name` (case-insensitive) to a non-blank value.
220    pub fn has_header(&self, name: &str) -> bool {
221        self.headers.iter().any(|(key, value)| {
222            key.eq_ignore_ascii_case(name)
223                && value
224                    .as_deref()
225                    .is_some_and(|value| !value.trim().is_empty())
226        })
227    }
228}
229
230/// Overrides for budget-based thinking, in tokens.
231#[derive(Clone, Debug, Default, PartialEq, Eq)]
232pub struct ThinkingBudgets {
233    /// Default 1024.
234    pub minimal: Option<u64>,
235    /// Default 2048.
236    pub low: Option<u64>,
237    /// Default 8192.
238    pub medium: Option<u64>,
239    /// Default 16384. Also used for `xhigh` and `max`.
240    pub high: Option<u64>,
241}
242
243/// One step of a streamed response.
244#[derive(Clone, Debug, PartialEq)]
245pub enum StreamEvent {
246    /// The response started, with the message so far and no content yet.
247    Start(AssistantMessage),
248    /// A content change, with the cumulative usage after it.
249    Update {
250        /// What changed.
251        event: AssistantMessageEvent,
252        /// Usage so far.
253        usage: Usage,
254    },
255    /// Finished with `stop`, `length`, `toolUse` or `deferred`.
256    Done(AssistantMessage),
257    /// Failed or aborted. The message holds what arrived and `errorMessage`.
258    Error(AssistantMessage),
259}
260
261/// The receiving end of a provider stream.
262#[derive(Debug)]
263pub struct EventStream {
264    receiver: mpsc::UnboundedReceiver<StreamEvent>,
265}
266
267impl EventStream {
268    /// A stream fed by `sender`.
269    pub fn channel() -> (EventSender, EventStream) {
270        let (sender, receiver) = mpsc::unbounded_channel();
271        (EventSender { sender }, EventStream { receiver })
272    }
273
274    /// The next event, or `None` after the final one.
275    pub async fn next(&mut self) -> Option<StreamEvent> {
276        self.receiver.recv().await
277    }
278
279    /// Drains the stream and returns the final message.
280    pub async fn result(mut self) -> Option<AssistantMessage> {
281        while let Some(event) = self.next().await {
282            if let StreamEvent::Done(message) | StreamEvent::Error(message) = event {
283                return Some(message);
284            }
285        }
286        None
287    }
288}
289
290/// The sending end of a provider stream. Sends after the receiver is gone are
291/// dropped.
292#[derive(Clone, Debug)]
293pub struct EventSender {
294    sender: mpsc::UnboundedSender<StreamEvent>,
295}
296
297impl EventSender {
298    /// Sends an event.
299    pub fn send(&self, event: StreamEvent) {
300        let _ = self.sender.send(event);
301    }
302
303    /// Sends a content update with the message's current usage.
304    pub fn update(&self, message: &AssistantMessage, event: AssistantMessageEvent) {
305        self.send(StreamEvent::Update {
306            event,
307            usage: message.usage.clone(),
308        });
309    }
310
311    /// Appends `block` to `output`, sends its start event and returns its index.
312    pub(crate) fn start(&self, output: &mut AssistantMessage, block: ContentBlock) -> usize {
313        let content_index = output.content.len();
314        let event = match &block {
315            ContentBlock::Text(_) => Some(AssistantMessageEvent::TextStart { content_index }),
316            ContentBlock::Thinking(_) => {
317                Some(AssistantMessageEvent::ThinkingStart { content_index })
318            }
319            ContentBlock::ToolCall(call) => Some(AssistantMessageEvent::ToolcallStart {
320                content_index,
321                id: call.id.clone(),
322                tool_name: call.name.clone(),
323            }),
324            ContentBlock::Image(_) => None,
325        };
326        output.content.push(block);
327        if let Some(event) = event {
328            self.update(output, event);
329        }
330        content_index
331    }
332
333    /// Sends the delta event of the block at `content_index`, first appending
334    /// `delta` to a text or thinking block. A tool call's arguments are left to
335    /// the caller.
336    pub(crate) fn delta(&self, output: &mut AssistantMessage, content_index: usize, delta: &str) {
337        let event = match output.content.get_mut(content_index) {
338            Some(ContentBlock::Text(text)) => {
339                text.text.push_str(delta);
340                AssistantMessageEvent::TextDelta {
341                    content_index,
342                    delta: delta.to_owned(),
343                }
344            }
345            Some(ContentBlock::Thinking(thinking)) => {
346                thinking.thinking.push_str(delta);
347                AssistantMessageEvent::ThinkingDelta {
348                    content_index,
349                    delta: delta.to_owned(),
350                }
351            }
352            Some(ContentBlock::ToolCall(_)) => AssistantMessageEvent::ToolcallDelta {
353                content_index,
354                delta: delta.to_owned(),
355            },
356            Some(ContentBlock::Image(_)) | None => return,
357        };
358        self.update(output, event);
359    }
360
361    /// Sends the end event of the block at `content_index`, with its final
362    /// content.
363    pub(crate) fn end(&self, output: &AssistantMessage, content_index: usize) {
364        let event = match output.content.get(content_index) {
365            Some(ContentBlock::Text(text)) => AssistantMessageEvent::TextEnd {
366                content_index,
367                content: text.text.clone(),
368            },
369            Some(ContentBlock::Thinking(thinking)) => AssistantMessageEvent::ThinkingEnd {
370                content_index,
371                content: thinking.thinking.clone(),
372            },
373            Some(ContentBlock::ToolCall(call)) => AssistantMessageEvent::ToolcallEnd {
374                content_index,
375                tool_call: call.clone(),
376            },
377            Some(ContentBlock::Image(_)) | None => return,
378        };
379        self.update(output, event);
380    }
381}
382
383/// An empty assistant message for `model`, as providers start one.
384pub fn new_output(model: &Model, now_ms: u64) -> AssistantMessage {
385    AssistantMessage {
386        content: Vec::new(),
387        api: model.api.clone(),
388        provider: model.provider.clone(),
389        model: model.id.clone(),
390        provider_thinking_level: None,
391        usage: Usage {
392            total_tokens: Some(0),
393            ..Usage::default()
394        },
395        stop_reason: StopReason::Pending,
396        timestamp: now_ms,
397        response_id: None,
398        response_model: None,
399        raw_stop_reason: None,
400        deferred: None,
401        error_message: None,
402        diagnostics: None,
403        end_turn: None,
404        thinking_level: None,
405    }
406}
407
408/// Ends a stream with an error: `aborted` when the request was cancelled.
409pub fn send_error(
410    sender: &EventSender,
411    mut output: AssistantMessage,
412    cancel: &CancellationToken,
413    message: String,
414) {
415    output.stop_reason = if cancel.is_cancelled() {
416        StopReason::Aborted
417    } else {
418        StopReason::Error
419    };
420    output.error_message = Some(message);
421    sender.send(StreamEvent::Error(output));
422}
423
424/// pi's end-of-stream check: an error when `cancel` fired, `pending` when no
425/// stop reason arrived, and the message's error, or a generic one, when it
426/// stopped with `error` or `aborted`.
427pub(crate) fn check_complete(
428    output: &AssistantMessage,
429    cancel: &CancellationToken,
430    pending: &str,
431) -> Result<(), String> {
432    if cancel.is_cancelled() {
433        return Err(crate::http::ABORTED_DURING_STREAM.to_owned());
434    }
435    match output.stop_reason {
436        StopReason::Pending => Err(pending.to_owned()),
437        StopReason::Aborted | StopReason::Error => Err(output
438            .error_message
439            .clone()
440            .filter(|message| !message.is_empty())
441            .unwrap_or_else(|| "An unknown error occurred".to_owned())),
442        _ => Ok(()),
443    }
444}
445
446pub use yapi_types::time::now_ms;
yapi_ai/stream.rs

yapi_ai/
stream.rs
