Skip to main content

ninox_server/routes/
brain.rs

1use ninox_core::{embeddings::Embedder, BrainEntry, BrainIndex, QueryFilters};
2use axum::{
3    extract::{Path, Query, State},
4    http::StatusCode,
5    routing::{get, post},
6    Json, Router,
7};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11#[derive(Clone)]
12pub struct BrainState {
13    pub brain: Arc<BrainIndex>,
14    pub embedder: Option<Arc<dyn Embedder>>,
15}
16
17pub fn brain_router(brain: Arc<BrainIndex>, embedder: Option<Arc<dyn Embedder>>) -> Router {
18    Router::new()
19        .route("/index", post(rebuild_index))
20        .route("/query", get(query_entries))
21        .route("/entry/*path", get(get_entry))
22        .with_state(BrainState { brain, embedder })
23}
24
25#[derive(Serialize)]
26struct IndexResponse {
27    count: usize,
28    embedded: usize,
29    cached: usize,
30}
31
32async fn rebuild_index(
33    State(state): State<BrainState>,
34) -> Result<Json<IndexResponse>, StatusCode> {
35    match state.brain.rebuild(state.embedder.as_deref()) {
36        Ok(stats) => Ok(Json(IndexResponse {
37            count: stats.indexed,
38            embedded: stats.embedded,
39            cached: stats.cached,
40        })),
41        Err(err) => {
42            tracing::error!("brain rebuild: {err}");
43            Err(StatusCode::INTERNAL_SERVER_ERROR)
44        }
45    }
46}
47
48#[derive(Deserialize)]
49struct QueryParams {
50    q: Option<String>,
51    #[serde(rename = "type")]
52    entry_type: Option<String>,
53    tag: Option<String>,
54}
55
56async fn query_entries(
57    State(state): State<BrainState>,
58    Query(params): Query<QueryParams>,
59) -> Result<Json<Vec<BrainEntry>>, StatusCode> {
60    let text = params.q.unwrap_or_default();
61    let filters = QueryFilters {
62        entry_type: params.entry_type,
63        tag: params.tag,
64    };
65    match state.brain.query(&text, state.embedder.as_deref(), filters) {
66        Ok(entries) => Ok(Json(entries)),
67        Err(err) => {
68            tracing::error!("brain query: {err}");
69            Err(StatusCode::INTERNAL_SERVER_ERROR)
70        }
71    }
72}
73
74async fn get_entry(
75    State(state): State<BrainState>,
76    Path(path): Path<String>,
77) -> Result<Json<BrainEntry>, StatusCode> {
78    match state.brain.get(&path) {
79        Ok(Some(entry)) => Ok(Json(entry)),
80        Ok(None) => Err(StatusCode::NOT_FOUND),
81        Err(err) => {
82            tracing::error!("brain get {path}: {err}");
83            Err(StatusCode::INTERNAL_SERVER_ERROR)
84        }
85    }
86}