oino_core/
provider.rs

1//! Provider dispatch: the [`StreamFn`] that talks only to provider extensions.
2//!
3//! Core has no provider code (charter). A request for a model goes to the provider component
4//! whose manifest `provides.apis` contains the model's `api`; its [`ProviderEvent`]s are
5//! assembled into `AssistantMessageEvent`s with live partials. Every stream ends with
6//! `done` or `error`.
7
8use std::sync::Arc;
9
10use futures::StreamExt;
11use futures::stream::{self, BoxStream};
12use oino_agent::{StreamFn, StreamRequest, now_ms};
13use oino_ai::{
14    AssistantMessage, AssistantMessageEvent, ErrorReason, JsonValue, PartialAssembler,
15    ProviderEvent, StopReason,
16};
17use serde_json::json;
18use tokio_util::sync::CancellationToken;
19
20use crate::extensions::api::ExtensionCallError;
21use crate::extensions::runner::{ExtensionSet, LoadedExtension};
22
23/// Error text when no component implements a model's API.
24#[must_use]
25pub fn no_provider_message(api: &str) -> String {
26    format!("No provider extension implements API \"{api}\"")
27}
28
29/// [`StreamFn`] over the provider components of an [`ExtensionSet`].
30#[derive(Clone, Debug)]
31pub struct ProviderStream {
32    extensions: Arc<ExtensionSet>,
33}
34
35impl ProviderStream {
36    /// Dispatches to the providers in `extensions`.
37    #[must_use]
38    pub const fn new(extensions: Arc<ExtensionSet>) -> Self {
39        Self { extensions }
40    }
41
42    fn component(&self, api: &str) -> Option<Arc<LoadedExtension>> {
43        self.extensions
44            .extensions()
45            .iter()
46            .find(|extension| {
47                extension.instance.provider().is_some()
48                    && extension.state().apis.iter().any(|a| a == api)
49            })
50            .cloned()
51    }
52}
53
54/// The provider request JSON (`{model, context, options, apiKey?}`).
55#[must_use]
56pub fn request_json(request: &StreamRequest) -> JsonValue {
57    let mut value = json!({
58        "model": request.model,
59        "context": request.context,
60        "options": request.options,
61    });
62    if let Some(key) = &request.api_key {
63        value["apiKey"] = json!(key);
64    }
65    value
66}
67
68fn failure(
69    mut message: AssistantMessage,
70    error: &str,
71    cancel: &CancellationToken,
72) -> AssistantMessageEvent {
73    let aborted = cancel.is_cancelled();
74    message.stop_reason = if aborted {
75        StopReason::Aborted
76    } else {
77        StopReason::Error
78    };
79    message.error_message = Some(error.to_owned());
80    AssistantMessageEvent::Error {
81        reason: if aborted {
82            ErrorReason::Aborted
83        } else {
84            ErrorReason::Error
85        },
86        error: message,
87    }
88}
89
90enum Phase {
91    Connect(Arc<LoadedExtension>, JsonValue),
92    Streaming(BoxStream<'static, Result<ProviderEvent, ExtensionCallError>>),
93    Done,
94}
95
96impl StreamFn for ProviderStream {
97    fn stream(&self, request: StreamRequest) -> BoxStream<'static, AssistantMessageEvent> {
98        let model = &request.model;
99        let template = AssistantMessage::empty(&model.api, &model.provider, &model.id, now_ms());
100        let cancel = request.cancel.clone();
101        let Some(component) = self.component(&model.api) else {
102            let event = failure(template, &no_provider_message(&model.api), &cancel);
103            return Box::pin(stream::once(async move { event }));
104        };
105        let body = request_json(&request);
106        let assembler = PartialAssembler::new(template);
107        let state = (Phase::Connect(component, body), assembler, cancel);
108        Box::pin(
109            stream::unfold(state, |(phase, mut assembler, cancel)| async move {
110                match phase {
111                    Phase::Done => None,
112                    Phase::Connect(component, body) => {
113                        let Some(provider) = component.instance.provider() else {
114                            let event = failure(
115                                assembler.partial().clone(),
116                                "Provider unavailable",
117                                &cancel,
118                            );
119                            return Some((vec![event], (Phase::Done, assembler, cancel)));
120                        };
121                        match provider.stream(body, cancel.clone()).await {
122                            Ok(events) => {
123                                Some((Vec::new(), (Phase::Streaming(events), assembler, cancel)))
124                            }
125                            Err(error) => {
126                                let event =
127                                    failure(assembler.partial().clone(), &error.message, &cancel);
128                                Some((vec![event], (Phase::Done, assembler, cancel)))
129                            }
130                        }
131                    }
132                    Phase::Streaming(mut events) => {
133                        let next = tokio::select! {
134                            biased;
135                            () = cancel.cancelled() => None,
136                            next = events.next() => Some(next),
137                        };
138                        match next {
139                            None => {
140                                let event = failure(
141                                    assembler.partial().clone(),
142                                    "Request was aborted",
143                                    &cancel,
144                                );
145                                Some((vec![event], (Phase::Done, assembler, cancel)))
146                            }
147                            Some(Some(Ok(event))) => {
148                                let terminal = matches!(
149                                    event,
150                                    ProviderEvent::Done { .. } | ProviderEvent::Error { .. }
151                                );
152                                let out = assembler.apply(event);
153                                let phase = if terminal {
154                                    Phase::Done
155                                } else {
156                                    Phase::Streaming(events)
157                                };
158                                Some((vec![out], (phase, assembler, cancel)))
159                            }
160                            Some(Some(Err(error))) => {
161                                let event =
162                                    failure(assembler.partial().clone(), &error.message, &cancel);
163                                Some((vec![event], (Phase::Done, assembler, cancel)))
164                            }
165                            Some(None) => {
166                                let event = failure(
167                                    assembler.partial().clone(),
168                                    "Provider stream ended without a final event",
169                                    &cancel,
170                                );
171                                Some((vec![event], (Phase::Done, assembler, cancel)))
172                            }
173                        }
174                    }
175                }
176            })
177            .flat_map(stream::iter),
178        )
179    }
180}
oino_core/provider.rs

oino_core/
provider.rs
