diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 585f613..3417fc1 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -96,24 +96,49 @@ Runtime that `ort` links needs glibc 2.38 or newer, which bookworm lacks. ## Code layout +The code is organised by domain. Each domain folder holds its model, its +use cases (a service) and its HTTP handlers: + ``` +src/datapoints/ ingest, delete and list datapoints + model.rs NewDatapoint, Datapoint, Meta, validation + ingest.rs single and batch ingest, idempotency, upsert by sourceId + delete.rs delete by sourceId or id + list.rs listing with filters and pagination + http.rs routes, request/response types +src/recall/ recall a project's datapoints for a query + select.rs the selection algorithm, pure over loaded metas + mod.rs the recall service + http.rs routes, request/response types +src/webhooks/ inbound webhooks + model.rs Webhook, Settings, filters, event mapping + service.rs manage webhooks, receive events + rate.rs per-webhook rate limit + http.rs routes, request/response types +src/app.rs Infra (store, index, embedder) and Services wiring +src/error.rs ServiceError, what a use case can fail with +src/http/ router, JSON extractor, HTTP error envelope +src/locks.rs in-process per-key locks src/main.rs wiring, download-only mode, server src/config.rs environment variables src/boot.rs index rebuild from S3 at startup -src/api/ handlers, request/response types, error mapping, key locks -src/domain.rs Datapoint, validation src/embed.rs Embedder trait + fastembed impl -src/store.rs S3 document store (content, meta, idempotency) +src/store.rs S3 document store (content, meta, idempotency, webhooks) src/index/mod.rs VectorIndex trait src/index/memory.rs in-memory impl src/index/s3vectors.rs Amazon S3 Vectors impl -src/recall.rs the recall algorithm, pure over loaded metas src/tokens.rs token counting ``` -Keep the recall algorithm in `recall.rs` pure: it takes already-loaded -metadata and scores and returns the selection. That is what makes it -unit-testable without S3. +Handlers in `http.rs` only parse the request, call the domain's service +and shape the response. Services know nothing about HTTP: they take plain +commands and return results or a `ServiceError`, which `src/http/error.rs` +maps to a status and error code. Business rules belong in the model or +the service, never in a handler. + +Keep the recall algorithm in `recall/select.rs` pure: it takes +already-loaded metadata and scores and returns the selection. That is +what makes it unit-testable without S3. ## Design docs diff --git a/src/api/datapoints.rs b/src/api/datapoints.rs deleted file mode 100644 index a978ee3..0000000 --- a/src/api/datapoints.rs +++ /dev/null @@ -1,252 +0,0 @@ -use axum::extract::rejection::QueryRejection; -use axum::extract::{Query, State}; -use chrono::{DateTime, Utc}; -use futures::{StreamExt, TryStreamExt}; -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; -use uuid::Uuid; - -use crate::api::ApiError; -use crate::api::state::AppState; -use crate::domain::{Meta, PROJECT_AUDIENCE, source_uuid, validate_audience}; - -pub const DEFAULT_LIMIT: usize = 50; -pub const MAX_LIMIT: usize = 200; -/// Most datapoints read per request. With narrow filters a page can come -/// back short, with a cursor to carry on from. -pub const MAX_SCAN: usize = 1000; -/// Datapoints read per round when filtering. -const CHUNK: usize = 100; -const FETCH_CONCURRENCY: usize = 16; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ListQuery { - pub project_id: Uuid, - pub limit: Option, - pub cursor: Option, - pub source_id: Option, - pub audience: Option, - pub label: Option, - pub occurred_from: Option>, - pub occurred_to: Option>, - #[serde(default)] - pub include_content: bool, -} - -/// A stored datapoint as listed: its metadata without the vector, and its -/// content only when asked for. -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DatapointView { - pub id: Uuid, - pub project_id: Uuid, - pub source_id: Option, - pub label: String, - pub timestamp: DateTime, - pub metadata: Map, - pub audience: Option, - pub content_tokens: usize, - pub content_sha256: Option, - pub embedding_model: String, - pub embedding_dim: usize, - #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, -} - -impl DatapointView { - fn new(meta: Meta, content: Option) -> DatapointView { - DatapointView { - id: meta.id, - project_id: meta.project_id, - source_id: meta.source_id, - label: meta.label, - timestamp: meta.timestamp, - metadata: meta.metadata, - audience: meta.audience, - content_tokens: meta.content_tokens, - content_sha256: meta.content_sha256, - embedding_model: meta.embedding.model, - embedding_dim: meta.embedding.dim, - content, - } - } -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DatapointList { - pub datapoints: Vec, - /// Pass as `cursor` for the next page; `null` on the last one. - pub next_cursor: Option, -} - -/// Which stored datapoints to list. -#[derive(Debug, Default)] -struct Filter { - /// `Some(None)` is the whole-project audience. - audience: Option>, - /// Lowercased. - label: Option, - from: Option>, - to: Option>, -} - -impl Filter { - fn from_query(q: &ListQuery) -> Result { - if let Some(audience) = &q.audience { - validate_audience(audience).map_err(ApiError::Validation)?; - } - if q.label.as_ref().is_some_and(|l| l.trim().is_empty()) { - return Err(ApiError::Validation( - "label must not be empty; omit it instead".into(), - )); - } - if let (Some(from), Some(to)) = (q.occurred_from, q.occurred_to) - && from > to - { - return Err(ApiError::RecallRangeInvalid); - } - Ok(Filter { - audience: q - .audience - .clone() - .map(|a| (a != PROJECT_AUDIENCE).then_some(a)), - label: q.label.as_ref().map(|l| l.to_lowercase()), - from: q.occurred_from, - to: q.occurred_to, - }) - } - - fn is_empty(&self) -> bool { - self.audience.is_none() && self.label.is_none() && self.from.is_none() && self.to.is_none() - } - - fn matches(&self, meta: &Meta) -> bool { - self.audience.as_ref().is_none_or(|a| *a == meta.audience) - && self - .label - .as_ref() - .is_none_or(|l| meta.label.to_lowercase().contains(l)) - && self.from.is_none_or(|from| meta.timestamp >= from) - && self.to.is_none_or(|to| meta.timestamp <= to) - } -} - -/// List a project's stored datapoints, read straight from storage, in id -/// order. Meant for debugging: it works while the index is still loading. -pub async fn list( - State(state): State, - query: Result, QueryRejection>, -) -> Result, ApiError> { - let Query(q) = query.map_err(|e| ApiError::Validation(e.body_text()))?; - let limit = q.limit.unwrap_or(DEFAULT_LIMIT); - if !(1..=MAX_LIMIT).contains(&limit) { - return Err(ApiError::Validation(format!( - "limit must be between 1 and {MAX_LIMIT}" - ))); - } - if q.source_id.as_ref().is_some_and(|s| s.trim().is_empty()) { - return Err(ApiError::Validation( - "sourceId must not be empty; omit it instead".into(), - )); - } - let cursor = q - .cursor - .as_deref() - .map(|c| { - c.parse::() - .map_err(|_| ApiError::Validation(format!("cursor {c:?} is not valid"))) - }) - .transpose()?; - let filter = Filter::from_query(&q)?; - let project = q.project_id; - - let (metas, next) = match &q.source_id { - Some(source_id) => { - let meta = state - .store - .get_meta(project, source_uuid(project, source_id)) - .await?; - ( - meta.into_iter().filter(|m| filter.matches(m)).collect(), - None, - ) - } - None => scan(&state, project, cursor, limit, &filter).await?, - }; - - let include_content = q.include_content; - let datapoints = futures::stream::iter(metas) - .map(|meta| { - let state = &state; - async move { - let content = if include_content { - state.store.get_content(project, meta.id).await? - } else { - None - }; - anyhow::Ok(DatapointView::new(meta, content)) - } - }) - .buffered(FETCH_CONCURRENCY) - .try_collect() - .await?; - Ok(axum::Json(DatapointList { - datapoints, - next_cursor: next.map(|id| id.to_string()), - })) -} - -/// Up to `limit` matching datapoints after `after`, reading at most -/// [`MAX_SCAN`], and the id to continue after when there may be more. -async fn scan( - state: &AppState, - project: Uuid, - mut after: Option, - limit: usize, - filter: &Filter, -) -> anyhow::Result<(Vec, Option)> { - let mut found = Vec::new(); - let mut scanned = 0; - loop { - // Without a filter every datapoint read is returned, so read no more - // than the page needs. - let want = if filter.is_empty() { - limit - found.len() - } else { - CHUNK - }; - let chunk = want.min(MAX_SCAN - scanned); - let ids = state - .store - .list_datapoint_ids(project, after, chunk) - .await?; - let exhausted = ids.len() < chunk; - let metas: Vec> = futures::stream::iter(ids.iter().copied()) - .map(|id| state.store.get_meta(project, id)) - .buffered(FETCH_CONCURRENCY) - .try_collect() - .await?; - let last = ids.last().copied(); - for (id, meta) in ids.into_iter().zip(metas) { - after = Some(id); - scanned += 1; - // Gone since listing, or filtered out. - let Some(meta) = meta.filter(|m| filter.matches(m)) else { - continue; - }; - found.push(meta); - if found.len() == limit { - let more = !(exhausted && Some(id) == last); - return Ok((found, more.then_some(id))); - } - } - if exhausted { - return Ok((found, None)); - } - if scanned >= MAX_SCAN { - return Ok((found, after)); - } - } -} diff --git a/src/api/delete.rs b/src/api/delete.rs deleted file mode 100644 index d1b9217..0000000 --- a/src/api/delete.rs +++ /dev/null @@ -1,90 +0,0 @@ -use axum::extract::State; -use futures::{StreamExt, TryStreamExt}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::api::ingest::datapoint_lock; -use crate::api::state::AppState; -use crate::api::{ApiError, Json}; -use crate::domain::source_uuid; - -/// Most ids accepted in one delete call. -pub const MAX_DELETE: usize = 1000; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct DeleteRequest { - pub project_id: Uuid, - #[serde(default)] - pub source_ids: Vec, - #[serde(default)] - pub ids: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DeleteResult { - pub id: Uuid, - #[serde(skip_serializing_if = "Option::is_none")] - pub source_id: Option, - /// False when nothing was stored under the id. - pub deleted: bool, -} - -#[derive(Debug, Serialize)] -pub struct DeleteResponse { - pub results: Vec, -} - -/// Delete datapoints by `sourceId` or server `id`. Results come back in -/// request order, `sourceIds` first. -pub async fn handler( - State(state): State, - Json(req): Json, -) -> Result, ApiError> { - state.ensure_ready()?; - let total = req.source_ids.len() + req.ids.len(); - if !(1..=MAX_DELETE).contains(&total) { - return Err(ApiError::Validation(format!( - "send between 1 and {MAX_DELETE} sourceIds and ids in total, got {total}" - ))); - } - if req.source_ids.iter().any(|s| s.trim().is_empty()) { - return Err(ApiError::Validation("sourceIds must not be empty".into())); - } - let project = req.project_id; - let targets = req - .source_ids - .into_iter() - .map(|s| (source_uuid(project, &s), Some(s))) - .chain(req.ids.into_iter().map(|id| (id, None))); - let results = futures::stream::iter(targets) - .map(|(id, source_id)| { - let state = &state; - async move { - let deleted = delete_one(state, project, id).await?; - anyhow::Ok(DeleteResult { - id, - source_id, - deleted, - }) - } - }) - .buffered(16) - .try_collect() - .await?; - Ok(axum::Json(DeleteResponse { results })) -} - -/// Remove from the index first, so a crash part way leaves a stored -/// datapoint that the next start puts back in the index, never a vector -/// pointing at nothing. -async fn delete_one(state: &AppState, project: Uuid, id: Uuid) -> anyhow::Result { - let _guard = state.locks.lock(datapoint_lock(project, id)).await; - if state.store.get_meta(project, id).await?.is_none() { - return Ok(false); - } - state.index.remove(project, id).await?; - state.store.delete_datapoint(project, id).await?; - Ok(true) -} diff --git a/src/api/error.rs b/src/api/error.rs deleted file mode 100644 index 18e4d10..0000000 --- a/src/api/error.rs +++ /dev/null @@ -1,171 +0,0 @@ -use axum::extract::rejection::JsonRejection; -use axum::http::StatusCode; -use axum::response::{IntoResponse, Response}; -use serde::Serialize; - -#[derive(Debug, Serialize)] -pub struct ItemError { - pub index: usize, - pub message: String, -} - -#[derive(Debug, thiserror::Error)] -pub enum ApiError { - #[error("{0}")] - Validation(String), - #[error("{message}")] - ValidationDetailed { - message: String, - details: Vec, - }, - #[error("occurredFrom must not be after occurredTo")] - RecallRangeInvalid, - #[error("send level or maxTokens, not both")] - BudgetConflict, - #[error("a batch holds at most 100 items")] - BatchTooLarge, - #[error("idempotencyKey was already used with a different request; use a new key")] - IdempotencyKeyReused, - #[error("request body is larger than 10 MiB")] - PayloadTooLarge, - #[error("{0}")] - NotFound(String), - #[error("this webhook is paused")] - WebhookPaused, - #[error("webhook body is larger than 1 MiB")] - WebhookBodyTooLarge, - #[error("too many events for this webhook; at most 600 a minute")] - RateLimited, - #[error("index is still loading from storage")] - Loading, - #[error("internal error")] - Internal(#[from] anyhow::Error), -} - -impl ApiError { - pub fn code(&self) -> &'static str { - match self { - ApiError::Validation(_) | ApiError::ValidationDetailed { .. } => "validation", - ApiError::RecallRangeInvalid => "recall-range-invalid", - ApiError::BudgetConflict => "budget-conflict", - ApiError::BatchTooLarge => "batch-too-large", - ApiError::IdempotencyKeyReused => "idempotency-key-reused", - ApiError::PayloadTooLarge => "payload-too-large", - ApiError::NotFound(_) => "not-found", - ApiError::WebhookPaused => "webhook-paused", - ApiError::WebhookBodyTooLarge => "webhook-body-too-large", - ApiError::RateLimited => "rate-limited", - ApiError::Loading => "loading", - ApiError::Internal(_) => "internal", - } - } - - pub fn status(&self) -> StatusCode { - match self { - ApiError::Validation(_) - | ApiError::ValidationDetailed { .. } - | ApiError::RecallRangeInvalid - | ApiError::BudgetConflict - | ApiError::BatchTooLarge => StatusCode::BAD_REQUEST, - ApiError::IdempotencyKeyReused => StatusCode::UNPROCESSABLE_ENTITY, - ApiError::PayloadTooLarge | ApiError::WebhookBodyTooLarge => { - StatusCode::PAYLOAD_TOO_LARGE - } - ApiError::NotFound(_) => StatusCode::NOT_FOUND, - ApiError::WebhookPaused => StatusCode::CONFLICT, - ApiError::RateLimited => StatusCode::TOO_MANY_REQUESTS, - ApiError::Loading => StatusCode::SERVICE_UNAVAILABLE, - ApiError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, - } - } - - pub fn to_body(&self) -> serde_json::Value { - let mut error = serde_json::json!({ "code": self.code(), "message": self.to_string() }); - if let ApiError::ValidationDetailed { details, .. } = self { - error["details"] = serde_json::to_value(details).unwrap_or_default(); - } - serde_json::json!({ "error": error }) - } -} - -impl IntoResponse for ApiError { - fn into_response(self) -> Response { - if let ApiError::Internal(e) = &self { - tracing::error!(error = ?e, "internal error"); - } - (self.status(), axum::Json(self.to_body())).into_response() - } -} - -impl From for ApiError { - fn from(r: JsonRejection) -> Self { - if r.status() == StatusCode::PAYLOAD_TOO_LARGE { - return ApiError::PayloadTooLarge; - } - ApiError::Validation(r.body_text()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use axum::http::StatusCode; - use axum::response::IntoResponse; - use http_body_util::BodyExt; - - async fn body_json(e: ApiError) -> (StatusCode, serde_json::Value) { - let resp = e.into_response(); - let status = resp.status(); - let bytes = resp.into_body().collect().await.unwrap().to_bytes(); - (status, serde_json::from_slice(&bytes).unwrap()) - } - - #[tokio::test] - async fn envelope_shape_and_status() { - let (status, body) = body_json(ApiError::RecallRangeInvalid).await; - assert_eq!(status, StatusCode::BAD_REQUEST); - assert_eq!(body["error"]["code"], "recall-range-invalid"); - assert!(body["error"]["message"].is_string()); - } - - #[tokio::test] - async fn internal_hides_details() { - let (status, body) = - body_json(ApiError::Internal(anyhow::anyhow!("secret s3 thing"))).await; - assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); - assert_eq!(body["error"]["message"], "internal error"); - assert_eq!(body["error"]["code"], "internal"); - } - - #[tokio::test] - async fn detailed_validation_includes_details() { - let e = ApiError::ValidationDetailed { - message: "bad items".into(), - details: vec![ItemError { - index: 2, - message: "label empty".into(), - }], - }; - let (status, body) = body_json(e).await; - assert_eq!(status, StatusCode::BAD_REQUEST); - assert_eq!(body["error"]["code"], "validation"); - assert_eq!(body["error"]["details"][0]["index"], 2); - } - - #[tokio::test] - async fn codes_and_statuses() { - assert_eq!(ApiError::Loading.status(), StatusCode::SERVICE_UNAVAILABLE); - assert_eq!(ApiError::Loading.code(), "loading"); - assert_eq!( - ApiError::IdempotencyKeyReused.code(), - "idempotency-key-reused" - ); - assert_eq!( - ApiError::IdempotencyKeyReused.status(), - StatusCode::UNPROCESSABLE_ENTITY - ); - assert_eq!(ApiError::BudgetConflict.code(), "budget-conflict"); - assert_eq!(ApiError::BatchTooLarge.code(), "batch-too-large"); - assert_eq!(ApiError::Validation("x".into()).code(), "validation"); - } -} diff --git a/src/api/ingest.rs b/src/api/ingest.rs deleted file mode 100644 index b1edaf5..0000000 --- a/src/api/ingest.rs +++ /dev/null @@ -1,337 +0,0 @@ -use std::collections::HashSet; - -use anyhow::Context; -use axum::extract::State; -use axum::http::StatusCode; -use chrono::Utc; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::api::error::ItemError; -use crate::api::state::AppState; -use crate::api::{ApiError, Json}; -use crate::domain::{Datapoint, Meta, NewDatapoint, StoredEmbedding, content_hash}; -use crate::embed::build_input; -use crate::store::MarkerOutcome; -use crate::tokens::count_tokens; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct IngestRequest { - pub idempotency_key: Option, - #[serde(flatten)] - pub datapoint: NewDatapoint, -} - -#[derive(Debug, Serialize)] -pub struct IngestResponse { - #[serde(flatten)] - pub datapoint: Datapoint, - pub deduplicated: bool, - pub replaced: bool, -} - -pub async fn handler( - State(state): State, - Json(req): Json, -) -> Result<(StatusCode, axum::Json), ApiError> { - state.ensure_ready()?; - if req.idempotency_key.is_some() && req.datapoint.source_id.is_some() { - return Err(ApiError::Validation( - "send idempotencyKey or sourceId, not both; resending a sourceId is already safe" - .into(), - )); - } - let (status, body) = ingest_one(&state, req.datapoint, req.idempotency_key.as_deref()).await?; - Ok((status, axum::Json(body))) -} - -/// An empty key would silently dedupe every later write that also sends one. -fn validate_key(key: Option<&str>) -> Result<(), ApiError> { - match key { - Some(k) if k.trim().is_empty() => Err(ApiError::Validation( - "idempotencyKey must not be empty; omit it instead".into(), - )), - _ => Ok(()), - } -} - -/// The in-process lock serialising writes and deletes of one datapoint. -pub fn datapoint_lock(project: Uuid, id: Uuid) -> String { - format!("{project}:datapoint:{id}") -} - -/// Store one datapoint. With a `sourceId`, the datapoint is stored under -/// it and `key` is ignored. With an idempotency key, a repeat of the same -/// request returns the original datapoint, a half-finished earlier attempt -/// is completed under the id it already reserved, and a different request -/// reusing the key is rejected. -pub async fn ingest_one( - state: &AppState, - new: NewDatapoint, - key: Option<&str>, -) -> Result<(StatusCode, IngestResponse), ApiError> { - validate_key(key)?; - new.validate().map_err(ApiError::Validation)?; - if new.source_id.is_some() { - return upsert(state, new).await; - } - let project = new.project_id; - let fingerprint = new.fingerprint(); - let mut id = new.id(); - - let _guard = match key { - Some(key) => Some(state.locks.lock(format!("{project}:key:{key}")).await), - None => None, - }; - if let Some(key) = key - && let MarkerOutcome::Exists { - id: existing_id, - fingerprint: recorded, - } = state - .store - .create_marker(project, key, id, &fingerprint) - .await? - { - // Markers from before fingerprints existed are trusted as before. - if recorded.is_some_and(|f| f != fingerprint) { - return Err(ApiError::IdempotencyKeyReused); - } - if let Some(meta) = state.store.get_meta(project, existing_id).await? { - let content = state - .store - .get_content(project, existing_id) - .await? - .with_context(|| format!("datapoint {existing_id} has metadata but no content"))?; - // The earlier attempt may have stopped after meta.json but - // before the index upsert. Upserting again is harmless. - if meta.embedding.model == state.embedder.model_name() { - state - .index - .upsert(project, existing_id, meta.embedding.vector.clone()) - .await?; - } - return Ok(( - StatusCode::OK, - IngestResponse { - datapoint: meta.into_datapoint(content), - deduplicated: true, - replaced: false, - }, - )); - } - // An earlier attempt reserved this id but did not finish. - id = existing_id; - } - - let datapoint = new - .into_parts(id, None, Utc::now()) - .map_err(ApiError::Validation)?; - let datapoint = write(state, datapoint).await?; - Ok(( - StatusCode::CREATED, - IngestResponse { - datapoint, - deduplicated: false, - replaced: false, - }, - )) -} - -/// Store a datapoint under its `sourceId`. A resend that changes nothing -/// returns the stored datapoint; one that changes anything replaces it. -async fn upsert( - state: &AppState, - new: NewDatapoint, -) -> Result<(StatusCode, IngestResponse), ApiError> { - let project = new.project_id; - let id = new.id(); - let _guard = state.locks.lock(datapoint_lock(project, id)).await; - let existing = state.store.get_meta(project, id).await?; - let datapoint = new - .into_parts(id, existing.as_ref(), Utc::now()) - .map_err(ApiError::Validation)?; - let (status, deduplicated, replaced, datapoint) = match existing { - Some(meta) if meta.matches(&datapoint, state.embedder.model_name()) => { - // An earlier send may have stopped before the index upsert. - state - .index - .upsert(project, id, meta.embedding.vector) - .await?; - (StatusCode::OK, true, false, datapoint) - } - Some(_) => (StatusCode::OK, false, true, write(state, datapoint).await?), - None => ( - StatusCode::CREATED, - false, - false, - write(state, datapoint).await?, - ), - }; - Ok(( - status, - IngestResponse { - datapoint, - deduplicated, - replaced, - }, - )) -} - -/// Embed, then write content, then metadata (the commit point), then index. -async fn write(state: &AppState, datapoint: Datapoint) -> Result { - let id = datapoint.id; - let input = build_input(&datapoint.label, &datapoint.content, &datapoint.metadata); - let embedder = state.embedder.clone(); - let vector = tokio::task::spawn_blocking(move || embedder.embed_passages(&[input])) - .await - .context("embedding task panicked")?? - .pop() - .context("embedder returned no vector")?; - - let meta = Meta { - id, - project_id: datapoint.project_id, - source_id: datapoint.source_id.clone(), - label: datapoint.label.clone(), - timestamp: datapoint.timestamp, - metadata: datapoint.metadata.clone(), - audience: datapoint.audience.clone(), - embedding: StoredEmbedding { - model: state.embedder.model_name().to_string(), - dim: vector.len(), - vector: vector.clone(), - }, - content_tokens: count_tokens(&datapoint.content), - content_sha256: Some(content_hash(&datapoint.content)), - }; - state - .store - .put_content(datapoint.project_id, id, &datapoint.content) - .await?; - state.store.put_meta(&meta).await?; - state.index.upsert(datapoint.project_id, id, vector).await?; - Ok(datapoint) -} - -/// Most items accepted in one batch. -pub const MAX_BATCH: usize = 100; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchRequest { - pub idempotency_key: Option, - pub items: Vec, -} - -#[derive(Debug, Serialize)] -pub struct ItemErrorBody { - pub code: String, - pub message: String, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -pub enum BatchItemResult { - Ok { - status: u16, - deduplicated: bool, - replaced: bool, - datapoint: Datapoint, - }, - Err { - status: u16, - error: ItemErrorBody, - }, -} - -#[derive(Debug, Serialize)] -pub struct BatchResponse { - pub results: Vec, -} - -/// Why each item would fail, checked before anything is written. An item -/// with a `sourceId` and no label needs a stored datapoint, or an earlier -/// item in the batch that creates it, to take its label from. -async fn invalid_items(state: &AppState, items: &[NewDatapoint]) -> anyhow::Result> { - let mut details = Vec::new(); - let mut labelled = HashSet::new(); - for (index, item) in items.iter().enumerate() { - if let Err(message) = item.validate() { - details.push(ItemError { index, message }); - continue; - } - let Some(source) = &item.source_id else { - continue; - }; - let id = item.id(); - if item.label.is_some() { - labelled.insert(id); - } else if !labelled.contains(&id) - && state.store.get_meta(item.project_id, id).await?.is_none() - { - details.push(ItemError { - index, - message: format!( - "label is required: no datapoint with sourceId {source:?} is stored" - ), - }); - } - } - Ok(details) -} - -pub async fn batch_handler( - State(state): State, - Json(req): Json, -) -> Result, ApiError> { - state.ensure_ready()?; - validate_key(req.idempotency_key.as_deref())?; - if req.items.len() > MAX_BATCH { - return Err(ApiError::BatchTooLarge); - } - let details = invalid_items(&state, &req.items).await?; - if !details.is_empty() { - return Err(ApiError::ValidationDetailed { - message: "invalid items".into(), - details, - }); - } - - let total = req.items.len(); - let mut results = Vec::with_capacity(total); - for (i, item) in req.items.into_iter().enumerate() { - let key = req.idempotency_key.as_ref().map(|k| format!("{k}#{i}")); - match ingest_one(&state, item, key.as_deref()).await { - Ok((status, resp)) => results.push(BatchItemResult::Ok { - status: status.as_u16(), - deduplicated: resp.deduplicated, - replaced: resp.replaced, - datapoint: resp.datapoint, - }), - Err(e) => { - if let ApiError::Internal(inner) = &e { - tracing::error!(error = ?inner, item = i, "batch item failed"); - } - results.push(BatchItemResult::Err { - status: e.status().as_u16(), - error: ItemErrorBody { - code: e.code().into(), - message: e.to_string(), - }, - }); - for _ in i + 1..total { - results.push(BatchItemResult::Err { - status: 0, - error: ItemErrorBody { - code: "skipped".into(), - message: "not attempted after an earlier item failed".into(), - }, - }); - } - break; - } - } - } - Ok(axum::Json(BatchResponse { results })) -} diff --git a/src/api/mod.rs b/src/api/mod.rs deleted file mode 100644 index aaefa25..0000000 --- a/src/api/mod.rs +++ /dev/null @@ -1,62 +0,0 @@ -pub mod datapoints; -pub mod delete; -pub mod error; -pub mod ingest; -pub mod locks; -pub mod rate; -pub mod recall; -pub mod state; -pub mod webhooks; - -use axum::Router; -use axum::extract::{DefaultBodyLimit, FromRequest, State}; -use axum::routing::{get, post}; - -use crate::webhook::MAX_BODY_BYTES as MAX_WEBHOOK_BODY_BYTES; -use tower_http::trace::TraceLayer; - -pub use error::ApiError; -pub use state::AppState; - -#[derive(FromRequest)] -#[from_request(via(axum::Json), rejection(ApiError))] -pub struct Json(pub T); - -/// Largest request body accepted, in bytes. -pub const MAX_BODY_BYTES: usize = 10 * 1024 * 1024; - -async fn healthz(State(state): State) -> axum::Json { - let status = if state.ensure_ready().is_ok() { - "ok" - } else { - "loading" - }; - axum::Json(serde_json::json!({ "status": status })) -} - -pub fn router(state: AppState) -> Router { - Router::new() - .route("/healthz", get(healthz)) - .route("/api/v1/ingest", post(ingest::handler)) - .route("/api/v1/ingest/batch", post(ingest::batch_handler)) - .route("/api/v1/recall", post(recall::handler)) - .route("/api/v1/datapoints", get(datapoints::list)) - .route("/api/v1/delete", post(delete::handler)) - .route( - "/api/v1/webhooks", - post(webhooks::create).get(webhooks::list), - ) - .route( - "/api/v1/webhooks/{id}", - get(webhooks::get) - .put(webhooks::update) - .delete(webhooks::delete), - ) - .route( - "/api/v1/webhooks/{id}/{token}", - post(webhooks::receive).layer(DefaultBodyLimit::max(MAX_WEBHOOK_BODY_BYTES)), - ) - .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) - .layer(TraceLayer::new_for_http()) - .with_state(state) -} diff --git a/src/api/recall.rs b/src/api/recall.rs deleted file mode 100644 index b2ac549..0000000 --- a/src/api/recall.rs +++ /dev/null @@ -1,230 +0,0 @@ -use std::collections::HashMap; - -use anyhow::Context; -use axum::extract::State; -use chrono::{DateTime, Utc}; -use futures::{StreamExt, TryStreamExt}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::api::state::AppState; -use crate::api::{ApiError, Json}; -use crate::domain::{Datapoint, Meta}; -use crate::recall::{Candidate, Level, SelectParams, Sort, select}; - -pub const DEFAULT_LIMIT: usize = 100; -pub const MAX_LIMIT: usize = 200; -/// Most hits fetched from the index before time filtering. -const MAX_SEARCH_K: usize = 1000; -const FETCH_CONCURRENCY: usize = 16; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct RecallRequest { - pub project_id: Uuid, - pub query: String, - pub limit: Option, - pub level: Option, - pub max_tokens: Option, - pub query_timestamp: Option>, - pub occurred_from: Option>, - pub occurred_to: Option>, - pub sort: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct RecallItem { - pub rank: usize, - pub score: f32, - #[serde(flatten)] - pub datapoint: Datapoint, - pub content_tokens: usize, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct RecallResponse { - pub results: Vec, - pub used_tokens: usize, - pub budget_tokens: usize, -} - -/// The token budget: `level` or `maxTokens`, never both; `low` by default. -pub fn resolve_budget(level: Option, max_tokens: Option) -> Result { - match (level, max_tokens) { - (Some(_), Some(_)) => Err(ApiError::BudgetConflict), - (_, Some(0)) => Err(ApiError::Validation("maxTokens must be at least 1".into())), - (_, Some(n)) => Ok(n), - (level, None) => Ok(level.unwrap_or_default().tokens()), - } -} - -pub fn validate(req: &RecallRequest) -> Result<(), ApiError> { - if req.query.trim().is_empty() { - return Err(ApiError::Validation("query must not be empty".into())); - } - if let Some(limit) = req.limit - && !(1..=MAX_LIMIT).contains(&limit) - { - return Err(ApiError::Validation(format!( - "limit must be between 1 and {MAX_LIMIT}" - ))); - } - if let (Some(from), Some(to)) = (req.occurred_from, req.occurred_to) - && from > to - { - return Err(ApiError::RecallRangeInvalid); - } - Ok(()) -} - -pub async fn handler( - State(state): State, - Json(req): Json, -) -> Result, ApiError> { - state.ensure_ready()?; - validate(&req)?; - let budget_tokens = resolve_budget(req.level, req.max_tokens)?; - let limit = req.limit.unwrap_or(DEFAULT_LIMIT); - let project = req.project_id; - - let embedder = state.embedder.clone(); - let query = req.query; - let qvec = tokio::task::spawn_blocking(move || embedder.embed_query(&query)) - .await - .context("embedding task panicked")??; - - let hits = state - .index - .search(project, &qvec, (limit * 4).min(MAX_SEARCH_K)) - .await?; - - let store = &state.store; - let metas: Vec<(f32, Option)> = futures::stream::iter(hits) - .map(|(id, score)| async move { store.get_meta(project, id).await.map(|m| (score, m)) }) - .buffer_unordered(FETCH_CONCURRENCY) - .try_collect() - .await?; - let mut by_id: HashMap = HashMap::new(); - let mut candidates = Vec::with_capacity(metas.len()); - for (score, meta) in metas { - let Some(meta) = meta else { continue }; - candidates.push(Candidate { - id: meta.id, - score, - timestamp: meta.timestamp, - content_tokens: meta.content_tokens, - }); - by_id.insert(meta.id, meta); - } - - let selected = select( - candidates, - &SelectParams { - limit, - budget_tokens, - query_timestamp: req.query_timestamp.unwrap_or_else(Utc::now), - occurred_from: req.occurred_from, - occurred_to: req.occurred_to, - sort: req.sort.unwrap_or_default(), - }, - ); - - let contents: HashMap> = - futures::stream::iter(selected.items.iter().map(|s| s.id).collect::>()) - .map(|id| async move { store.get_content(project, id).await.map(|c| (id, c)) }) - .buffer_unordered(FETCH_CONCURRENCY) - .try_collect() - .await?; - - let mut used_tokens = selected.used_tokens; - let mut results = Vec::with_capacity(selected.items.len()); - for s in selected.items { - let meta = by_id - .remove(&s.id) - .context("selected id without metadata")?; - let content_tokens = meta.content_tokens; - let Some(content) = contents.get(&s.id).cloned().flatten() else { - // Only reachable if the content object was removed outside the - // service. Skip it rather than fail every query that finds it. - tracing::warn!(project = %project, id = %s.id, "metadata without content; skipping"); - used_tokens -= content_tokens; - continue; - }; - results.push(RecallItem { - rank: s.rank, - score: s.score, - datapoint: meta.into_datapoint(content), - content_tokens, - }); - } - - Ok(axum::Json(RecallResponse { - results, - used_tokens, - budget_tokens, - })) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn req() -> RecallRequest { - serde_json::from_value(serde_json::json!({ - "projectId": "019a0000-0000-7000-8000-000000000000", - "query": "q" - })) - .unwrap() - } - - #[test] - fn default_budget_is_8000() { - assert_eq!(resolve_budget(None, None).unwrap(), 8000); - assert_eq!(resolve_budget(Some(Level::High), None).unwrap(), 24000); - assert_eq!(resolve_budget(None, Some(50)).unwrap(), 50); - } - - #[test] - fn both_is_conflict() { - assert!(matches!( - resolve_budget(Some(Level::Low), Some(5)), - Err(ApiError::BudgetConflict) - )); - } - - #[test] - fn zero_max_tokens_rejected() { - assert!(matches!( - resolve_budget(None, Some(0)), - Err(ApiError::Validation(_)) - )); - } - - #[test] - fn range_invalid() { - let mut r = req(); - r.occurred_from = Some("2026-10-02T00:00:00Z".parse().unwrap()); - r.occurred_to = Some("2026-10-01T00:00:00Z".parse().unwrap()); - assert!(matches!(validate(&r), Err(ApiError::RecallRangeInvalid))); - r.occurred_to = r.occurred_from; - assert!(validate(&r).is_ok()); - } - - #[test] - fn limit_bounds() { - for (limit, ok) in [(0, false), (201, false), (1, true), (200, true)] { - let mut r = req(); - r.limit = Some(limit); - assert_eq!(validate(&r).is_ok(), ok, "limit {limit}"); - } - } - - #[test] - fn blank_query_rejected() { - let mut r = req(); - r.query = " \n".into(); - assert!(matches!(validate(&r), Err(ApiError::Validation(_)))); - } -} diff --git a/src/api/state.rs b/src/api/state.rs deleted file mode 100644 index cbe7562..0000000 --- a/src/api/state.rs +++ /dev/null @@ -1,32 +0,0 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use crate::api::ApiError; -use crate::api::locks::KeyLocks; -use crate::api::rate::RateLimiter; -use crate::embed::Embedder; -use crate::index::VectorIndex; -use crate::store::S3Store; - -/// Everything a request handler needs. -#[derive(Clone)] -pub struct AppState { - pub store: S3Store, - pub index: Arc, - pub embedder: Arc, - /// False until the index has been rebuilt from storage. - pub ready: Arc, - pub locks: KeyLocks, - /// Events accepted per webhook per minute. - pub webhook_rate: RateLimiter, -} - -impl AppState { - pub fn ensure_ready(&self) -> Result<(), ApiError> { - if self.ready.load(Ordering::SeqCst) { - Ok(()) - } else { - Err(ApiError::Loading) - } - } -} diff --git a/src/api/webhooks.rs b/src/api/webhooks.rs deleted file mode 100644 index bc7ce3b..0000000 --- a/src/api/webhooks.rs +++ /dev/null @@ -1,244 +0,0 @@ -use std::time::Duration; - -use axum::body::Bytes; -use axum::extract::rejection::{BytesRejection, QueryRejection}; -use axum::extract::{Path, Query, State}; -use axum::http::StatusCode; -use chrono::Utc; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::api::ingest::ingest_one; -use crate::api::state::AppState; -use crate::api::{ApiError, Json}; -use crate::webhook::{MAX_PER_PROJECT, Settings, Webhook, WebhookView, is_webhook_id}; - -/// Most events one webhook accepts per minute. -pub const EVENTS_PER_MINUTE: u32 = 600; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct WebhookRequest { - pub project_id: Uuid, - #[serde(flatten)] - pub settings: Settings, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ProjectQuery { - pub project_id: Uuid, -} - -#[derive(Debug, Serialize)] -pub struct CreatedWebhook { - #[serde(flatten)] - pub webhook: WebhookView, - /// Where to send events. The token in it is the only credential and is - /// not shown again. - pub path: String, -} - -#[derive(Debug, Serialize)] -pub struct WebhookList { - pub webhooks: Vec, -} - -#[derive(Debug, Serialize)] -pub struct DeletedWebhook { - pub id: String, - pub deleted: bool, -} - -fn project_of(query: Result, QueryRejection>) -> Result { - query - .map(|Query(q)| q.project_id) - .map_err(|e| ApiError::Validation(e.body_text())) -} - -fn not_found(id: &str) -> ApiError { - ApiError::NotFound(format!("no webhook {id:?} in this project")) -} - -/// The webhook, if it exists and belongs to `project`. -async fn load(state: &AppState, project: Uuid, id: &str) -> Result { - if !is_webhook_id(id) { - return Err(not_found(id)); - } - match state.store.get_webhook(id).await? { - Some(w) if w.project_id == project => Ok(w), - _ => Err(not_found(id)), - } -} - -/// The in-process lock serialising changes to one project's webhooks. -fn webhooks_lock(project: Uuid) -> String { - format!("{project}:webhooks") -} - -pub async fn create( - State(state): State, - Json(req): Json, -) -> Result<(StatusCode, axum::Json), ApiError> { - req.settings.validate().map_err(ApiError::Validation)?; - let project = req.project_id; - let _guard = state.locks.lock(webhooks_lock(project)).await; - if state.store.list_webhooks(project).await?.len() >= MAX_PER_PROJECT { - return Err(ApiError::Validation(format!( - "a project holds at most {MAX_PER_PROJECT} webhooks; delete one first" - ))); - } - let (webhook, token) = Webhook::create(project, req.settings, Utc::now()); - state.store.put_webhook(&webhook).await?; - let path = format!("/api/v1/webhooks/{}/{token}", webhook.id); - Ok(( - StatusCode::CREATED, - axum::Json(CreatedWebhook { - webhook: webhook.view(), - path, - }), - )) -} - -pub async fn list( - State(state): State, - query: Result, QueryRejection>, -) -> Result, ApiError> { - let project = project_of(query)?; - let webhooks = state.store.list_webhooks(project).await?; - Ok(axum::Json(WebhookList { - webhooks: webhooks.into_iter().map(Webhook::view).collect(), - })) -} - -pub async fn get( - State(state): State, - Path(id): Path, - query: Result, QueryRejection>, -) -> Result, ApiError> { - let project = project_of(query)?; - Ok(axum::Json(load(&state, project, &id).await?.view())) -} - -/// Replace a webhook's settings. Its id and token stay the same. -pub async fn update( - State(state): State, - Path(id): Path, - Json(req): Json, -) -> Result, ApiError> { - req.settings.validate().map_err(ApiError::Validation)?; - let _guard = state.locks.lock(webhooks_lock(req.project_id)).await; - let mut webhook = load(&state, req.project_id, &id).await?; - webhook.settings = req.settings; - state.store.put_webhook(&webhook).await?; - Ok(axum::Json(webhook.view())) -} - -/// Delete a webhook. Datapoints it created stay. -pub async fn delete( - State(state): State, - Path(id): Path, - query: Result, QueryRejection>, -) -> Result, ApiError> { - let project = project_of(query)?; - let _guard = state.locks.lock(webhooks_lock(project)).await; - let deleted = match load(&state, project, &id).await { - Ok(_) => { - state.store.delete_webhook(&id).await?; - state.webhook_rate.forget(&id); - true - } - Err(ApiError::NotFound(_)) => false, - Err(e) => return Err(e), - }; - Ok(axum::Json(DeletedWebhook { id, deleted })) -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct EventResponse { - pub event_id: String, - /// `ingested` or `dropped`. - pub verdict: &'static str, - /// Why the event was dropped: the first filter it failed, or the field - /// it lacked. - pub dropped_by: Option, - pub datapoint_id: Option, - /// True when the event changed nothing that was already stored. - pub unchanged: bool, -} - -/// Turn an event into a datapoint. The token in the URL is the only -/// credential, so a wrong id and a wrong token look the same. -pub async fn receive( - State(state): State, - Path((id, token)): Path<(String, String)>, - body: Result, -) -> Result<(StatusCode, axum::Json), ApiError> { - let unknown = || ApiError::NotFound("no such webhook".into()); - if !is_webhook_id(&id) { - return Err(unknown()); - } - let webhook = match state.store.get_webhook(&id).await? { - Some(w) if w.token_matches(&token) => w, - _ => return Err(unknown()), - }; - let body = body.map_err(|e| { - if e.status() == StatusCode::PAYLOAD_TOO_LARGE { - ApiError::WebhookBodyTooLarge - } else { - ApiError::Validation(e.body_text()) - } - })?; - if webhook.settings.paused { - return Err(ApiError::WebhookPaused); - } - if !state - .webhook_rate - .allow(&id, EVENTS_PER_MINUTE, Duration::from_secs(60)) - { - return Err(ApiError::RateLimited); - } - state.ensure_ready()?; - - let event_id = format!("evt_{}", Uuid::now_v7().simple()); - let new = match webhook.map_event(&body) { - Ok(new) => new, - Err(reason) => { - tracing::info!( - webhook = id, - event = event_id, - reason, - "webhook event dropped" - ); - return Ok(( - StatusCode::OK, - axum::Json(EventResponse { - event_id, - verdict: "dropped", - dropped_by: Some(reason), - datapoint_id: None, - unchanged: false, - }), - )); - } - }; - // Without an id from the event, an identical event is keyed on what it - // maps to, so a resend is not stored twice. - let key = new - .source_id - .is_none() - .then(|| format!("webhook:{id}:{}", new.fingerprint())); - let (_, resp) = ingest_one(&state, new, key.as_deref()).await?; - tracing::info!(webhook = id, event = event_id, datapoint = %resp.datapoint.id, "webhook event ingested"); - Ok(( - StatusCode::ACCEPTED, - axum::Json(EventResponse { - event_id, - verdict: "ingested", - dropped_by: None, - datapoint_id: Some(resp.datapoint.id), - unchanged: resp.deduplicated, - }), - )) -} diff --git a/src/app.rs b/src/app.rs new file mode 100644 index 0000000..d6fed49 --- /dev/null +++ b/src/app.rs @@ -0,0 +1,64 @@ +//! Wires each domain's service to the shared infrastructure. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use crate::datapoints::DatapointService; +use crate::embed::Embedder; +use crate::error::ServiceError; +use crate::index::VectorIndex; +use crate::locks::KeyLocks; +use crate::recall::RecallService; +use crate::store::S3Store; +use crate::webhooks::WebhookService; + +/// The infrastructure the services run on. +#[derive(Clone)] +pub struct Infra { + pub store: S3Store, + pub index: Arc, + pub embedder: Arc, + /// False until the index has been rebuilt from storage. + pub ready: Arc, +} + +impl Infra { + pub fn is_ready(&self) -> bool { + self.ready.load(Ordering::SeqCst) + } + + pub(crate) fn ensure_ready(&self) -> Result<(), ServiceError> { + if self.is_ready() { + Ok(()) + } else { + Err(ServiceError::Loading) + } + } +} + +/// Every service, wired to one set of infrastructure. +#[derive(Clone)] +pub struct Services { + pub datapoints: DatapointService, + pub recall: RecallService, + pub webhooks: WebhookService, + infra: Infra, +} + +impl Services { + pub fn new(infra: Infra) -> Services { + let locks = KeyLocks::default(); + let datapoints = DatapointService::new(infra.clone(), locks.clone()); + Services { + recall: RecallService::new(infra.clone()), + webhooks: WebhookService::new(infra.clone(), locks, datapoints.clone()), + datapoints, + infra, + } + } + + /// Whether the index has been rebuilt from storage. + pub fn is_ready(&self) -> bool { + self.infra.is_ready() + } +} diff --git a/src/boot.rs b/src/boot.rs index 91a40d6..03d9cf1 100644 --- a/src/boot.rs +++ b/src/boot.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use anyhow::Context; -use crate::domain::{Meta, StoredEmbedding, content_hash}; +use crate::datapoints::{Meta, StoredEmbedding, content_hash}; use crate::embed::{Embedder, build_input}; use crate::index::VectorIndex; use crate::store::S3Store; diff --git a/src/datapoints/delete.rs b/src/datapoints/delete.rs new file mode 100644 index 0000000..4612396 --- /dev/null +++ b/src/datapoints/delete.rs @@ -0,0 +1,76 @@ +use futures::{StreamExt, TryStreamExt}; +use uuid::Uuid; + +use super::{DatapointService, datapoint_lock}; +use crate::datapoints::source_uuid; +use crate::error::ServiceError; + +/// Most ids accepted in one delete call. +pub const MAX_DELETE: usize = 1000; + +/// Datapoints to delete, by the caller's `sourceId` or by server id. +#[derive(Debug)] +pub struct DeleteCommand { + pub project_id: Uuid, + pub source_ids: Vec, + pub ids: Vec, +} + +#[derive(Debug)] +pub struct Deleted { + pub id: Uuid, + pub source_id: Option, + /// False when nothing was stored under the id. + pub deleted: bool, +} + +impl DatapointService { + /// Delete datapoints by `sourceId` or server `id`. Results come back in + /// command order, `source_ids` first. + pub async fn delete(&self, cmd: DeleteCommand) -> Result, ServiceError> { + self.infra.ensure_ready()?; + let total = cmd.source_ids.len() + cmd.ids.len(); + if !(1..=MAX_DELETE).contains(&total) { + return Err(ServiceError::Validation(format!( + "send between 1 and {MAX_DELETE} sourceIds and ids in total, got {total}" + ))); + } + if cmd.source_ids.iter().any(|s| s.trim().is_empty()) { + return Err(ServiceError::Validation( + "sourceIds must not be empty".into(), + )); + } + let project = cmd.project_id; + let targets = cmd + .source_ids + .into_iter() + .map(|s| (source_uuid(project, &s), Some(s))) + .chain(cmd.ids.into_iter().map(|id| (id, None))); + let results = futures::stream::iter(targets) + .map(|(id, source_id)| async move { + let deleted = self.delete_one(project, id).await?; + anyhow::Ok(Deleted { + id, + source_id, + deleted, + }) + }) + .buffered(16) + .try_collect() + .await?; + Ok(results) + } + + /// Remove from the index first, so a crash part way leaves a stored + /// datapoint that the next start puts back in the index, never a vector + /// pointing at nothing. + async fn delete_one(&self, project: Uuid, id: Uuid) -> anyhow::Result { + let _guard = self.locks.lock(datapoint_lock(project, id)).await; + if self.infra.store.get_meta(project, id).await?.is_none() { + return Ok(false); + } + self.infra.index.remove(project, id).await?; + self.infra.store.delete_datapoint(project, id).await?; + Ok(true) + } +} diff --git a/src/datapoints/http.rs b/src/datapoints/http.rs new file mode 100644 index 0000000..03aa86c --- /dev/null +++ b/src/datapoints/http.rs @@ -0,0 +1,305 @@ +use axum::Router; +use axum::extract::rejection::QueryRejection; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::routing::{get, post}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use uuid::Uuid; + +use super::{ + BatchItem, Datapoint, DeleteCommand, Deleted, IngestOutcome, Ingested, ListQuery, Listed, + NewDatapoint, +}; +use crate::app::Services; +use crate::error::ServiceError; +use crate::http::{ApiError, Json}; + +pub fn routes() -> Router { + Router::new() + .route("/api/v1/ingest", post(ingest)) + .route("/api/v1/ingest/batch", post(ingest_batch)) + .route("/api/v1/delete", post(delete)) + .route("/api/v1/datapoints", get(list)) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct IngestRequest { + pub idempotency_key: Option, + #[serde(flatten)] + pub datapoint: NewDatapoint, +} + +#[derive(Debug, Serialize)] +pub struct IngestResponse { + #[serde(flatten)] + pub datapoint: Datapoint, + pub deduplicated: bool, + pub replaced: bool, +} + +impl From for IngestResponse { + fn from(i: Ingested) -> Self { + IngestResponse { + datapoint: i.datapoint, + deduplicated: i.outcome == IngestOutcome::Deduplicated, + replaced: i.outcome == IngestOutcome::Replaced, + } + } +} + +fn status_of(outcome: IngestOutcome) -> StatusCode { + match outcome { + IngestOutcome::Created => StatusCode::CREATED, + IngestOutcome::Deduplicated | IngestOutcome::Replaced => StatusCode::OK, + } +} + +pub async fn ingest( + State(services): State, + Json(req): Json, +) -> Result<(StatusCode, axum::Json), ApiError> { + let ingested = services + .datapoints + .ingest(req.datapoint, req.idempotency_key.as_deref()) + .await?; + Ok((status_of(ingested.outcome), axum::Json(ingested.into()))) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchRequest { + pub idempotency_key: Option, + pub items: Vec, +} + +#[derive(Debug, Serialize)] +pub struct ItemErrorBody { + pub code: String, + pub message: String, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +pub enum BatchItemResult { + Ok { + status: u16, + deduplicated: bool, + replaced: bool, + datapoint: Datapoint, + }, + Err { + status: u16, + error: ItemErrorBody, + }, +} + +#[derive(Debug, Serialize)] +pub struct BatchResponse { + pub results: Vec, +} + +fn batch_item_result(index: usize, item: BatchItem) -> BatchItemResult { + match item { + BatchItem::Ingested(ingested) => { + let status = status_of(ingested.outcome).as_u16(); + let resp = IngestResponse::from(ingested); + BatchItemResult::Ok { + status, + deduplicated: resp.deduplicated, + replaced: resp.replaced, + datapoint: resp.datapoint, + } + } + BatchItem::Failed(e) => { + if let ServiceError::Internal(inner) = &e { + tracing::error!(error = ?inner, item = index, "batch item failed"); + } + let e = ApiError::from(e); + BatchItemResult::Err { + status: e.status().as_u16(), + error: ItemErrorBody { + code: e.code().into(), + message: e.to_string(), + }, + } + } + BatchItem::Skipped => BatchItemResult::Err { + status: 0, + error: ItemErrorBody { + code: "skipped".into(), + message: "not attempted after an earlier item failed".into(), + }, + }, + } +} + +pub async fn ingest_batch( + State(services): State, + Json(req): Json, +) -> Result, ApiError> { + let items = services + .datapoints + .ingest_batch(req.items, req.idempotency_key.as_deref()) + .await?; + let results = items + .into_iter() + .enumerate() + .map(|(i, item)| batch_item_result(i, item)) + .collect(); + Ok(axum::Json(BatchResponse { results })) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteRequest { + pub project_id: Uuid, + #[serde(default)] + pub source_ids: Vec, + #[serde(default)] + pub ids: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteResult { + pub id: Uuid, + #[serde(skip_serializing_if = "Option::is_none")] + pub source_id: Option, + /// False when nothing was stored under the id. + pub deleted: bool, +} + +#[derive(Debug, Serialize)] +pub struct DeleteResponse { + pub results: Vec, +} + +impl From for DeleteCommand { + fn from(r: DeleteRequest) -> Self { + DeleteCommand { + project_id: r.project_id, + source_ids: r.source_ids, + ids: r.ids, + } + } +} + +impl From for DeleteResult { + fn from(d: Deleted) -> Self { + DeleteResult { + id: d.id, + source_id: d.source_id, + deleted: d.deleted, + } + } +} + +/// Delete datapoints by `sourceId` or server `id`. Results come back in +/// request order, `sourceIds` first. +pub async fn delete( + State(services): State, + Json(req): Json, +) -> Result, ApiError> { + let deleted = services.datapoints.delete(req.into()).await?; + Ok(axum::Json(DeleteResponse { + results: deleted.into_iter().map(DeleteResult::from).collect(), + })) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ListParams { + pub project_id: Uuid, + pub limit: Option, + pub cursor: Option, + pub source_id: Option, + pub audience: Option, + pub label: Option, + pub occurred_from: Option>, + pub occurred_to: Option>, + #[serde(default)] + pub include_content: bool, +} + +impl From for ListQuery { + fn from(p: ListParams) -> Self { + ListQuery { + project_id: p.project_id, + limit: p.limit, + cursor: p.cursor, + source_id: p.source_id, + audience: p.audience, + label: p.label, + occurred_from: p.occurred_from, + occurred_to: p.occurred_to, + include_content: p.include_content, + } + } +} + +/// A stored datapoint as listed: its metadata without the vector, and its +/// content only when asked for. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DatapointView { + pub id: Uuid, + pub project_id: Uuid, + pub source_id: Option, + pub label: String, + pub timestamp: DateTime, + pub metadata: Map, + pub audience: Option, + pub content_tokens: usize, + pub content_sha256: Option, + pub embedding_model: String, + pub embedding_dim: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, +} + +impl From for DatapointView { + fn from(Listed { meta, content }: Listed) -> Self { + DatapointView { + id: meta.id, + project_id: meta.project_id, + source_id: meta.source_id, + label: meta.label, + timestamp: meta.timestamp, + metadata: meta.metadata, + audience: meta.audience, + content_tokens: meta.content_tokens, + content_sha256: meta.content_sha256, + embedding_model: meta.embedding.model, + embedding_dim: meta.embedding.dim, + content, + } + } +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DatapointList { + pub datapoints: Vec, + /// Pass as `cursor` for the next page; `null` on the last one. + pub next_cursor: Option, +} + +/// List a project's stored datapoints. Works while the index is loading. +pub async fn list( + State(services): State, + query: Result, QueryRejection>, +) -> Result, ApiError> { + let Query(params) = query.map_err(|e| ApiError::Malformed(e.body_text()))?; + let page = services.datapoints.list(params.into()).await?; + Ok(axum::Json(DatapointList { + datapoints: page + .datapoints + .into_iter() + .map(DatapointView::from) + .collect(), + next_cursor: page.next_cursor, + })) +} diff --git a/src/datapoints/ingest.rs b/src/datapoints/ingest.rs new file mode 100644 index 0000000..051c440 --- /dev/null +++ b/src/datapoints/ingest.rs @@ -0,0 +1,275 @@ +use std::collections::HashSet; + +use anyhow::Context; +use chrono::Utc; + +use super::{DatapointService, datapoint_lock}; +use crate::datapoints::{Datapoint, Meta, NewDatapoint, StoredEmbedding, content_hash}; +use crate::embed::build_input; +use crate::error::{ItemError, ServiceError}; +use crate::store::MarkerOutcome; +use crate::tokens::count_tokens; + +/// Most items accepted in one batch. +pub const MAX_BATCH: usize = 100; + +/// What an ingest did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IngestOutcome { + /// A new datapoint was stored. + Created, + /// The datapoint was already stored as sent; nothing changed. + Deduplicated, + /// The stored datapoint with the same `sourceId` was replaced. + Replaced, +} + +#[derive(Debug)] +pub struct Ingested { + pub datapoint: Datapoint, + pub outcome: IngestOutcome, +} + +/// How one item of a batch went. +#[derive(Debug)] +pub enum BatchItem { + Ingested(Ingested), + Failed(ServiceError), + /// Not attempted, because an earlier item failed. + Skipped, +} + +/// An empty key would silently dedupe every later write that also sends one. +fn validate_key(key: Option<&str>) -> Result<(), ServiceError> { + match key { + Some(k) if k.trim().is_empty() => Err(ServiceError::Validation( + "idempotencyKey must not be empty; omit it instead".into(), + )), + _ => Ok(()), + } +} + +impl DatapointService { + /// Store one datapoint, made safe to retry by its `sourceId` or by an + /// idempotency `key`, not both. + pub async fn ingest( + &self, + new: NewDatapoint, + key: Option<&str>, + ) -> Result { + self.infra.ensure_ready()?; + if key.is_some() && new.source_id.is_some() { + return Err(ServiceError::Validation( + "send idempotencyKey or sourceId, not both; resending a sourceId is already safe" + .into(), + )); + } + self.ingest_one(new, key).await + } + + /// Store up to [`MAX_BATCH`] datapoints in order, stopping at the first + /// failure. Every item is validated before anything is written. Item + /// `i` uses the idempotency key `"{key}#{i}"`. + pub async fn ingest_batch( + &self, + items: Vec, + key: Option<&str>, + ) -> Result, ServiceError> { + self.infra.ensure_ready()?; + validate_key(key)?; + if items.len() > MAX_BATCH { + return Err(ServiceError::BatchTooLarge); + } + let details = self.invalid_items(&items).await?; + if !details.is_empty() { + return Err(ServiceError::ValidationDetailed { + message: "invalid items".into(), + details, + }); + } + + let total = items.len(); + let mut results = Vec::with_capacity(total); + for (i, item) in items.into_iter().enumerate() { + let item_key = key.map(|k| format!("{k}#{i}")); + match self.ingest_one(item, item_key.as_deref()).await { + Ok(ingested) => results.push(BatchItem::Ingested(ingested)), + Err(e) => { + results.push(BatchItem::Failed(e)); + results.extend((i + 1..total).map(|_| BatchItem::Skipped)); + break; + } + } + } + Ok(results) + } + + /// Why each item would fail, checked before anything is written. An item + /// with a `sourceId` and no label needs a stored datapoint, or an earlier + /// item in the batch that creates it, to take its label from. + async fn invalid_items(&self, items: &[NewDatapoint]) -> anyhow::Result> { + let mut details = Vec::new(); + let mut labelled = HashSet::new(); + for (index, item) in items.iter().enumerate() { + if let Err(message) = item.validate() { + details.push(ItemError { index, message }); + continue; + } + let Some(source) = &item.source_id else { + continue; + }; + let id = item.id(); + if item.label.is_some() { + labelled.insert(id); + } else if !labelled.contains(&id) + && self + .infra + .store + .get_meta(item.project_id, id) + .await? + .is_none() + { + details.push(ItemError { + index, + message: format!( + "label is required: no datapoint with sourceId {source:?} is stored" + ), + }); + } + } + Ok(details) + } + + /// Store one datapoint. With a `sourceId`, the datapoint is stored under + /// it and `key` is ignored. With an idempotency key, a repeat of the same + /// request returns the original datapoint, a half-finished earlier attempt + /// is completed under the id it already reserved, and a different request + /// reusing the key is rejected. + pub(crate) async fn ingest_one( + &self, + new: NewDatapoint, + key: Option<&str>, + ) -> Result { + validate_key(key)?; + new.validate().map_err(ServiceError::Validation)?; + if new.source_id.is_some() { + return self.upsert(new).await; + } + let infra = &self.infra; + let project = new.project_id; + let fingerprint = new.fingerprint(); + let mut id = new.id(); + + let _guard = match key { + Some(key) => Some(self.locks.lock(format!("{project}:key:{key}")).await), + None => None, + }; + if let Some(key) = key + && let MarkerOutcome::Exists { + id: existing_id, + fingerprint: recorded, + } = infra + .store + .create_marker(project, key, id, &fingerprint) + .await? + { + // Markers from before fingerprints existed are trusted as before. + if recorded.is_some_and(|f| f != fingerprint) { + return Err(ServiceError::IdempotencyKeyReused); + } + if let Some(meta) = infra.store.get_meta(project, existing_id).await? { + let content = infra + .store + .get_content(project, existing_id) + .await? + .with_context(|| { + format!("datapoint {existing_id} has metadata but no content") + })?; + // The earlier attempt may have stopped after meta.json but + // before the index upsert. Upserting again is harmless. + if meta.embedding.model == infra.embedder.model_name() { + infra + .index + .upsert(project, existing_id, meta.embedding.vector.clone()) + .await?; + } + return Ok(Ingested { + datapoint: meta.into_datapoint(content), + outcome: IngestOutcome::Deduplicated, + }); + } + // An earlier attempt reserved this id but did not finish. + id = existing_id; + } + + let datapoint = new + .into_parts(id, None, Utc::now()) + .map_err(ServiceError::Validation)?; + Ok(Ingested { + datapoint: self.write(datapoint).await?, + outcome: IngestOutcome::Created, + }) + } + + /// Store a datapoint under its `sourceId`. A resend that changes nothing + /// returns the stored datapoint; one that changes anything replaces it. + async fn upsert(&self, new: NewDatapoint) -> Result { + let project = new.project_id; + let id = new.id(); + let _guard = self.locks.lock(datapoint_lock(project, id)).await; + let existing = self.infra.store.get_meta(project, id).await?; + let datapoint = new + .into_parts(id, existing.as_ref(), Utc::now()) + .map_err(ServiceError::Validation)?; + let (outcome, datapoint) = match existing { + Some(meta) if meta.matches(&datapoint, self.infra.embedder.model_name()) => { + // An earlier send may have stopped before the index upsert. + self.infra + .index + .upsert(project, id, meta.embedding.vector) + .await?; + (IngestOutcome::Deduplicated, datapoint) + } + Some(_) => (IngestOutcome::Replaced, self.write(datapoint).await?), + None => (IngestOutcome::Created, self.write(datapoint).await?), + }; + Ok(Ingested { datapoint, outcome }) + } + + /// Embed, then write content, then metadata (the commit point), then index. + async fn write(&self, datapoint: Datapoint) -> anyhow::Result { + let infra = &self.infra; + let id = datapoint.id; + let input = build_input(&datapoint.label, &datapoint.content, &datapoint.metadata); + let embedder = infra.embedder.clone(); + let vector = tokio::task::spawn_blocking(move || embedder.embed_passages(&[input])) + .await + .context("embedding task panicked")?? + .pop() + .context("embedder returned no vector")?; + + let meta = Meta { + id, + project_id: datapoint.project_id, + source_id: datapoint.source_id.clone(), + label: datapoint.label.clone(), + timestamp: datapoint.timestamp, + metadata: datapoint.metadata.clone(), + audience: datapoint.audience.clone(), + embedding: StoredEmbedding { + model: infra.embedder.model_name().to_string(), + dim: vector.len(), + vector: vector.clone(), + }, + content_tokens: count_tokens(&datapoint.content), + content_sha256: Some(content_hash(&datapoint.content)), + }; + infra + .store + .put_content(datapoint.project_id, id, &datapoint.content) + .await?; + infra.store.put_meta(&meta).await?; + infra.index.upsert(datapoint.project_id, id, vector).await?; + Ok(datapoint) + } +} diff --git a/src/datapoints/list.rs b/src/datapoints/list.rs new file mode 100644 index 0000000..b445480 --- /dev/null +++ b/src/datapoints/list.rs @@ -0,0 +1,208 @@ +use chrono::{DateTime, Utc}; +use futures::{StreamExt, TryStreamExt}; +use uuid::Uuid; + +use super::DatapointService; +use crate::datapoints::{Meta, PROJECT_AUDIENCE, source_uuid, validate_audience}; +use crate::error::ServiceError; + +pub const DEFAULT_PAGE_LIMIT: usize = 50; +pub const MAX_PAGE_LIMIT: usize = 200; +/// Most datapoints read per request. With narrow filters a page can come +/// back short, with a cursor to carry on from. +pub const MAX_SCAN: usize = 1000; +/// Datapoints read per round when filtering. +const CHUNK: usize = 100; +const FETCH_CONCURRENCY: usize = 16; + +/// Which of a project's datapoints to list, and how many. +#[derive(Debug)] +pub struct ListQuery { + pub project_id: Uuid, + pub limit: Option, + pub cursor: Option, + pub source_id: Option, + pub audience: Option, + pub label: Option, + pub occurred_from: Option>, + pub occurred_to: Option>, + pub include_content: bool, +} + +/// A stored datapoint as listed: its metadata, and its content only when +/// asked for. +#[derive(Debug)] +pub struct Listed { + pub meta: Meta, + pub content: Option, +} + +#[derive(Debug)] +pub struct Page { + pub datapoints: Vec, + /// Where the next page starts; `None` on the last one. + pub next_cursor: Option, +} + +/// Which stored datapoints to list. +#[derive(Debug, Default)] +struct Filter { + /// `Some(None)` is the whole-project audience. + audience: Option>, + /// Lowercased. + label: Option, + from: Option>, + to: Option>, +} + +impl Filter { + fn from_query(q: &ListQuery) -> Result { + if let Some(audience) = &q.audience { + validate_audience(audience).map_err(ServiceError::Validation)?; + } + if q.label.as_ref().is_some_and(|l| l.trim().is_empty()) { + return Err(ServiceError::Validation( + "label must not be empty; omit it instead".into(), + )); + } + if let (Some(from), Some(to)) = (q.occurred_from, q.occurred_to) + && from > to + { + return Err(ServiceError::RecallRangeInvalid); + } + Ok(Filter { + audience: q + .audience + .clone() + .map(|a| (a != PROJECT_AUDIENCE).then_some(a)), + label: q.label.as_ref().map(|l| l.to_lowercase()), + from: q.occurred_from, + to: q.occurred_to, + }) + } + + fn is_empty(&self) -> bool { + self.audience.is_none() && self.label.is_none() && self.from.is_none() && self.to.is_none() + } + + fn matches(&self, meta: &Meta) -> bool { + self.audience.as_ref().is_none_or(|a| *a == meta.audience) + && self + .label + .as_ref() + .is_none_or(|l| meta.label.to_lowercase().contains(l)) + && self.from.is_none_or(|from| meta.timestamp >= from) + && self.to.is_none_or(|to| meta.timestamp <= to) + } +} + +impl DatapointService { + /// List a project's stored datapoints, read straight from storage, in id + /// order. Meant for debugging: it works while the index is still loading. + pub async fn list(&self, q: ListQuery) -> Result { + let limit = q.limit.unwrap_or(DEFAULT_PAGE_LIMIT); + if !(1..=MAX_PAGE_LIMIT).contains(&limit) { + return Err(ServiceError::Validation(format!( + "limit must be between 1 and {MAX_PAGE_LIMIT}" + ))); + } + if q.source_id.as_ref().is_some_and(|s| s.trim().is_empty()) { + return Err(ServiceError::Validation( + "sourceId must not be empty; omit it instead".into(), + )); + } + let cursor = q + .cursor + .as_deref() + .map(|c| { + c.parse::() + .map_err(|_| ServiceError::Validation(format!("cursor {c:?} is not valid"))) + }) + .transpose()?; + let filter = Filter::from_query(&q)?; + let project = q.project_id; + let store = &self.infra.store; + + let (metas, next) = match &q.source_id { + Some(source_id) => { + let meta = store + .get_meta(project, source_uuid(project, source_id)) + .await?; + ( + meta.into_iter().filter(|m| filter.matches(m)).collect(), + None, + ) + } + None => self.scan(project, cursor, limit, &filter).await?, + }; + + let include_content = q.include_content; + let datapoints = futures::stream::iter(metas) + .map(|meta| async move { + let content = if include_content { + store.get_content(project, meta.id).await? + } else { + None + }; + anyhow::Ok(Listed { meta, content }) + }) + .buffered(FETCH_CONCURRENCY) + .try_collect() + .await?; + Ok(Page { + datapoints, + next_cursor: next.map(|id| id.to_string()), + }) + } + + /// Up to `limit` matching datapoints after `after`, reading at most + /// [`MAX_SCAN`], and the id to continue after when there may be more. + async fn scan( + &self, + project: Uuid, + mut after: Option, + limit: usize, + filter: &Filter, + ) -> anyhow::Result<(Vec, Option)> { + let store = &self.infra.store; + let mut found = Vec::new(); + let mut scanned = 0; + loop { + // Without a filter every datapoint read is returned, so read no + // more than the page needs. + let want = if filter.is_empty() { + limit - found.len() + } else { + CHUNK + }; + let chunk = want.min(MAX_SCAN - scanned); + let ids = store.list_datapoint_ids(project, after, chunk).await?; + let exhausted = ids.len() < chunk; + let metas: Vec> = futures::stream::iter(ids.iter().copied()) + .map(|id| store.get_meta(project, id)) + .buffered(FETCH_CONCURRENCY) + .try_collect() + .await?; + let last = ids.last().copied(); + for (id, meta) in ids.into_iter().zip(metas) { + after = Some(id); + scanned += 1; + // Gone since listing, or filtered out. + let Some(meta) = meta.filter(|m| filter.matches(m)) else { + continue; + }; + found.push(meta); + if found.len() == limit { + let more = !(exhausted && Some(id) == last); + return Ok((found, more.then_some(id))); + } + } + if exhausted { + return Ok((found, None)); + } + if scanned >= MAX_SCAN { + return Ok((found, after)); + } + } + } +} diff --git a/src/datapoints/mod.rs b/src/datapoints/mod.rs new file mode 100644 index 0000000..93ff80f --- /dev/null +++ b/src/datapoints/mod.rs @@ -0,0 +1,37 @@ +//! Datapoints: the memories a project stores. Ingest stores or replaces +//! them, delete removes them, and listing shows what is stored. + +pub mod http; + +mod delete; +mod ingest; +mod list; +mod model; + +pub use delete::{DeleteCommand, Deleted, MAX_DELETE}; +pub use ingest::{BatchItem, IngestOutcome, Ingested, MAX_BATCH}; +pub use list::{DEFAULT_PAGE_LIMIT, ListQuery, Listed, MAX_PAGE_LIMIT, MAX_SCAN, Page}; +pub use model::*; + +use uuid::Uuid; + +use crate::app::Infra; +use crate::locks::KeyLocks; + +/// Stores, replaces, deletes and lists a project's datapoints. +#[derive(Clone)] +pub struct DatapointService { + infra: Infra, + locks: KeyLocks, +} + +impl DatapointService { + pub(crate) fn new(infra: Infra, locks: KeyLocks) -> DatapointService { + DatapointService { infra, locks } + } +} + +/// The in-process lock serialising writes and deletes of one datapoint. +fn datapoint_lock(project: Uuid, id: Uuid) -> String { + format!("{project}:datapoint:{id}") +} diff --git a/src/domain.rs b/src/datapoints/model.rs similarity index 100% rename from src/domain.rs rename to src/datapoints/model.rs diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..2ac5f50 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,37 @@ +/// Why one item of a batch is invalid. +#[derive(Debug)] +pub struct ItemError { + pub index: usize, + pub message: String, +} + +/// Why a use case failed, in domain terms. The API maps these to statuses +/// and error codes. +#[derive(Debug, thiserror::Error)] +pub enum ServiceError { + #[error("{0}")] + Validation(String), + #[error("{message}")] + ValidationDetailed { + message: String, + details: Vec, + }, + #[error("occurredFrom must not be after occurredTo")] + RecallRangeInvalid, + #[error("send level or maxTokens, not both")] + BudgetConflict, + #[error("a batch holds at most 100 items")] + BatchTooLarge, + #[error("idempotencyKey was already used with a different request; use a new key")] + IdempotencyKeyReused, + #[error("{0}")] + NotFound(String), + #[error("this webhook is paused")] + WebhookPaused, + #[error("too many events for this webhook; at most 600 a minute")] + RateLimited, + #[error("index is still loading from storage")] + Loading, + #[error("internal error")] + Internal(#[from] anyhow::Error), +} diff --git a/src/http/error.rs b/src/http/error.rs new file mode 100644 index 0000000..d6baf01 --- /dev/null +++ b/src/http/error.rs @@ -0,0 +1,173 @@ +use axum::extract::rejection::JsonRejection; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; + +use crate::error::ServiceError; + +/// A failed request: a service error, or a request that never reached a +/// service. +#[derive(Debug, thiserror::Error)] +pub enum ApiError { + #[error(transparent)] + Service(#[from] ServiceError), + /// A body, path or query string that does not parse. + #[error("{0}")] + Malformed(String), + #[error("request body is larger than 10 MiB")] + PayloadTooLarge, + #[error("webhook body is larger than 1 MiB")] + WebhookBodyTooLarge, +} + +impl ApiError { + pub fn code(&self) -> &'static str { + match self { + ApiError::Service(e) => match e { + ServiceError::Validation(_) | ServiceError::ValidationDetailed { .. } => { + "validation" + } + ServiceError::RecallRangeInvalid => "recall-range-invalid", + ServiceError::BudgetConflict => "budget-conflict", + ServiceError::BatchTooLarge => "batch-too-large", + ServiceError::IdempotencyKeyReused => "idempotency-key-reused", + ServiceError::NotFound(_) => "not-found", + ServiceError::WebhookPaused => "webhook-paused", + ServiceError::RateLimited => "rate-limited", + ServiceError::Loading => "loading", + ServiceError::Internal(_) => "internal", + }, + ApiError::Malformed(_) => "validation", + ApiError::PayloadTooLarge => "payload-too-large", + ApiError::WebhookBodyTooLarge => "webhook-body-too-large", + } + } + + pub fn status(&self) -> StatusCode { + match self { + ApiError::Service(e) => match e { + ServiceError::Validation(_) + | ServiceError::ValidationDetailed { .. } + | ServiceError::RecallRangeInvalid + | ServiceError::BudgetConflict + | ServiceError::BatchTooLarge => StatusCode::BAD_REQUEST, + ServiceError::IdempotencyKeyReused => StatusCode::UNPROCESSABLE_ENTITY, + ServiceError::NotFound(_) => StatusCode::NOT_FOUND, + ServiceError::WebhookPaused => StatusCode::CONFLICT, + ServiceError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + ServiceError::Loading => StatusCode::SERVICE_UNAVAILABLE, + ServiceError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + }, + ApiError::Malformed(_) => StatusCode::BAD_REQUEST, + ApiError::PayloadTooLarge | ApiError::WebhookBodyTooLarge => { + StatusCode::PAYLOAD_TOO_LARGE + } + } + } + + pub fn to_body(&self) -> serde_json::Value { + let mut error = serde_json::json!({ "code": self.code(), "message": self.to_string() }); + if let ApiError::Service(ServiceError::ValidationDetailed { details, .. }) = self { + error["details"] = details + .iter() + .map(|d| serde_json::json!({ "index": d.index, "message": d.message })) + .collect(); + } + serde_json::json!({ "error": error }) + } +} + +impl IntoResponse for ApiError { + fn into_response(self) -> Response { + if let ApiError::Service(ServiceError::Internal(e)) = &self { + tracing::error!(error = ?e, "internal error"); + } + (self.status(), axum::Json(self.to_body())).into_response() + } +} + +impl From for ApiError { + fn from(r: JsonRejection) -> Self { + if r.status() == StatusCode::PAYLOAD_TOO_LARGE { + return ApiError::PayloadTooLarge; + } + ApiError::Malformed(r.body_text()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::ItemError; + use axum::http::StatusCode; + use axum::response::IntoResponse; + use http_body_util::BodyExt; + + async fn body_json(e: impl Into) -> (StatusCode, serde_json::Value) { + let resp = e.into().into_response(); + let status = resp.status(); + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + (status, serde_json::from_slice(&bytes).unwrap()) + } + + #[tokio::test] + async fn envelope_shape_and_status() { + let (status, body) = body_json(ServiceError::RecallRangeInvalid).await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["error"]["code"], "recall-range-invalid"); + assert!(body["error"]["message"].is_string()); + } + + #[tokio::test] + async fn internal_hides_details() { + let (status, body) = + body_json(ServiceError::Internal(anyhow::anyhow!("secret s3 thing"))).await; + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(body["error"]["message"], "internal error"); + assert_eq!(body["error"]["code"], "internal"); + } + + #[tokio::test] + async fn detailed_validation_includes_details() { + let e = ServiceError::ValidationDetailed { + message: "bad items".into(), + details: vec![ItemError { + index: 2, + message: "label empty".into(), + }], + }; + let (status, body) = body_json(e).await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["error"]["code"], "validation"); + assert_eq!(body["error"]["details"][0]["index"], 2); + assert_eq!(body["error"]["details"][0]["message"], "label empty"); + } + + #[tokio::test] + async fn codes_and_statuses() { + let api = ApiError::from; + assert_eq!( + api(ServiceError::Loading).status(), + StatusCode::SERVICE_UNAVAILABLE + ); + assert_eq!(api(ServiceError::Loading).code(), "loading"); + assert_eq!( + api(ServiceError::IdempotencyKeyReused).code(), + "idempotency-key-reused" + ); + assert_eq!( + api(ServiceError::IdempotencyKeyReused).status(), + StatusCode::UNPROCESSABLE_ENTITY + ); + assert_eq!(api(ServiceError::BudgetConflict).code(), "budget-conflict"); + assert_eq!(api(ServiceError::BatchTooLarge).code(), "batch-too-large"); + assert_eq!( + api(ServiceError::Validation("x".into())).code(), + "validation" + ); + assert_eq!(ApiError::Malformed("x".into()).code(), "validation"); + assert_eq!( + ApiError::PayloadTooLarge.status(), + StatusCode::PAYLOAD_TOO_LARGE + ); + } +} diff --git a/src/http/mod.rs b/src/http/mod.rs new file mode 100644 index 0000000..ca4e037 --- /dev/null +++ b/src/http/mod.rs @@ -0,0 +1,38 @@ +//! What every domain's HTTP handlers share: the router, the JSON extractor +//! and the error envelope. The handlers live with their domain, in its +//! `http` module. + +mod error; + +use axum::Router; +use axum::extract::{DefaultBodyLimit, FromRequest, State}; +use axum::routing::get; +use tower_http::trace::TraceLayer; + +use crate::app::Services; +use crate::{datapoints, recall, webhooks}; + +pub use error::ApiError; + +#[derive(FromRequest)] +#[from_request(via(axum::Json), rejection(ApiError))] +pub struct Json(pub T); + +/// Largest request body accepted, in bytes. +pub const MAX_BODY_BYTES: usize = 10 * 1024 * 1024; + +async fn healthz(State(services): State) -> axum::Json { + let status = if services.is_ready() { "ok" } else { "loading" }; + axum::Json(serde_json::json!({ "status": status })) +} + +pub fn router(services: Services) -> Router { + Router::new() + .route("/healthz", get(healthz)) + .merge(datapoints::http::routes()) + .merge(recall::http::routes()) + .merge(webhooks::http::routes()) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) + .layer(TraceLayer::new_for_http()) + .with_state(services) +} diff --git a/src/lib.rs b/src/lib.rs index 6051f44..815e2a6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,14 @@ -pub mod api; +pub mod app; pub mod aws_error; pub mod boot; pub mod config; -pub mod domain; +pub mod datapoints; pub mod embed; +pub mod error; +pub mod http; pub mod index; +pub mod locks; pub mod recall; pub mod store; pub mod tokens; -pub mod webhook; +pub mod webhooks; diff --git a/src/api/locks.rs b/src/locks.rs similarity index 100% rename from src/api/locks.rs rename to src/locks.rs diff --git a/src/main.rs b/src/main.rs index 9033106..5dd2388 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use anyhow::Context; -use muninn::api::AppState; -use muninn::api::locks::KeyLocks; +use muninn::app::{Infra, Services}; use muninn::boot; use muninn::config::{Config, VectorStore}; use muninn::embed::Embedder; @@ -56,24 +55,22 @@ async fn main() -> anyhow::Result<()> { VectorStore::Memory => Arc::new(MemoryIndex::new()), }; - let state = AppState { + let infra = Infra { store, index, embedder: Arc::new(embedder), ready: Arc::new(AtomicBool::new(false)), - locks: KeyLocks::default(), - webhook_rate: Default::default(), }; let listener = tokio::net::TcpListener::bind(config.bind).await?; tracing::info!(bind = %config.bind, "listening; rebuilding index"); boot::spawn_rebuild( - state.store.clone(), - state.index.clone(), - state.embedder.clone(), - state.ready.clone(), + infra.store.clone(), + infra.index.clone(), + infra.embedder.clone(), + infra.ready.clone(), ); - axum::serve(listener, muninn::api::router(state)) + axum::serve(listener, muninn::http::router(Services::new(infra))) .with_graceful_shutdown(shutdown_signal()) .await?; tracing::info!("shut down"); diff --git a/src/recall/http.rs b/src/recall/http.rs new file mode 100644 index 0000000..1e86c73 --- /dev/null +++ b/src/recall/http.rs @@ -0,0 +1,86 @@ +use axum::Router; +use axum::extract::State; +use axum::routing::post; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use super::{Level, RecallQuery, RecalledItem, Sort}; +use crate::app::Services; +use crate::datapoints::Datapoint; +use crate::http::{ApiError, Json}; + +pub fn routes() -> Router { + Router::new().route("/api/v1/recall", post(handler)) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RecallRequest { + pub project_id: Uuid, + pub query: String, + pub limit: Option, + pub level: Option, + pub max_tokens: Option, + pub query_timestamp: Option>, + pub occurred_from: Option>, + pub occurred_to: Option>, + pub sort: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RecallItem { + pub rank: usize, + pub score: f32, + #[serde(flatten)] + pub datapoint: Datapoint, + pub content_tokens: usize, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RecallResponse { + pub results: Vec, + pub used_tokens: usize, + pub budget_tokens: usize, +} + +impl From for RecallQuery { + fn from(r: RecallRequest) -> Self { + RecallQuery { + project_id: r.project_id, + query: r.query, + limit: r.limit, + level: r.level, + max_tokens: r.max_tokens, + query_timestamp: r.query_timestamp, + occurred_from: r.occurred_from, + occurred_to: r.occurred_to, + sort: r.sort, + } + } +} + +impl From for RecallItem { + fn from(i: RecalledItem) -> Self { + RecallItem { + rank: i.rank, + score: i.score, + datapoint: i.datapoint, + content_tokens: i.content_tokens, + } + } +} + +pub async fn handler( + State(services): State, + Json(req): Json, +) -> Result, ApiError> { + let recalled = services.recall.recall(req.into()).await?; + Ok(axum::Json(RecallResponse { + results: recalled.results.into_iter().map(RecallItem::from).collect(), + used_tokens: recalled.used_tokens, + budget_tokens: recalled.budget_tokens, + })) +} diff --git a/src/recall/mod.rs b/src/recall/mod.rs new file mode 100644 index 0000000..11d53dd --- /dev/null +++ b/src/recall/mod.rs @@ -0,0 +1,258 @@ +//! Recall: the stored datapoints most relevant to a query, within a token +//! budget. + +pub mod http; + +mod select; + +pub use select::{Candidate, Level, SelectParams, Selected, Selection, Sort, select}; + +use std::collections::HashMap; + +use anyhow::Context; +use chrono::{DateTime, Utc}; +use futures::{StreamExt, TryStreamExt}; +use uuid::Uuid; + +use crate::app::Infra; +use crate::datapoints::{Datapoint, Meta}; +use crate::error::ServiceError; + +pub const DEFAULT_LIMIT: usize = 100; +pub const MAX_LIMIT: usize = 200; +/// Most hits fetched from the index before time filtering. +const MAX_SEARCH_K: usize = 1000; +const FETCH_CONCURRENCY: usize = 16; + +/// What to recall from a project. +#[derive(Debug)] +pub struct RecallQuery { + pub project_id: Uuid, + pub query: String, + pub limit: Option, + pub level: Option, + pub max_tokens: Option, + pub query_timestamp: Option>, + pub occurred_from: Option>, + pub occurred_to: Option>, + pub sort: Option, +} + +#[derive(Debug)] +pub struct RecalledItem { + pub rank: usize, + pub score: f32, + pub datapoint: Datapoint, + pub content_tokens: usize, +} + +#[derive(Debug)] +pub struct RecallResult { + pub results: Vec, + pub used_tokens: usize, + pub budget_tokens: usize, +} + +/// The token budget: `level` or `maxTokens`, never both; `low` by default. +pub fn resolve_budget( + level: Option, + max_tokens: Option, +) -> Result { + match (level, max_tokens) { + (Some(_), Some(_)) => Err(ServiceError::BudgetConflict), + (_, Some(0)) => Err(ServiceError::Validation( + "maxTokens must be at least 1".into(), + )), + (_, Some(n)) => Ok(n), + (level, None) => Ok(level.unwrap_or_default().tokens()), + } +} + +pub fn validate(q: &RecallQuery) -> Result<(), ServiceError> { + if q.query.trim().is_empty() { + return Err(ServiceError::Validation("query must not be empty".into())); + } + if let Some(limit) = q.limit + && !(1..=MAX_LIMIT).contains(&limit) + { + return Err(ServiceError::Validation(format!( + "limit must be between 1 and {MAX_LIMIT}" + ))); + } + if let (Some(from), Some(to)) = (q.occurred_from, q.occurred_to) + && from > to + { + return Err(ServiceError::RecallRangeInvalid); + } + Ok(()) +} + +/// Finds the stored datapoints most relevant to a query, within a token +/// budget. +#[derive(Clone)] +pub struct RecallService { + infra: Infra, +} + +impl RecallService { + pub(crate) fn new(infra: Infra) -> RecallService { + RecallService { infra } + } + + pub async fn recall(&self, q: RecallQuery) -> Result { + self.infra.ensure_ready()?; + validate(&q)?; + let budget_tokens = resolve_budget(q.level, q.max_tokens)?; + let limit = q.limit.unwrap_or(DEFAULT_LIMIT); + let project = q.project_id; + + let embedder = self.infra.embedder.clone(); + let query = q.query; + let qvec = tokio::task::spawn_blocking(move || embedder.embed_query(&query)) + .await + .context("embedding task panicked")??; + + let hits = self + .infra + .index + .search(project, &qvec, (limit * 4).min(MAX_SEARCH_K)) + .await?; + + let store = &self.infra.store; + let metas: Vec<(f32, Option)> = futures::stream::iter(hits) + .map(|(id, score)| async move { store.get_meta(project, id).await.map(|m| (score, m)) }) + .buffer_unordered(FETCH_CONCURRENCY) + .try_collect() + .await?; + let mut by_id: HashMap = HashMap::new(); + let mut candidates = Vec::with_capacity(metas.len()); + for (score, meta) in metas { + let Some(meta) = meta else { continue }; + candidates.push(Candidate { + id: meta.id, + score, + timestamp: meta.timestamp, + content_tokens: meta.content_tokens, + }); + by_id.insert(meta.id, meta); + } + + let selected = select( + candidates, + &SelectParams { + limit, + budget_tokens, + query_timestamp: q.query_timestamp.unwrap_or_else(Utc::now), + occurred_from: q.occurred_from, + occurred_to: q.occurred_to, + sort: q.sort.unwrap_or_default(), + }, + ); + + let contents: HashMap> = + futures::stream::iter(selected.items.iter().map(|s| s.id).collect::>()) + .map(|id| async move { store.get_content(project, id).await.map(|c| (id, c)) }) + .buffer_unordered(FETCH_CONCURRENCY) + .try_collect() + .await?; + + let mut used_tokens = selected.used_tokens; + let mut results = Vec::with_capacity(selected.items.len()); + for s in selected.items { + let meta = by_id + .remove(&s.id) + .context("selected id without metadata")?; + let content_tokens = meta.content_tokens; + let Some(content) = contents.get(&s.id).cloned().flatten() else { + // Only reachable if the content object was removed outside the + // service. Skip it rather than fail every query that finds it. + tracing::warn!(project = %project, id = %s.id, "metadata without content; skipping"); + used_tokens -= content_tokens; + continue; + }; + results.push(RecalledItem { + rank: s.rank, + score: s.score, + datapoint: meta.into_datapoint(content), + content_tokens, + }); + } + + Ok(RecallResult { + results, + used_tokens, + budget_tokens, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn query() -> RecallQuery { + RecallQuery { + project_id: Uuid::now_v7(), + query: "q".into(), + limit: None, + level: None, + max_tokens: None, + query_timestamp: None, + occurred_from: None, + occurred_to: None, + sort: None, + } + } + + #[test] + fn default_budget_is_8000() { + assert_eq!(resolve_budget(None, None).unwrap(), 8000); + assert_eq!(resolve_budget(Some(Level::High), None).unwrap(), 24000); + assert_eq!(resolve_budget(None, Some(50)).unwrap(), 50); + } + + #[test] + fn both_is_conflict() { + assert!(matches!( + resolve_budget(Some(Level::Low), Some(5)), + Err(ServiceError::BudgetConflict) + )); + } + + #[test] + fn zero_max_tokens_rejected() { + assert!(matches!( + resolve_budget(None, Some(0)), + Err(ServiceError::Validation(_)) + )); + } + + #[test] + fn range_invalid() { + let mut q = query(); + q.occurred_from = Some("2026-10-02T00:00:00Z".parse().unwrap()); + q.occurred_to = Some("2026-10-01T00:00:00Z".parse().unwrap()); + assert!(matches!( + validate(&q), + Err(ServiceError::RecallRangeInvalid) + )); + q.occurred_to = q.occurred_from; + assert!(validate(&q).is_ok()); + } + + #[test] + fn limit_bounds() { + for (limit, ok) in [(0, false), (201, false), (1, true), (200, true)] { + let mut q = query(); + q.limit = Some(limit); + assert_eq!(validate(&q).is_ok(), ok, "limit {limit}"); + } + } + + #[test] + fn blank_query_rejected() { + let mut q = query(); + q.query = " \n".into(); + assert!(matches!(validate(&q), Err(ServiceError::Validation(_)))); + } +} diff --git a/src/recall.rs b/src/recall/select.rs similarity index 100% rename from src/recall.rs rename to src/recall/select.rs diff --git a/src/store.rs b/src/store.rs index 4f34f12..126c4dc 100644 --- a/src/store.rs +++ b/src/store.rs @@ -12,8 +12,8 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::config::Config; -use crate::domain::Meta; -use crate::webhook::Webhook; +use crate::datapoints::Meta; +use crate::webhooks::Webhook; /// Object key layout. See ARCHITECTURE.md. pub mod keys { diff --git a/src/webhooks/http.rs b/src/webhooks/http.rs new file mode 100644 index 0000000..fadcbe9 --- /dev/null +++ b/src/webhooks/http.rs @@ -0,0 +1,187 @@ +use axum::Router; +use axum::body::Bytes; +use axum::extract::rejection::{BytesRejection, QueryRejection}; +use axum::extract::{DefaultBodyLimit, Path, Query, State}; +use axum::http::StatusCode; +use axum::routing::{self, post}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use super::{EventOutcome, MAX_BODY_BYTES, Settings, Webhook, WebhookView}; +use crate::app::Services; +use crate::http::{ApiError, Json}; + +pub fn routes() -> Router { + Router::new() + .route("/api/v1/webhooks", post(create).get(list)) + .route( + "/api/v1/webhooks/{id}", + routing::get(get).put(update).delete(delete), + ) + .route( + "/api/v1/webhooks/{id}/{token}", + post(receive).layer(DefaultBodyLimit::max(MAX_BODY_BYTES)), + ) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WebhookRequest { + pub project_id: Uuid, + #[serde(flatten)] + pub settings: Settings, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProjectQuery { + pub project_id: Uuid, +} + +#[derive(Debug, Serialize)] +pub struct CreatedWebhook { + #[serde(flatten)] + pub webhook: WebhookView, + /// Where to send events. The token in it is the only credential and is + /// not shown again. + pub path: String, +} + +#[derive(Debug, Serialize)] +pub struct WebhookList { + pub webhooks: Vec, +} + +#[derive(Debug, Serialize)] +pub struct DeletedWebhook { + pub id: String, + pub deleted: bool, +} + +fn project_of(query: Result, QueryRejection>) -> Result { + query + .map(|Query(q)| q.project_id) + .map_err(|e| ApiError::Malformed(e.body_text())) +} + +pub async fn create( + State(services): State, + Json(req): Json, +) -> Result<(StatusCode, axum::Json), ApiError> { + let (webhook, token) = services + .webhooks + .create(req.project_id, req.settings) + .await?; + let path = format!("/api/v1/webhooks/{}/{token}", webhook.id); + Ok(( + StatusCode::CREATED, + axum::Json(CreatedWebhook { + webhook: webhook.view(), + path, + }), + )) +} + +pub async fn list( + State(services): State, + query: Result, QueryRejection>, +) -> Result, ApiError> { + let project = project_of(query)?; + let webhooks = services.webhooks.list(project).await?; + Ok(axum::Json(WebhookList { + webhooks: webhooks.into_iter().map(Webhook::view).collect(), + })) +} + +pub async fn get( + State(services): State, + Path(id): Path, + query: Result, QueryRejection>, +) -> Result, ApiError> { + let project = project_of(query)?; + Ok(axum::Json( + services.webhooks.get(project, &id).await?.view(), + )) +} + +/// Replace a webhook's settings. Its id and token stay the same. +pub async fn update( + State(services): State, + Path(id): Path, + Json(req): Json, +) -> Result, ApiError> { + let webhook = services + .webhooks + .update(req.project_id, &id, req.settings) + .await?; + Ok(axum::Json(webhook.view())) +} + +/// Delete a webhook. Datapoints it created stay. +pub async fn delete( + State(services): State, + Path(id): Path, + query: Result, QueryRejection>, +) -> Result, ApiError> { + let project = project_of(query)?; + let deleted = services.webhooks.delete(project, &id).await?; + Ok(axum::Json(DeletedWebhook { id, deleted })) +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct EventResponse { + pub event_id: String, + /// `ingested` or `dropped`. + pub verdict: &'static str, + /// Why the event was dropped: the first filter it failed, or the field + /// it lacked. + pub dropped_by: Option, + pub datapoint_id: Option, + /// True when the event changed nothing that was already stored. + pub unchanged: bool, +} + +/// Turn an event into a datapoint. The webhook is checked before the body, +/// so a wrong id or token is a `404` whatever was sent. +pub async fn receive( + State(services): State, + Path((id, token)): Path<(String, String)>, + body: Result, +) -> Result<(StatusCode, axum::Json), ApiError> { + let webhook = services.webhooks.authenticate(&id, &token).await?; + let body = body.map_err(|e| { + if e.status() == StatusCode::PAYLOAD_TOO_LARGE { + ApiError::WebhookBodyTooLarge + } else { + ApiError::Malformed(e.body_text()) + } + })?; + let event = services.webhooks.receive(&webhook, &body).await?; + let (status, resp) = match event.outcome { + EventOutcome::Dropped { reason } => ( + StatusCode::OK, + EventResponse { + event_id: event.id, + verdict: "dropped", + dropped_by: Some(reason), + datapoint_id: None, + unchanged: false, + }, + ), + EventOutcome::Ingested { + datapoint_id, + unchanged, + } => ( + StatusCode::ACCEPTED, + EventResponse { + event_id: event.id, + verdict: "ingested", + dropped_by: None, + datapoint_id: Some(datapoint_id), + unchanged, + }, + ), + }; + Ok((status, axum::Json(resp))) +} diff --git a/src/webhooks/mod.rs b/src/webhooks/mod.rs new file mode 100644 index 0000000..93a4f84 --- /dev/null +++ b/src/webhooks/mod.rs @@ -0,0 +1,11 @@ +//! Inbound webhooks: a third party POSTs to a secret URL and the +//! webhook's settings turn the body into a datapoint. + +pub mod http; + +mod model; +mod rate; +mod service; + +pub use model::*; +pub use service::*; diff --git a/src/webhook.rs b/src/webhooks/model.rs similarity index 99% rename from src/webhook.rs rename to src/webhooks/model.rs index f128066..473a982 100644 --- a/src/webhook.rs +++ b/src/webhooks/model.rs @@ -1,13 +1,10 @@ -//! Inbound webhooks: a third party POSTs to a secret URL and the -//! webhook's settings turn the body into a datapoint. - use chrono::{DateTime, TimeZone, Utc}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value, json}; use sha2::{Digest, Sha256}; use uuid::Uuid; -use crate::domain::{NewDatapoint, validate_audience}; +use crate::datapoints::{NewDatapoint, validate_audience}; /// Most webhooks one project holds. pub const MAX_PER_PROJECT: usize = 20; diff --git a/src/api/rate.rs b/src/webhooks/rate.rs similarity index 100% rename from src/api/rate.rs rename to src/webhooks/rate.rs diff --git a/src/webhooks/service.rs b/src/webhooks/service.rs new file mode 100644 index 0000000..67c5f80 --- /dev/null +++ b/src/webhooks/service.rs @@ -0,0 +1,192 @@ +use std::time::Duration; + +use chrono::Utc; +use uuid::Uuid; + +use super::rate::RateLimiter; +use super::{MAX_PER_PROJECT, Settings, Webhook, is_webhook_id}; +use crate::app::Infra; +use crate::datapoints::{DatapointService, IngestOutcome}; +use crate::error::ServiceError; +use crate::locks::KeyLocks; + +/// Most events one webhook accepts per minute. +pub const EVENTS_PER_MINUTE: u32 = 600; + +/// What became of an inbound event. +#[derive(Debug)] +pub struct Event { + pub id: String, + pub outcome: EventOutcome, +} + +#[derive(Debug)] +pub enum EventOutcome { + /// Not ingested: the first filter the event failed, or the field it + /// lacked. + Dropped { reason: String }, + Ingested { + datapoint_id: Uuid, + /// True when the event changed nothing that was already stored. + unchanged: bool, + }, +} + +/// Manages a project's webhooks and turns their events into datapoints. +#[derive(Clone)] +pub struct WebhookService { + infra: Infra, + locks: KeyLocks, + /// Events accepted per webhook per minute. + rate: RateLimiter, + datapoints: DatapointService, +} + +fn not_found(id: &str) -> ServiceError { + ServiceError::NotFound(format!("no webhook {id:?} in this project")) +} + +/// The in-process lock serialising changes to one project's webhooks. +fn webhooks_lock(project: Uuid) -> String { + format!("{project}:webhooks") +} + +impl WebhookService { + pub(crate) fn new( + infra: Infra, + locks: KeyLocks, + datapoints: DatapointService, + ) -> WebhookService { + WebhookService { + infra, + locks, + rate: RateLimiter::default(), + datapoints, + } + } + + /// A new webhook and the token for its URL. The token is the only + /// credential and cannot be recovered later. + pub async fn create( + &self, + project: Uuid, + settings: Settings, + ) -> Result<(Webhook, String), ServiceError> { + settings.validate().map_err(ServiceError::Validation)?; + let _guard = self.locks.lock(webhooks_lock(project)).await; + if self.infra.store.list_webhooks(project).await?.len() >= MAX_PER_PROJECT { + return Err(ServiceError::Validation(format!( + "a project holds at most {MAX_PER_PROJECT} webhooks; delete one first" + ))); + } + let (webhook, token) = Webhook::create(project, settings, Utc::now()); + self.infra.store.put_webhook(&webhook).await?; + Ok((webhook, token)) + } + + pub async fn list(&self, project: Uuid) -> Result, ServiceError> { + Ok(self.infra.store.list_webhooks(project).await?) + } + + /// The webhook, if it exists and belongs to `project`. + pub async fn get(&self, project: Uuid, id: &str) -> Result { + if !is_webhook_id(id) { + return Err(not_found(id)); + } + match self.infra.store.get_webhook(id).await? { + Some(w) if w.project_id == project => Ok(w), + _ => Err(not_found(id)), + } + } + + /// Replace a webhook's settings. Its id and token stay the same. + pub async fn update( + &self, + project: Uuid, + id: &str, + settings: Settings, + ) -> Result { + settings.validate().map_err(ServiceError::Validation)?; + let _guard = self.locks.lock(webhooks_lock(project)).await; + let mut webhook = self.get(project, id).await?; + webhook.settings = settings; + self.infra.store.put_webhook(&webhook).await?; + Ok(webhook) + } + + /// Delete a webhook. Datapoints it created stay. False when there was + /// no such webhook. + pub async fn delete(&self, project: Uuid, id: &str) -> Result { + let _guard = self.locks.lock(webhooks_lock(project)).await; + match self.get(project, id).await { + Ok(_) => { + self.infra.store.delete_webhook(id).await?; + self.rate.forget(id); + Ok(true) + } + Err(ServiceError::NotFound(_)) => Ok(false), + Err(e) => Err(e), + } + } + + /// The webhook an event URL names. The token is the only credential, so + /// a wrong id and a wrong token look the same. + pub async fn authenticate(&self, id: &str, token: &str) -> Result { + let unknown = || ServiceError::NotFound("no such webhook".into()); + if !is_webhook_id(id) { + return Err(unknown()); + } + match self.infra.store.get_webhook(id).await? { + Some(w) if w.token_matches(token) => Ok(w), + _ => Err(unknown()), + } + } + + /// Turn an event for an authenticated webhook into a datapoint. + pub async fn receive(&self, webhook: &Webhook, body: &[u8]) -> Result { + let id = &webhook.id; + if webhook.settings.paused { + return Err(ServiceError::WebhookPaused); + } + if !self + .rate + .allow(id, EVENTS_PER_MINUTE, Duration::from_secs(60)) + { + return Err(ServiceError::RateLimited); + } + self.infra.ensure_ready()?; + + let event_id = format!("evt_{}", Uuid::now_v7().simple()); + let new = match webhook.map_event(body) { + Ok(new) => new, + Err(reason) => { + tracing::info!( + webhook = id, + event = event_id, + reason, + "webhook event dropped" + ); + return Ok(Event { + id: event_id, + outcome: EventOutcome::Dropped { reason }, + }); + } + }; + // Without an id from the event, an identical event is keyed on what it + // maps to, so a resend is not stored twice. + let key = new + .source_id + .is_none() + .then(|| format!("webhook:{id}:{}", new.fingerprint())); + let ingested = self.datapoints.ingest_one(new, key.as_deref()).await?; + let datapoint_id = ingested.datapoint.id; + tracing::info!(webhook = id, event = event_id, datapoint = %datapoint_id, "webhook event ingested"); + Ok(Event { + id: event_id, + outcome: EventOutcome::Ingested { + datapoint_id, + unchanged: ingested.outcome == IngestOutcome::Deduplicated, + }, + }) + } +} diff --git a/tests/it/boot.rs b/tests/it/boot.rs index e46087e..ee258dd 100644 --- a/tests/it/boot.rs +++ b/tests/it/boot.rs @@ -4,7 +4,7 @@ use std::sync::atomic::Ordering; use axum::http::StatusCode; use chrono::Utc; use muninn::boot::rebuild_index; -use muninn::domain::{Meta, StoredEmbedding}; +use muninn::datapoints::{Meta, StoredEmbedding}; use std::sync::Arc; use muninn::embed::{Embedder, FakeEmbedder}; diff --git a/tests/it/common/mod.rs b/tests/it/common/mod.rs index 1ef310c..175800c 100644 --- a/tests/it/common/mod.rs +++ b/tests/it/common/mod.rs @@ -7,8 +7,7 @@ use std::sync::atomic::AtomicBool; use axum::body::Body; use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; -use muninn::api::locks::KeyLocks; -use muninn::api::state::AppState; +use muninn::app::{Infra, Services}; use muninn::config::Config; use muninn::embed::FakeEmbedder; use muninn::index::MemoryIndex; @@ -56,17 +55,15 @@ pub fn project() -> Uuid { } /// The full router over a fresh bucket, a fake embedder and an in-memory -/// index, marked ready. -pub async fn app() -> (axum::Router, AppState) { - let state = AppState { +/// index, marked ready, and the infrastructure behind it. +pub async fn app() -> (axum::Router, Infra) { + let infra = Infra { store: store().await, index: Arc::new(MemoryIndex::new()), embedder: Arc::new(FakeEmbedder::new(64)), ready: Arc::new(AtomicBool::new(true)), - locks: KeyLocks::default(), - webhook_rate: Default::default(), }; - (muninn::api::router(state.clone()), state) + (muninn::http::router(Services::new(infra.clone())), infra) } pub async fn post(router: &axum::Router, path: &str, body: Value) -> (StatusCode, Value) { @@ -120,7 +117,7 @@ async fn send(router: &axum::Router, req: Request) -> (StatusCode, Value) } /// How many vectors the index holds for a project. -pub async fn index_len(state: &AppState, project: Uuid) -> usize { +pub async fn index_len(state: &Infra, project: Uuid) -> usize { let probe = vec![0.0; state.embedder.dim()]; state .index diff --git a/tests/it/delete.rs b/tests/it/delete.rs index 13e8689..e8fb4c3 100644 --- a/tests/it/delete.rs +++ b/tests/it/delete.rs @@ -43,7 +43,7 @@ async fn delete_by_source_id_and_id() { assert!(r[2].get("sourceId").is_none()); assert_eq!(common::index_len(&state, p).await, 1); - let id = muninn::domain::source_uuid(p, "crm-1"); + let id = muninn::datapoints::source_uuid(p, "crm-1"); assert!(state.store.get_meta(p, id).await.unwrap().is_none()); assert!(state.store.get_content(p, id).await.unwrap().is_none()); diff --git a/tests/it/ingest.rs b/tests/it/ingest.rs index a28f410..c88c1c4 100644 --- a/tests/it/ingest.rs +++ b/tests/it/ingest.rs @@ -2,7 +2,7 @@ use crate::common; use std::sync::atomic::Ordering; use axum::http::StatusCode; -use muninn::domain::NewDatapoint; +use muninn::datapoints::NewDatapoint; use muninn::store::MarkerOutcome; use serde_json::json; use uuid::Uuid; @@ -327,11 +327,11 @@ async fn dedupe_restores_missing_index_entry() { // Same S3 data, but the index lost the entry (request cancelled after // meta.json was written, before the index upsert). - let fresh = muninn::api::AppState { + let fresh = muninn::app::Infra { index: std::sync::Arc::new(muninn::index::MemoryIndex::new()), ..state.clone() }; - let app2 = muninn::api::router(fresh.clone()); + let app2 = muninn::http::router(muninn::app::Services::new(fresh.clone())); let (status, resp) = common::post(&app2, "/api/v1/ingest", body).await; assert_eq!(status, StatusCode::OK); assert_eq!(resp["deduplicated"], true); @@ -399,7 +399,7 @@ async fn source_id_creates_then_dedupes_identical_resend() { assert_eq!(b1["sourceId"], "crm-1"); assert_eq!(b1["replaced"], false); let id: Uuid = b1["id"].as_str().unwrap().parse().unwrap(); - assert_eq!(id, muninn::domain::source_uuid(p, "crm-1")); + assert_eq!(id, muninn::datapoints::source_uuid(p, "crm-1")); let (s2, b2) = common::post(&app, "/api/v1/ingest", body).await; assert_eq!(s2, StatusCode::OK, "{b2}"); @@ -536,7 +536,7 @@ async fn batch_with_source_ids_upserts_and_checks_labels_up_front() { .await; assert_eq!(status, StatusCode::BAD_REQUEST, "{body}"); assert_eq!(body["error"]["details"][0]["index"], 1); - let id = muninn::domain::source_uuid(p, "a"); + let id = muninn::datapoints::source_uuid(p, "a"); assert_eq!( state.store.get_content(p, id).await.unwrap().as_deref(), Some("two"), diff --git a/tests/it/store.rs b/tests/it/store.rs index b33b6a3..ad40072 100644 --- a/tests/it/store.rs +++ b/tests/it/store.rs @@ -1,6 +1,6 @@ use crate::common; use chrono::Utc; -use muninn::domain::{Meta, StoredEmbedding}; +use muninn::datapoints::{Meta, StoredEmbedding}; use muninn::store::MarkerOutcome; use uuid::Uuid; diff --git a/tests/it/webhooks.rs b/tests/it/webhooks.rs index b1d4e2c..d38f70f 100644 --- a/tests/it/webhooks.rs +++ b/tests/it/webhooks.rs @@ -94,7 +94,7 @@ async fn create_validation_and_limit() { let (status, _) = common::get(&app, "/api/v1/webhooks").await; assert_eq!(status, StatusCode::BAD_REQUEST, "projectId is required"); - for i in 0..muninn::webhook::MAX_PER_PROJECT { + for i in 0..muninn::webhooks::MAX_PER_PROJECT { create( &app, p, @@ -155,7 +155,7 @@ async fn events_become_datapoints() { .parse::() .unwrap(); let source_id = format!("{id}:SUP-1042"); - assert_eq!(datapoint, muninn::domain::source_uuid(p, &source_id)); + assert_eq!(datapoint, muninn::datapoints::source_uuid(p, &source_id)); let meta = state.store.get_meta(p, datapoint).await.unwrap().unwrap(); assert_eq!(meta.label, "Login broken"); @@ -258,7 +258,7 @@ async fn refusals() { assert_eq!(body["error"]["code"], "not-found"); } - let big = vec![b'a'; muninn::webhook::MAX_BODY_BYTES + 1]; + let big = vec![b'a'; muninn::webhooks::MAX_BODY_BYTES + 1]; let (status, body) = common::post_raw(&app, &path, "text/plain", big).await; assert_eq!(status, StatusCode::PAYLOAD_TOO_LARGE); assert_eq!(body["error"]["code"], "webhook-body-too-large"); @@ -279,7 +279,7 @@ async fn rate_limited_per_webhook() { ) .await; // Dropped events are cheap and still count. - for _ in 0..muninn::api::webhooks::EVENTS_PER_MINUTE { + for _ in 0..muninn::webhooks::EVENTS_PER_MINUTE { let (status, _) = send(&app, &path, json!({ "a": "x" })).await; assert_eq!(status, StatusCode::OK); }