docling_rag/
api.rs

1//! REST API over a [`Pipeline`]: document info and search in every retrieval mode.
2//!
3//! Authentication is a static API-key list from config (`RAG_API_KEYS`), accepted
4//! as `X-Api-Key: <key>` or `Authorization: Bearer <key>`. Auth is fail-closed:
5//! [`router`] errors when the key list is empty. `GET /health` is public.
6//!
7//! Endpoints (all under auth except `/health`):
8//!
9//! | Method | Path                  | Description                                   |
10//! |--------|-----------------------|-----------------------------------------------|
11//! | GET    | `/health`             | liveness probe (public)                       |
12//! | GET    | `/api/stats`          | document / chunk counts                       |
13//! | GET    | `/api/documents`      | all documents with metadata + metrics         |
14//! | GET    | `/api/documents/{id}` | one document by id                            |
15//! | GET    | `/api/search`         | `?q=…&mode=hybrid&k=5` (also accepts POST)    |
16//! | POST   | `/api/search`         | `{"query", "mode?", "top_k?", "answer?"}`     |
17//!
18//! Search modes: `vector`, `bm25`, `hybrid`, `multi-query`, `hyde`. With
19//! `answer=true` the LLM synthesizes a grounded answer (needs `OPENROUTER_API_KEY`).
20
21use crate::model::RetrievalMode;
22use crate::pipeline::Pipeline;
23use crate::{RagError, Result};
24use axum::extract::{Path, Query, Request, State};
25use axum::http::{header, StatusCode};
26use axum::middleware::{self, Next};
27use axum::response::{IntoResponse, Response};
28use axum::routing::get;
29use axum::{Json, Router};
30use serde::Deserialize;
31use serde_json::json;
32use std::collections::HashSet;
33use std::str::FromStr;
34use std::sync::Arc;
35
36struct AppState {
37    pipeline: Pipeline,
38    keys: HashSet<String>,
39}
40
41/// Build the router. Errors if `keys` is empty (auth is fail-closed).
42pub fn router(pipeline: Pipeline, keys: Vec<String>) -> Result<Router> {
43    if keys.is_empty() {
44        return Err(RagError::config(
45            "RAG_API_KEYS must contain at least one key to start the REST API",
46        ));
47    }
48    let state = Arc::new(AppState {
49        pipeline,
50        keys: keys.into_iter().collect(),
51    });
52
53    let protected = Router::new()
54        .route("/api/stats", get(stats))
55        .route("/api/documents", get(list_documents))
56        .route("/api/documents/{id}", get(get_document))
57        .route("/api/search", get(search_get).post(search_post))
58        .layer(middleware::from_fn_with_state(state.clone(), auth));
59
60    Ok(Router::new()
61        .route("/health", get(|| async { Json(json!({"status": "ok"})) }))
62        .merge(protected)
63        .with_state(state))
64}
65
66/// Bind `addr` and serve until the process is stopped.
67pub async fn serve(pipeline: Pipeline, addr: &str, keys: Vec<String>) -> Result<()> {
68    let app = router(pipeline, keys)?;
69    let listener = tokio::net::TcpListener::bind(addr)
70        .await
71        .map_err(|e| RagError::config(format!("cannot bind {addr}: {e}")))?;
72    tracing::info!(%addr, "REST API listening");
73    axum::serve(listener, app)
74        .await
75        .map_err(|e| RagError::config(format!("server error: {e}")))
76}
77
78async fn auth(State(state): State<Arc<AppState>>, req: Request, next: Next) -> Response {
79    let headers = req.headers();
80    let provided = headers
81        .get("x-api-key")
82        .and_then(|v| v.to_str().ok())
83        .or_else(|| {
84            headers
85                .get(header::AUTHORIZATION)
86                .and_then(|v| v.to_str().ok())
87                .and_then(|v| v.strip_prefix("Bearer "))
88        });
89    match provided {
90        Some(key) if state.keys.contains(key) => next.run(req).await,
91        _ => err(StatusCode::UNAUTHORIZED, "invalid or missing API key").into_response(),
92    }
93}
94
95type ApiResult = std::result::Result<Response, (StatusCode, Json<serde_json::Value>)>;
96
97fn err(code: StatusCode, msg: impl std::fmt::Display) -> (StatusCode, Json<serde_json::Value>) {
98    (code, Json(json!({"error": msg.to_string()})))
99}
100
101fn internal(e: RagError) -> (StatusCode, Json<serde_json::Value>) {
102    err(StatusCode::INTERNAL_SERVER_ERROR, e)
103}
104
105async fn stats(State(state): State<Arc<AppState>>) -> ApiResult {
106    let store = state.pipeline.store();
107    let documents = store.count_documents().await.map_err(internal)?;
108    let chunks = store.count_chunks().await.map_err(internal)?;
109    Ok(Json(json!({"documents": documents, "chunks": chunks})).into_response())
110}
111
112async fn list_documents(State(state): State<Arc<AppState>>) -> ApiResult {
113    let docs = state
114        .pipeline
115        .store()
116        .list_documents()
117        .await
118        .map_err(internal)?;
119    Ok(Json(json!({"documents": docs})).into_response())
120}
121
122async fn get_document(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> ApiResult {
123    let docs = state
124        .pipeline
125        .store()
126        .list_documents()
127        .await
128        .map_err(internal)?;
129    match docs.into_iter().find(|d| d.id == id) {
130        Some(doc) => Ok(Json(doc).into_response()),
131        None => Err(err(
132            StatusCode::NOT_FOUND,
133            format!("no document with id '{id}'"),
134        )),
135    }
136}
137
138/// Search parameters, shared by the GET (query-string) and POST (JSON) forms.
139#[derive(Debug, Deserialize)]
140struct SearchParams {
141    /// The search query (`q` also accepted on GET).
142    #[serde(alias = "q")]
143    query: String,
144    /// vector | bm25 | hybrid | multi-query | hyde. Defaults to the configured mode.
145    mode: Option<String>,
146    /// Number of results (default: configured top_k).
147    #[serde(alias = "k")]
148    top_k: Option<usize>,
149    /// Also synthesize an LLM answer grounded in the results.
150    #[serde(default)]
151    answer: bool,
152}
153
154async fn search_get(
155    State(state): State<Arc<AppState>>,
156    Query(params): Query<SearchParams>,
157) -> ApiResult {
158    run_search(state, params).await
159}
160
161async fn search_post(
162    State(state): State<Arc<AppState>>,
163    Json(params): Json<SearchParams>,
164) -> ApiResult {
165    run_search(state, params).await
166}
167
168async fn run_search(state: Arc<AppState>, params: SearchParams) -> ApiResult {
169    if params.query.trim().is_empty() {
170        return Err(err(StatusCode::BAD_REQUEST, "query must not be empty"));
171    }
172    let mode = match &params.mode {
173        Some(m) => RetrievalMode::from_str(m).map_err(|e| err(StatusCode::BAD_REQUEST, e))?,
174        None => state.pipeline.config().retrieval_mode,
175    };
176    let k = params
177        .top_k
178        .unwrap_or(state.pipeline.config().top_k)
179        .clamp(1, 100);
180
181    if params.answer {
182        let a = state
183            .pipeline
184            .answer(&params.query, mode, k)
185            .await
186            .map_err(|e| match e {
187                RagError::Llm(_) => err(StatusCode::BAD_REQUEST, e),
188                other => internal(other),
189            })?;
190        return Ok(Json(json!({
191            "query": params.query,
192            "mode": mode.to_string(),
193            "answer": a.text,
194            "results": a.sources,
195        }))
196        .into_response());
197    }
198
199    let hits = state
200        .pipeline
201        .query(mode, &params.query, k)
202        .await
203        .map_err(|e| match e {
204            RagError::Llm(_) => err(StatusCode::BAD_REQUEST, e),
205            other => internal(other),
206        })?;
207    Ok(Json(json!({
208        "query": params.query,
209        "mode": mode.to_string(),
210        "results": hits,
211    }))
212    .into_response())
213}
docling_rag/api.rs

docling_rag/
api.rs
