From 7f1290a86ae0c3f16ae69d8218253bd3d42c40c0 Mon Sep 17 00:00:00 2001 From: Jacob Date: Sat, 24 Oct 2026 16:18:32 -0400 Subject: [PATCH] feat: extraction models --- src/server/embeddings.rs | 1 - src/server/extraction.rs | 217 +++++++++++++++++++++++++++++++++++++++ src/server/mod.rs | 2 + 3 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 src/server/extraction.rs diff --git a/src/server/embeddings.rs b/src/server/embeddings.rs index e248506..4cac59a 100644 --- a/src/server/embeddings.rs +++ b/src/server/embeddings.rs @@ -1,4 +1,3 @@ -use crate::CloudflareCredentials; use crate::openai; use crate::server::AppState; use crate::server::Upstream; diff --git a/src/server/extraction.rs b/src/server/extraction.rs new file mode 100644 index 0000000..f879146 --- /dev/null +++ b/src/server/extraction.rs @@ -0,0 +1,217 @@ +use crate::CloudflareCredentials; +use crate::openai; +use crate::server::AppState; +use crate::server::Upstream; +use crate::server::take_random_by_swapping; +use axum::Json; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::response::Response; +use axum_extra::routing::TypedPath; +use ecow::EcoString; +use headers::Header; +use headers::authorization::Bearer; +use mitsein::vec1::Vec1; +use reqwest::Certificate; +use secrecy::ExposeSecret; +use serde::Deserialize; +use serde_with::DeserializeFromStr; +use serde_with::SerializeDisplay; +use std::collections::VecDeque; +use strum::EnumString; +use strum::VariantArray; + +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/extraction")] +pub struct V1ExtractionPath; + +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] +#[strum(serialize_all = "lowercase")] +pub enum Shortcut { + M, +} + +#[derive(Debug, Clone)] +enum Provider { + Cloudflare, +} + +impl Shortcut { + fn suggestions(&self) -> Vec<(Provider, EcoString)> { + use Provider as P; + use Shortcut as S; + + match self { + S::M => { + vec![(P::Cloudflare, "@cf/cloudflare/clairic".into())] + } + } + } +} + +pub struct ExtractionResponse(pub openai::extraction::ExtractionResponse); + +impl IntoResponse for ExtractionResponse { + fn into_response(self) -> Response { + let Self(extraction) = self; + + Json(extraction).into_response() + } +} + +pub enum ExtractionError { + NoneSucceeded, +} + +impl IntoResponse for ExtractionError { + fn into_response(self) -> Response { + match self { + ExtractionError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(), + } + } +} + +#[tracing::instrument(skip(state, request))] +pub async fn create_extraction( + _: V1ExtractionPath, + State(state): State, + Json(request): Json>, +) -> Result { + let shortcut = request.model.clone(); + let mut model_upstreams = VecDeque::new(); + + for (provider, model) in shortcut.suggestions() { + let upstreams: Vec = + match provider { + Provider::Cloudflare => state + .app_args + .cloudflare_workers_ai_credentials + .iter() + .cloned() + .map( + |CloudflareCredentials { + account_id, + api_key, + }| { + Upstream::builder() + .url( + format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/{model}") + .parse() + .expect("TODO"), + ) + .api_key(api_key) + .build() + }, + ) + .collect(), + }; + + if let Ok(upstreams) = Vec1::try_from(upstreams) { + model_upstreams.push_back((model, upstreams)); + } + } + + while let Some((model, upstreams)) = model_upstreams.pop_front() { + let (upstream, upstreams) = take_random_by_swapping(upstreams, &mut rand::rng()); + if let Some(upstreams) = upstreams { + model_upstreams.push_back((model.clone(), upstreams)); + } + + tracing::info!(?model, ?upstream); + + let openai::extraction::ExtractionRequest { + model: _shortcut, + state, + fields, + images, + reference_date, + options, + } = request.clone(); + + let request = openai::extraction::ExtractionRequest { + model, + state, + fields, + images, + reference_date, + options, + }; + + let client = reqwest::Client::builder() + .tls_certs_only( + webpki_root_certs::TLS_SERVER_ROOT_CERTS + .iter() + .filter_map(|cert| Certificate::from_der(cert).ok()), + ) + .build() + .expect("TODO"); + + let mut upstream_request = client.post(upstream.url).json(&request); + + if let Some(api_key) = upstream.api_key { + let authorization = + headers::Authorization::bearer(api_key.expose_secret()).expect("TODO"); + upstream_request = upstream_request.header(headers::Authorization::::name(), { + let mut vec = vec![]; + + authorization.encode(&mut vec); + + let [one] = vec.try_into().expect("TODO"); + + one + }); + } + + let upstream_response = match upstream_request.send().await { + Ok(upstream_response) => upstream_response, + Err(upstream_error) => { + tracing::error!(?upstream_error); + continue; + } + }; + + if let Err(upstream_error) = upstream_response.error_for_status_ref() { + tracing::error!(?upstream_error); + match upstream_response.text().await { + Ok(body) => tracing::error!(?body), + Err(body_error) => tracing::error!(?body_error), + }; + continue; + }; + + let response: serde_json::Value = match upstream_response.json().await { + Ok(response) => response, + Err(json_error) => { + tracing::error!(?json_error); + continue; + } + }; + + match openai::extraction::ExtractionResponse::deserialize( + response.clone(), // TODO: not need to clone + ) { + Ok(response) => return Ok(ExtractionResponse(response)), + Err(structure_error) => { + #[derive(Deserialize)] + struct CloudflareRunResponse { + result: openai::extraction::ExtractionResponse, + success: bool, + errors: Vec, // TODO + messages: Vec, // TODO + } + + match CloudflareRunResponse::deserialize(response) { + Ok(response) => return Ok(ExtractionResponse(response.result)), + Err(cloudflare_structure_error) => { + tracing::error!(?structure_error); + tracing::error!(?cloudflare_structure_error); + } + } + } + } + } + Err(ExtractionError::NoneSucceeded) +} diff --git a/src/server/mod.rs b/src/server/mod.rs index 6acc429..7334cd8 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -13,6 +13,7 @@ use url::Url; mod chat; mod decisions; mod embeddings; +mod extraction; mod frontend; mod models; mod system_one; @@ -46,6 +47,7 @@ pub fn create_router() -> Router { .typed_post(chat::completions::create_chat_completion) .typed_post(decisions::create_decision) .typed_post(embeddings::create_embeddings) + .typed_post(extraction::create_extraction) .typed_post(system_one::create_system_one) .typed_get(models::list_models) .typed_get(frontend::spa)