bitview: reorg part 9

crates/brk_server/src/api/metrics/data.rs

@@ -0,0 +1,126 @@
+use std::time::Duration;
+
+use axum::{
+    Json,
+    body::Body,
+    extract::{Query, State},
+    http::{HeaderMap, StatusCode, Uri},
+    response::{IntoResponse, Response},
+};
+use brk_error::{Error, Result};
+use brk_interface::{Format, Output, Params};
+use quick_cache::sync::GuardResult;
+use vecdb::Stamp;
+
+use crate::{HeaderMapExtended, ResponseExtended};
+
+use super::AppState;
+
+const MAX_WEIGHT: usize = 65 * 10_000;
+
+pub async fn handler(
+    uri: Uri,
+    headers: HeaderMap,
+    query: Query<Params>,
+    State(app_state): State<AppState>,
+) -> Response {
+    match req_to_response_res(uri, headers, query, app_state) {
+        Ok(response) => response,
+        Err(error) => {
+            let mut response =
+                (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response();
+            response.headers_mut().insert_cors();
+            response
+        }
+    }
+}
+
+fn req_to_response_res(
+    uri: Uri,
+    headers: HeaderMap,
+    Query(params): Query<Params>,
+    AppState {
+        interface, cache, ..
+    }: AppState,
+) -> Result<Response> {
+    let vecs = interface.search(&params)?;
+
+    if vecs.is_empty() {
+        return Ok(Json(vec![] as Vec<usize>).into_response());
+    }
+
+    let from = params.from();
+    let to = params.to();
+    let format = params.format();
+
+    // TODO: From and to should be capped here
+
+    let weight = vecs
+        .iter()
+        .map(|(_, v)| v.range_weight(from, to))
+        .sum::<usize>();
+
+    if weight > MAX_WEIGHT {
+        return Err(Error::String(format!(
+            "Request is too heavy, max weight is {MAX_WEIGHT} bytes"
+        )));
+    }
+
+    // TODO: height should be from vec, but good enough for now
+    let etag = vecs
+        .first()
+        .unwrap()
+        .1
+        .etag(Stamp::from(interface.get_height()), to);
+
+    if headers
+        .get_if_none_match()
+        .is_some_and(|prev_etag| etag == prev_etag)
+    {
+        return Ok(Response::new_not_modified());
+    }
+
+    let guard_res = cache.get_value_or_guard(
+        &format!("{}{}{etag}", uri.path(), uri.query().unwrap_or("")),
+        Some(Duration::from_millis(50)),
+    );
+
+    let mut response = if let GuardResult::Value(v) = guard_res {
+        Response::new(Body::from(v))
+    } else {
+        match interface.format(vecs, &params.rest)? {
+            Output::CSV(s) => {
+                if let GuardResult::Guard(g) = guard_res {
+                    g.insert(s.clone().into())
+                        .map_err(|_| Error::QuickCacheError)?;
+                }
+                s.into_response()
+            }
+            Output::Json(v) => {
+                let json = v.to_vec();
+                if let GuardResult::Guard(g) = guard_res {
+                    g.insert(json.clone().into())
+                        .map_err(|_| Error::QuickCacheError)?;
+                }
+                json.into_response()
+            }
+        }
+    };
+
+    let headers = response.headers_mut();
+
+    headers.insert_cors();
+
+    headers.insert_etag(&etag);
+    headers.insert_cache_control_must_revalidate();
+
+    match format {
+        Format::CSV => {
+            headers.insert_content_disposition_attachment();
+            headers.insert_content_type_text_csv()
+        }
+        Format::JSON => headers.insert_content_type_application_json(),
+    }
+
+    Ok(response)
+}

crates/brk_server/src/api/metrics/mod.rs

@@ -0,0 +1,111 @@
+use axum::{
+    Json, Router,
+    extract::{Path, Query, State},
+    http::{HeaderMap, Uri},
+    response::{IntoResponse, Response},
+    routing::get,
+};
+use brk_interface::{Index, PaginatedIndexParam, PaginationParam, Params, ParamsOpt};
+
+use super::AppState;
+
+mod data;
+
+pub trait ApiMetricsRoutes {
+    fn add_api_metrics_routes(self) -> Self;
+}
+
+const TO_SEPARATOR: &str = "_to_";
+
+impl ApiMetricsRoutes for Router<AppState> {
+    fn add_api_metrics_routes(self) -> Self {
+        self.route(
+            "/api/metrics/count",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(sonic_rs::json!({
+                    "distinct": app_state.interface.distinct_metric_count(),
+                    "total": app_state.interface.total_metric_count(),
+                }))
+                .into_response()
+            }),
+        )
+        .route(
+            "/api/metrics/indexes",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_accepted_indexes()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/metrics",
+            get(
+                async |State(app_state): State<AppState>,
+                       Query(pagination): Query<PaginationParam>|
+                       -> Response {
+                    Json(app_state.interface.get_metrics(pagination)).into_response()
+                },
+            ),
+        )
+        .route(
+            "/api/vecs/index-to-metrics",
+            get(
+                async |State(app_state): State<AppState>,
+                       Query(paginated_index): Query<PaginatedIndexParam>|
+                       -> Response {
+                    Json(app_state.interface.get_index_to_vecids(paginated_index)).into_response()
+                },
+            ),
+        )
+        .route(
+            "/api/metrics/{metric}",
+            get(
+                async |State(app_state): State<AppState>, Path(metric): Path<String>| -> Response {
+                    // If not found do fuzzy search but here or in interface ?
+                    Json(app_state.interface.metric_to_indexes(metric)).into_response()
+                },
+            ),
+        )
+        .route(
+            "/api/metrics/{metric}/{index}",
+            get(
+                async |State(app_state): State<AppState>,
+                       Path((metric, index)): Path<(String, Index)>|
+                       -> Response {
+                    // If not found do fuzzy search but here or in interface ?
+                    Json(
+                        format!("{metric}/{index}"), // app_state
+                                                     //     .interface
+                                                     //     .metric_to_indexes(metric.replace("-", "_")),
+                    )
+                    .into_response()
+                },
+            ),
+        )
+        // DEPRECATED
+        .route("/api/vecs/query", get(data::handler))
+        // DEPRECATED
+        .route(
+            "/api/vecs/{variant}",
+            get(
+                async |uri: Uri,
+                       headers: HeaderMap,
+                       Path(variant): Path<String>,
+                       Query(params_opt): Query<ParamsOpt>,
+                       state: State<AppState>|
+                       -> Response {
+                    let variant = variant.replace("-", "_");
+                    let mut split = variant.split(TO_SEPARATOR);
+
+                    if let Ok(index) = Index::try_from(split.next().unwrap()) {
+                        let params = Params::from((
+                            (index, split.collect::<Vec<_>>().join(TO_SEPARATOR)),
+                            params_opt,
+                        ));
+                        data::handler(uri, headers, Query(params), state).await
+                    } else {
+                        "Bad variant".into_response()
+                    }
+                },
+            ),
+        )
+    }
+}