use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::PROVIDER_NAME;
use super::wire::{REQUEST_ID_HEADER, stamp};
use crate::error::{EncodeError, ProviderError};
use crate::json_utils::Lenient;
use crate::model::{ModelInfo, ModelList};
use crate::operation::{ModelListing, ModelPage};
use crate::providers::internal::wire::classify_untyped_line;
use crate::providers::openai::OpenAIConfig;
use crate::wire::{Body, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
pub(crate) const MODEL_LISTING_PATH: &str = "/models";
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Models {
pub provider: OpenAIConfig,
}
impl Models {
pub fn new(provider: OpenAIConfig) -> Self {
Self { provider }
}
}
fn page(event: Value) -> Result<ModelPage, ProviderError> {
let text =
|entry: &Value, pointer: &str| entry.at(pointer).and_then(Value::as_str).map(str::to_owned);
let models = event
.arr("data")
.iter()
.filter_map(|entry| {
let mut model = ModelInfo::from_id(entry.str("id")?);
model.name = text(entry, "/name");
model.owned_by = text(entry, "/vendor");
model.r#type = text(entry, "/capabilities/type");
Some(model)
})
.collect();
Ok(ModelPage {
models: ModelList::new(models),
next: None,
})
}
impl Wire for Models {
type Op = ModelListing;
type Payload = crate::wire::Encoded;
type Frame = crate::wire::WireFrame;
type Event = Value;
type Decoder<'id> = DocumentDecoder<ModelListing, Value>;
type Reassembler = crate::wire::document::Unreassembled;
fn describe(&self) -> Descriptor<'_> {
Descriptor::new(PROVIDER_NAME)
}
fn encode(&self, _cursor: Option<String>, _mode: Mode) -> Result<Encoded, EncodeError> {
let uri = format!(
"{}{MODEL_LISTING_PATH}",
self.provider.base_url.trim_end_matches('/')
);
let mut request = http::Request::get(uri)
.header(http::header::CONTENT_TYPE, "application/json")
.body(Body::empty())?;
stamp(&mut request, self.provider.api_key.expose(), "user", false)?;
Ok(Encoded::new(request, Framing::Whole).with_request_id_header(REQUEST_ID_HEADER))
}
fn decoder<'id>(&self) -> Self::Decoder<'id> {
DocumentDecoder::new(|body| classify_untyped_line(body.as_bytes()), page)
}
}