use core::net::{IpAddr, Ipv4Addr, SocketAddr};
extern crate alloc;
use alloc::ffi::CString;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use crate::common::site::Site;
use crate::{
CancelToken, Context, OrtResult, Read, Write, chunked,
common::{buf_read, config, resolver},
http,
input::args,
};
pub fn run(
api_key: &str,
_cancel_token: CancelToken, settings: config::Settings,
opts: args::ListOpts,
site: &'static Site,
mut w: impl Write,
) -> OrtResult<()> {
let addrs: Vec<_> = if settings.dns.is_empty() {
let c_host = CString::new(site.host).unwrap();
let ips = unsafe { resolver::resolve(c_host.as_ptr())? };
ips.into_iter()
.map(|ip| SocketAddr::new(IpAddr::V4(ip), 443))
.collect()
} else {
settings
.dns
.into_iter()
.map(|a| {
let ip_addr = a.parse::<Ipv4Addr>().unwrap();
SocketAddr::new(IpAddr::V4(ip_addr), 443)
})
.collect()
};
let reader = http::list_models(api_key, site.host, site.list_url, addrs)
.context("list_models connect")?;
let mut reader = buf_read::OrtBufReader::new(reader);
let is_chunked = http::skip_header(&mut reader)?;
if opts.is_json {
if is_chunked {
let mut chunked = chunked::read(reader);
while let Some(chunk) = chunked.next_chunk() {
let chunk = chunk?;
w.write_all(chunk.as_bytes()).context("write models JSON")?;
}
} else {
let mut buf: [u8; 4096] = [0; 4096];
loop {
let bytes_read = reader.read(&mut buf).context("read models body")?;
if bytes_read == 0 {
break;
}
w.write_all(&buf[..bytes_read])
.context("write models JSON")?;
}
}
w.flush().context("flush models JSON")?;
} else {
let mut slugs = Vec::with_capacity(400);
if is_chunked {
let mut partial = String::new();
let mut chunked = chunked::read(reader);
while let Some(chunk) = chunked.next_chunk() {
let chunk = chunk?;
for (pos, section) in chunk.split(r#""id":""#).enumerate() {
let maybe_next_id = if pos == 0 && !partial.is_empty() {
partial.push_str(section);
until_quote(&partial)
} else if pos == 0 {
continue;
} else {
until_quote(section)
};
match maybe_next_id {
Some(slug) => {
slugs.push(slug.to_string());
partial.clear();
}
None => {
partial.push_str(section);
}
}
}
}
} else {
let mut models = String::with_capacity(512 * 1024);
reader
.read(unsafe { models.as_mut_vec().as_mut_slice() })
.context("read models body")?;
for slug in models.split(r#""id":""#).skip(1).filter_map(until_quote) {
slugs.push(slug.to_string());
}
};
slugs.sort();
for s in slugs {
let _ = w.write(s.as_bytes());
let _ = w.write(b"\n");
}
}
Ok(())
}
fn until_quote(s: &str) -> Option<&str> {
let mut qp = 0;
let len = s.len();
let b = s.as_bytes();
while qp < len && b[qp] != b'"' {
qp += 1;
}
if qp == len { None } else { Some(&s[..qp]) }
}