feat: extraction models

src/server/embeddings.rs

@@ -1,4 +1,3 @@
-use crate::CloudflareCredentials;
 use crate::openai;
 use crate::server::AppState;
 use crate::server::Upstream;

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<String>);
+
+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<AppState>,
+    Json(request): Json<openai::extraction::ExtractionRequest<Shortcut>>,
+) -> Result<ExtractionResponse, ExtractionError> {
+    let shortcut = request.model.clone();
+    let mut model_upstreams = VecDeque::new();
+
+    for (provider, model) in shortcut.suggestions() {
+        let upstreams: Vec<Upstream> =
+            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::<Bearer>::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<String>,
+                    success: bool,
+                    errors: Vec<serde_json::Value>,   // TODO
+                    messages: Vec<serde_json::Value>, // 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)
+}

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<AppState> {
         .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)