use crate::chat::stream::{Event, EventStream};
use crate::chat::{ChatRequest, ChatResponse};
use crate::error::{Error, Result};
use crate::model::{ModelCapabilities, ModelId, Reach};
use crate::observe;
use crate::provider::{Access, Provider};
use crate::registry::Registry;
use crate::secret::Secret;
use crate::transport::{ByteStream, HttpRequest, HttpTransport};
use async_trait::async_trait;
use futures_core::Stream;
use serde_json::Value;
use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
#[derive(Debug, Clone, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct SseFrame {
pub event: String,
pub data: String,
}
impl SseFrame {
pub fn json(&self) -> Option<Value> {
serde_json::from_str(&self.data).ok()
}
}
pub trait Protocol: Send + Sync {
fn id(&self) -> &str;
fn chat_url(&self, base_url: &str, model: &ModelId) -> String;
fn headers(&self, key: &Secret) -> Result<Vec<(String, String)>>;
fn body(&self, request: &ChatRequest) -> Result<Value>;
fn read(&self, body: &Value, asked_for: &ModelId) -> Result<ChatResponse>;
fn catalogue_url(&self, _base_url: &str) -> Option<String> {
None
}
fn read_catalogue(&self, _body: &Value) -> Result<Vec<ModelId>> {
Err(Error::Unsupported(format!(
"{} has no model catalogue",
self.id()
)))
}
fn stream_body(&self, _request: &ChatRequest) -> Result<Option<Value>> {
Ok(None)
}
fn stream_url(&self, base_url: &str, model: &ModelId) -> String {
self.chat_url(base_url, model)
}
fn read_event(&self, _frame: &SseFrame, _asked_for: &ModelId) -> Result<Vec<Event>> {
Ok(Vec::new())
}
}
pub struct ApiProvider<P: Protocol> {
protocol: P,
transport: Arc<dyn HttpTransport>,
key: Secret,
base_url: String,
reach: Reach,
registry: Arc<Registry>,
}
impl<P: Protocol> ApiProvider<P> {
pub fn new(
protocol: P,
base_url: impl Into<String>,
transport: Arc<dyn HttpTransport>,
key: Secret,
reach: Reach,
registry: Arc<Registry>,
) -> Self {
Self {
protocol,
transport,
key,
base_url: base_url.into().trim_end_matches('/').to_string(),
reach,
registry,
}
}
pub fn reach(&self) -> Reach {
self.reach
}
pub fn protocol(&self) -> &P {
&self.protocol
}
async fn call(&self, request: HttpRequest) -> Result<Value> {
let mut request = request;
request.headers = self.protocol.headers(&self.key)?;
let response = self.transport.send(request).await?;
response.check()?;
serde_json::from_slice(&response.body)
.map_err(|e| Error::Unreadable(format!("the reply was not JSON: {e}")))
}
}
#[async_trait]
impl<P: Protocol> Provider for ApiProvider<P> {
fn id(&self) -> &str {
self.protocol.id()
}
fn capabilities(&self, model: &ModelId) -> Option<ModelCapabilities> {
self.registry.capabilities(model)
}
async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
let span = observe::calling(self.protocol.id(), &request.model, self.reach);
observe::inside(span.clone(), async move {
let body = serde_json::to_vec(&self.protocol.body(&request)?)
.map_err(|e| Error::InvalidRequest(format!("building the request: {e}")))?;
let parsed = self
.call(HttpRequest::new(
self.protocol.chat_url(&self.base_url, &request.model),
body,
))
.await?;
let reply = self.protocol.read(&parsed, &request.model)?;
observe::measured(&span, reply.usage.coverage());
Ok(reply)
})
.await
}
async fn catalogue(&self) -> Result<Vec<ModelId>> {
let Some(url) = self.protocol.catalogue_url(&self.base_url) else {
return Err(Error::Unsupported(format!(
"{} has no model catalogue",
self.protocol.id()
)));
};
let parsed = self.call(HttpRequest::get(url)).await?;
self.protocol.read_catalogue(&parsed)
}
async fn stream(&self, request: ChatRequest) -> Result<EventStream<'_>> {
let Some(body) = self.protocol.stream_body(&request)? else {
return Ok(crate::provider::replay_stream(&self.chat(request).await?));
};
let body = serde_json::to_vec(&body)
.map_err(|e| Error::InvalidRequest(format!("building the request: {e}")))?;
let mut http = HttpRequest::new(
self.protocol.stream_url(&self.base_url, &request.model),
body,
);
http.headers = self.protocol.headers(&self.key)?;
Ok(Box::pin(Frames {
bytes: self.transport.send_streaming(http).await?,
protocol: &self.protocol,
asked_for: request.model,
buffer: Vec::new(),
ready: VecDeque::new(),
done: false,
}))
}
async fn validate(&self, model: &ModelId) -> Access {
let id = self.protocol.id();
let Some(url) = self.protocol.catalogue_url(&self.base_url) else {
return Access::unknown(format!(
"{id} has no model list to ask for, and the only other question it answers costs a call"
));
};
let parsed = match self.call(HttpRequest::get(url)).await {
Ok(parsed) => parsed,
Err(e) => return refusal_or_doubt(id, &e),
};
let listed = match self.protocol.read_catalogue(&parsed) {
Ok(listed) => listed,
Err(e) => {
return Access::unknown(format!(
"{id} sent a model list that could not be read: {e}"
))
}
};
if listed.contains(model) {
return Access::Ready;
}
Access::denied(format!(
"{id} answered with {} models and none of them is {model}. Either this account cannot reach it or the list spells it differently",
listed.len()
))
}
}
struct Frames<'a, P: Protocol> {
bytes: ByteStream,
protocol: &'a P,
asked_for: ModelId,
buffer: Vec<u8>,
ready: VecDeque<Event>,
done: bool,
}
impl<P: Protocol> Frames<'_, P> {
fn absorb(&mut self, chunk: &[u8]) -> Result<()> {
self.buffer.extend_from_slice(chunk);
while let Some(end) = frame_end(&self.buffer) {
let raw = self.buffer.drain(..end.0).collect::<Vec<u8>>();
self.buffer.drain(..end.1);
if let Some(frame) = parse_frame(&raw) {
self.ready
.extend(self.protocol.read_event(&frame, &self.asked_for)?);
}
}
Ok(())
}
}
impl<P: Protocol> Stream for Frames<'_, P> {
type Item = Result<Event>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if let Some(event) = this.ready.pop_front() {
return Poll::Ready(Some(Ok(event)));
}
if this.done {
return Poll::Ready(None);
}
match this.bytes.as_mut().poll_next(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Some(Err(e))) => {
this.done = true;
return Poll::Ready(Some(Err(e)));
}
Poll::Ready(Some(Ok(chunk))) => {
if let Err(e) = this.absorb(&chunk) {
this.done = true;
return Poll::Ready(Some(Err(e)));
}
}
Poll::Ready(None) => {
this.done = true;
let trailing = std::mem::take(&mut this.buffer);
if let Some(frame) = parse_frame(&trailing) {
match this.protocol.read_event(&frame, &this.asked_for) {
Ok(events) => this.ready.extend(events),
Err(e) => return Poll::Ready(Some(Err(e))),
}
}
}
}
}
}
}
fn frame_end(buffer: &[u8]) -> Option<(usize, usize)> {
let mut i = 0;
while i < buffer.len() {
if buffer[i..].starts_with(b"\r\n\r\n") {
return Some((i, 4));
}
if buffer[i..].starts_with(b"\n\n") {
return Some((i, 2));
}
i += 1;
}
None
}
fn parse_frame(raw: &[u8]) -> Option<SseFrame> {
let text = String::from_utf8_lossy(raw);
let mut frame = SseFrame::default();
let mut data: Vec<&str> = Vec::new();
for line in text.lines() {
let line = line.trim_end_matches('\r');
if line.starts_with(':') || line.is_empty() {
continue;
}
let Some((field, value)) = line.split_once(':') else {
continue;
};
let value = value.strip_prefix(' ').unwrap_or(value);
match field {
"event" => frame.event = value.to_string(),
"data" => data.push(value),
_ => {}
}
}
if data.is_empty() {
return None;
}
frame.data = data.join("\n");
Some(frame)
}
fn refusal_or_doubt(id: &str, e: &Error) -> Access {
match e {
Error::Auth(said) => Access::denied(format!("{id} rejected the credential: {said}")),
Error::NotFound(said) => Access::denied(format!("{id} has no such endpoint: {said}")),
other => Access::unknown(format!("{id} could not be asked: {other}")),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chat::{ContentBlock, Message, Role, StopReason};
use crate::cost::Usage;
use crate::transport::{HttpResponse, Method};
use crate::Provider;
use std::sync::Mutex;
struct Test {
lists: bool,
}
impl Protocol for Test {
fn id(&self) -> &str {
"test-protocol"
}
fn chat_url(&self, base_url: &str, _model: &ModelId) -> String {
format!("{base_url}/chat")
}
fn headers(&self, _key: &Secret) -> Result<Vec<(String, String)>> {
Ok(Vec::new())
}
fn body(&self, _request: &ChatRequest) -> Result<Value> {
Ok(serde_json::json!({}))
}
fn read(&self, _body: &Value, asked_for: &ModelId) -> Result<ChatResponse> {
Ok(ChatResponse::new(
Message {
role: Role::Assistant,
content: vec![ContentBlock::Text("ok".into())],
},
StopReason::EndTurn,
Usage::absent(),
asked_for.clone(),
))
}
fn catalogue_url(&self, base_url: &str) -> Option<String> {
self.lists.then(|| format!("{base_url}/models"))
}
fn read_catalogue(&self, body: &Value) -> Result<Vec<ModelId>> {
Ok(body
.get("data")
.and_then(Value::as_array)
.ok_or_else(|| Error::Unreadable("no data array".into()))?
.iter()
.filter_map(|m| m.get("id").and_then(Value::as_str))
.map(ModelId::from)
.collect())
}
}
struct Scripted {
replies: Mutex<Vec<Result<HttpResponse>>>,
seen: Mutex<Vec<(Method, String)>>,
}
impl Scripted {
fn new(replies: Vec<Result<HttpResponse>>) -> Arc<Self> {
Arc::new(Self {
replies: Mutex::new(replies),
seen: Mutex::new(Vec::new()),
})
}
fn urls(&self) -> Vec<(Method, String)> {
self.seen.lock().map(|s| s.clone()).unwrap_or_default()
}
}
#[async_trait]
impl crate::transport::HttpTransport for Scripted {
async fn send(&self, request: HttpRequest) -> Result<HttpResponse> {
if let Ok(mut seen) = self.seen.lock() {
seen.push((request.method, request.url.clone()));
}
let mut replies = self
.replies
.lock()
.map_err(|_| Error::Transient("poisoned".into()))?;
if replies.is_empty() {
return Err(Error::Transient("the test scripted no more replies".into()));
}
match replies.remove(0) {
Ok(reply) => Ok(reply),
Err(e) => Err(Error::Transient(e.to_string())),
}
}
}
fn listing(models: &[&str]) -> Result<HttpResponse> {
let body = serde_json::json!({
"data": models.iter().map(|m| serde_json::json!({ "id": m })).collect::<Vec<_>>()
});
Ok(HttpResponse::new(
200,
serde_json::to_vec(&body).unwrap_or_default(),
))
}
fn status(code: u16) -> Result<HttpResponse> {
Ok(HttpResponse::new(
code,
b"the vendor said something".to_vec(),
))
}
fn provider(lists: bool, transport: Arc<Scripted>) -> ApiProvider<Test> {
ApiProvider::new(
Test { lists },
"https://example.invalid",
transport,
Secret::new("key", "sk-test"),
Reach::FirstPartyApi,
Arc::new(Registry::empty("test", Reach::FirstPartyApi)),
)
}
#[tokio::test]
async fn a_protocol_with_no_model_list_answers_unknown_rather_than_denied() {
let transport = Scripted::new(Vec::new());
let access = provider(false, transport.clone())
.validate(&"any".into())
.await;
assert!(access.is_unknown(), "{access:?}");
assert!(
transport.urls().is_empty(),
"nothing should have been sent: {:?}",
transport.urls()
);
}
#[tokio::test]
async fn a_model_the_vendor_lists_is_ready() {
let transport = Scripted::new(vec![listing(&["gpt-test", "gpt-other"])]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert_eq!(access, Access::Ready, "{access:?}");
}
#[tokio::test]
async fn a_model_the_vendor_does_not_list_is_denied_and_says_which_two_things_it_could_be() {
let transport = Scripted::new(vec![listing(&["gpt-other"])]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_denied(), "{access:?}");
let said = access.detail().unwrap_or_default();
assert!(said.contains("gpt-test"), "{said}");
assert!(said.contains("cannot reach it"), "{said}");
assert!(said.contains("spells it differently"), "{said}");
}
#[tokio::test]
async fn a_rejected_credential_is_denied_rather_than_unknown() {
for code in [401, 403] {
let transport = Scripted::new(vec![status(code)]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_denied(), "{code} gave {access:?}");
assert!(
access.detail().unwrap_or_default().contains("credential"),
"{code}: {access}"
);
}
}
#[tokio::test]
async fn an_endpoint_that_is_not_there_is_denied() {
let transport = Scripted::new(vec![status(404)]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_denied(), "{access:?}");
}
#[tokio::test]
async fn a_moment_rather_than_an_answer_is_unknown() {
for code in [429, 500, 502, 503] {
let transport = Scripted::new(vec![status(code)]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_unknown(), "{code} gave {access:?}");
}
}
#[tokio::test]
async fn a_transport_that_could_not_connect_is_unknown() {
let transport = Scripted::new(vec![Err(Error::Transient("connection refused".into()))]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_unknown(), "{access:?}");
}
#[tokio::test]
async fn a_model_list_that_could_not_be_read_is_unknown_rather_than_denied() {
let transport = Scripted::new(vec![Ok(HttpResponse::new(
200,
b"{\"models\":[]}".to_vec(),
))]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_unknown(), "{access:?}");
assert!(
access
.detail()
.unwrap_or_default()
.contains("could not be read"),
"{access}"
);
}
#[tokio::test]
async fn a_reply_that_was_not_json_at_all_is_unknown() {
let transport = Scripted::new(vec![Ok(HttpResponse::new(200, b"<html>".to_vec()))]);
let access = provider(true, transport).validate(&"gpt-test".into()).await;
assert!(access.is_unknown(), "{access:?}");
}
#[tokio::test]
async fn validate_sends_one_get_and_never_touches_the_chat_endpoint() {
let transport = Scripted::new(vec![listing(&["gpt-test"])]);
let _ = provider(true, transport.clone())
.validate(&"gpt-test".into())
.await;
let seen = transport.urls();
assert_eq!(seen.len(), 1, "{seen:?}");
assert_eq!(seen[0].0, Method::Get, "{seen:?}");
assert!(seen[0].1.ends_with("/models"), "{seen:?}");
assert!(!seen[0].1.contains("/chat"), "{seen:?}");
}
#[tokio::test]
async fn the_answer_is_not_remembered_between_calls() {
let transport = Scripted::new(vec![status(401), listing(&["gpt-test"])]);
let provider = provider(true, transport.clone());
let first = provider.validate(&"gpt-test".into()).await;
let second = provider.validate(&"gpt-test".into()).await;
assert!(first.is_denied(), "{first:?}");
assert_eq!(second, Access::Ready, "{second:?}");
assert_eq!(transport.urls().len(), 2);
}
}